Skip to main content

alien_bindings/providers/kv/
gcp_firestore.rs

1use crate::error::{ErrorData, Result};
2use crate::traits::{Binding, Kv, KvEntry, PutCondition, PutOptions, ScanResult};
3use alien_error::{AlienError, Context, IntoAlienError};
4use alien_gcp_clients::firestore::{
5    CollectionSelector, CompositeFilter, CompositeFilterOperator, Cursor, Direction, Document,
6    FieldFilter, FieldFilterOperator, FieldReference, Filter, FirestoreApi, FirestoreClient, Order,
7    QueryType, RunQueryRequest, StructuredQuery, Value,
8};
9use async_trait::async_trait;
10use base64::{self, Engine};
11use chrono::{DateTime, Utc};
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14use std::fmt::{Debug, Formatter};
15
16use super::{decode_version, encode_version, validate_key, validate_value};
17
18/// Firestore document for KV storage
19#[derive(Debug, Clone, Serialize, Deserialize)]
20struct KvDocument {
21    value: String, // Base64-encoded binary data
22    created_at: DateTime<Utc>,
23    expires_at: Option<DateTime<Utc>>, // For TTL policy
24}
25
26#[derive(Debug, Serialize, Deserialize)]
27#[serde(rename_all = "camelCase")]
28struct CursorState {
29    version: u8,
30    prefix: String,
31    last_document_name: String,
32}
33
34/// GCP Firestore implementation of the KV trait
35pub struct GcpFirestoreKv {
36    client: FirestoreClient,
37    project_id: String,
38    database_id: String,
39    collection_name: String,
40}
41
42impl Debug for GcpFirestoreKv {
43    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
44        f.debug_struct("GcpFirestoreKv")
45            .field("project_id", &self.project_id)
46            .field("database_id", &self.database_id)
47            .field("collection_name", &self.collection_name)
48            .finish()
49    }
50}
51
52impl GcpFirestoreKv {
53    pub fn new(
54        client: FirestoreClient,
55        project_id: String,
56        database_id: String,
57        collection_name: String,
58    ) -> Result<Self> {
59        Ok(Self {
60            client,
61            project_id,
62            database_id,
63            collection_name,
64        })
65    }
66
67    /// Checks if an item has expired based on TTL
68    fn is_expired(&self, expires_at: Option<DateTime<Utc>>) -> bool {
69        if let Some(expiry) = expires_at {
70            Utc::now() >= expiry
71        } else {
72            false
73        }
74    }
75
76    fn document_name(&self, key: &str) -> String {
77        format!(
78            "projects/{}/databases/{}/documents/{}/{}",
79            self.project_id, self.database_id, self.collection_name, key
80        )
81    }
82
83    fn encode_cursor(state: &CursorState) -> Result<String> {
84        let json =
85            serde_json::to_vec(state)
86                .into_alien_error()
87                .context(ErrorData::InvalidInput {
88                    operation_context: "Firestore KV scan cursor encoding".to_string(),
89                    details: "Failed to serialize cursor state".to_string(),
90                    field_name: Some("cursor".to_string()),
91                })?;
92        Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(json))
93    }
94
95    fn decode_cursor(&self, prefix: &str, cursor: &str) -> Result<CursorState> {
96        let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
97            .decode(cursor)
98            .into_alien_error()
99            .context(ErrorData::InvalidInput {
100                operation_context: "Firestore KV scan cursor decoding".to_string(),
101                details: "Invalid cursor encoding".to_string(),
102                field_name: Some("cursor".to_string()),
103            })?;
104        let state: CursorState = serde_json::from_slice(&decoded)
105            .into_alien_error()
106            .context(ErrorData::InvalidInput {
107                operation_context: "Firestore KV scan cursor decoding".to_string(),
108                details: "Invalid cursor JSON".to_string(),
109                field_name: Some("cursor".to_string()),
110            })?;
111        if state.version != 1
112            || state.prefix != prefix
113            || !state
114                .last_document_name
115                .starts_with(&self.document_name(prefix))
116        {
117            return Err(AlienError::new(ErrorData::InvalidInput {
118                operation_context: "Firestore KV scan cursor validation".to_string(),
119                details: "Cursor does not belong to this prefix scan".to_string(),
120                field_name: Some("cursor".to_string()),
121            }));
122        }
123        Ok(state)
124    }
125
126    /// Converts a KV document to Firestore Document format
127    fn kv_document_to_firestore(&self, _key: &str, kv_doc: &KvDocument) -> Document {
128        let mut fields = HashMap::new();
129
130        fields.insert(
131            "value".to_string(),
132            Value::StringValue(kv_doc.value.clone()),
133        );
134        fields.insert(
135            "created_at".to_string(),
136            Value::TimestampValue(kv_doc.created_at.to_rfc3339()),
137        );
138
139        if let Some(expires_at) = kv_doc.expires_at {
140            fields.insert(
141                "expires_at".to_string(),
142                Value::TimestampValue(expires_at.to_rfc3339()),
143            );
144        }
145
146        Document::builder().fields(fields).build()
147    }
148
149    /// Converts a KV document to Firestore Document format with name (for updates)
150    fn kv_document_to_firestore_with_name(&self, key: &str, kv_doc: &KvDocument) -> Document {
151        let mut fields = HashMap::new();
152
153        fields.insert(
154            "value".to_string(),
155            Value::StringValue(kv_doc.value.clone()),
156        );
157        fields.insert(
158            "created_at".to_string(),
159            Value::TimestampValue(kv_doc.created_at.to_rfc3339()),
160        );
161
162        if let Some(expires_at) = kv_doc.expires_at {
163            fields.insert(
164                "expires_at".to_string(),
165                Value::TimestampValue(expires_at.to_rfc3339()),
166            );
167        }
168
169        Document::builder()
170            .name(format!(
171                "projects/{}/databases/{}/documents/{}/{}",
172                self.project_id, self.database_id, self.collection_name, key
173            ))
174            .fields(fields)
175            .build()
176    }
177
178    /// Converts a Firestore Document to KV document
179    /// Attempt to take over a logically-expired document after a
180    /// conditional create conflict. Returns `Ok(true)` only when THIS
181    /// caller replaced the expired document (an `updateTime` precondition
182    /// makes the replace atomic — a racing taker bumps the update time and
183    /// this patch loses); a live document, a lost race, or a deletion in
184    /// between all resolve to `Ok(false)`.
185    async fn try_take_over_expired(&self, key: &str, document: &Document) -> Result<bool> {
186        use alien_client_core::ErrorData as CloudErrorData;
187        use alien_gcp_clients::gcp::firestore::{Precondition, PreconditionType};
188
189        let document_path = format!("{}/{}", self.collection_name, key);
190        let existing = match self
191            .client
192            .get_document(
193                self.database_id.clone(),
194                document_path.clone(),
195                None,
196                None,
197                None,
198            )
199            .await
200        {
201            Ok(existing) => existing,
202            Err(e)
203                if matches!(
204                    e.error.as_ref(),
205                    Some(CloudErrorData::RemoteResourceNotFound { .. })
206                ) =>
207            {
208                return Ok(false);
209            }
210            Err(e) => {
211                return Err(crate::error::map_cloud_client_error(
212                    e,
213                    format!("Failed to read existing document for key '{}'", key),
214                    Some(key.to_string()),
215                ));
216            }
217        };
218
219        let kv_doc = self.firestore_to_kv_document(&existing)?;
220        if !self.is_expired(kv_doc.expires_at) {
221            return Ok(false);
222        }
223        let Some(update_time) = existing.update_time.clone() else {
224            return Ok(false);
225        };
226
227        match self
228            .client
229            .patch_document(
230                self.database_id.clone(),
231                document_path,
232                document.clone(),
233                None,
234                None,
235                Some(Precondition {
236                    condition: PreconditionType::UpdateTime(update_time),
237                }),
238            )
239            .await
240        {
241            Ok(_) => Ok(true),
242            Err(e) => match e.error.as_ref() {
243                // A lost race: the precondition mismatch maps to Conflict
244                // (FAILED_PRECONDITION → RemoteResourceConflict in
245                // map_gcp_error), and a deletion in between maps to NotFound.
246                Some(CloudErrorData::RemoteResourceConflict { .. })
247                | Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(false),
248                _ => Err(crate::error::map_cloud_client_error(
249                    e,
250                    format!("Failed to take over expired document for key '{}'", key),
251                    Some(key.to_string()),
252                )),
253            },
254        }
255    }
256
257    fn firestore_to_kv_document(&self, doc: &Document) -> Result<KvDocument> {
258        let fields = doc.fields.as_ref().ok_or_else(|| {
259            AlienError::new(ErrorData::UnexpectedResponseFormat {
260                provider: "gcp".to_string(),
261                binding_name: "firestore".to_string(),
262                field: "fields".to_string(),
263                response_json: serde_json::to_string(doc).unwrap_or_default(),
264            })
265        })?;
266
267        let value = match fields.get("value") {
268            Some(Value::StringValue(v)) => v.clone(),
269            _ => {
270                return Err(AlienError::new(ErrorData::UnexpectedResponseFormat {
271                    provider: "gcp".to_string(),
272                    binding_name: "firestore".to_string(),
273                    field: "value".to_string(),
274                    response_json: serde_json::to_string(doc).unwrap_or_default(),
275                }))
276            }
277        };
278
279        let created_at = match fields.get("created_at") {
280            Some(Value::TimestampValue(t)) => DateTime::parse_from_rfc3339(t)
281                .map_err(|_| {
282                    AlienError::new(ErrorData::UnexpectedResponseFormat {
283                        provider: "gcp".to_string(),
284                        binding_name: "firestore".to_string(),
285                        field: "created_at".to_string(),
286                        response_json: serde_json::to_string(doc).unwrap_or_default(),
287                    })
288                })?
289                .with_timezone(&Utc),
290            _ => {
291                return Err(AlienError::new(ErrorData::UnexpectedResponseFormat {
292                    provider: "gcp".to_string(),
293                    binding_name: "firestore".to_string(),
294                    field: "created_at".to_string(),
295                    response_json: serde_json::to_string(doc).unwrap_or_default(),
296                }))
297            }
298        };
299
300        let expires_at = match fields.get("expires_at") {
301            Some(Value::TimestampValue(t)) => Some(
302                DateTime::parse_from_rfc3339(t)
303                    .map_err(|_| {
304                        AlienError::new(ErrorData::UnexpectedResponseFormat {
305                            provider: "gcp".to_string(),
306                            binding_name: "firestore".to_string(),
307                            field: "expires_at".to_string(),
308                            response_json: serde_json::to_string(doc).unwrap_or_default(),
309                        })
310                    })?
311                    .with_timezone(&Utc),
312            ),
313            _ => None,
314        };
315
316        Ok(KvDocument {
317            value,
318            created_at,
319            expires_at,
320        })
321    }
322}
323
324impl Binding for GcpFirestoreKv {}
325
326#[async_trait]
327impl Kv for GcpFirestoreKv {
328    async fn get(&self, key: &str) -> Result<Option<KvEntry>> {
329        validate_key(key)?;
330
331        let document_id = key;
332        let document_path = format!("{}/{}", self.collection_name, document_id);
333
334        match self
335            .client
336            .get_document(self.database_id.clone(), document_path, None, None, None)
337            .await
338        {
339            Ok(doc) => {
340                let kv_doc = self.firestore_to_kv_document(&doc)?;
341
342                // Check TTL expiry (logical expiry contract)
343                if self.is_expired(kv_doc.expires_at) {
344                    return Ok(None); // Logically expired
345                }
346
347                let value = base64::engine::general_purpose::STANDARD
348                    .decode(&kv_doc.value)
349                    .into_alien_error()
350                    .context(ErrorData::KvOperationFailed {
351                        operation: "get".to_string(),
352                        key: key.to_string(),
353                        reason: "Failed to decode base64 value".to_string(),
354                    })?;
355
356                let update_time = doc.update_time.clone().ok_or_else(|| {
357                    AlienError::new(ErrorData::UnexpectedResponseFormat {
358                        provider: "gcp".to_string(),
359                        binding_name: "firestore".to_string(),
360                        field: "updateTime".to_string(),
361                        response_json: serde_json::to_string(&doc).unwrap_or_default(),
362                    })
363                })?;
364                Ok(Some(KvEntry {
365                    key: key.to_string(),
366                    value,
367                    version: encode_version(
368                        key,
369                        update_time,
370                        kv_doc
371                            .expires_at
372                            .map(|expires_at| expires_at.timestamp_millis()),
373                    )?,
374                }))
375            }
376            Err(e) => {
377                // Check if this is a "not found" error
378                match &e.error {
379                    Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. }) => {
380                        Ok(None) // Document doesn't exist
381                    }
382                    _ => Err(crate::error::map_cloud_client_error(
383                        e,
384                        "Failed to get Firestore document".to_string(),
385                        Some(key.to_string()),
386                    )),
387                }
388            }
389        }
390    }
391
392    async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
393        validate_key(key)?;
394        validate_value(&value)?;
395
396        let options = options.unwrap_or_default();
397
398        let encoded_value = base64::engine::general_purpose::STANDARD.encode(&value);
399        let kv_doc = KvDocument {
400            value: encoded_value,
401            created_at: Utc::now(),
402            expires_at: options.ttl.map(|d| Utc::now() + d),
403        };
404
405        let document = self.kv_document_to_firestore(key, &kv_doc);
406
407        if matches!(options.condition, PutCondition::Absent) {
408            let document_id = key.to_string();
409            match self
410                .client
411                .create_document(
412                    self.database_id.clone(),
413                    self.collection_name.clone(),
414                    Some(document_id),
415                    document.clone(),
416                    None,
417                )
418                .await
419            {
420                Ok(_) => Ok(true),
421                Err(e) => {
422                    // Check if this is a conflict (document already exists)
423                    match &e.error {
424                        Some(alien_client_core::ErrorData::RemoteResourceConflict { .. }) => {
425                            // Expired documents count as ABSENT, matching the
426                            // local provider's atomic takeover: Firestore's
427                            // TTL deletion can lag the logical expiry by
428                            // hours, and without a takeover a conditional
429                            // put (e.g. a command lease after the previous
430                            // holder died) stays blocked until then.
431                            self.try_take_over_expired(key, &document).await
432                        }
433                        _ => Err(crate::error::map_cloud_client_error(
434                            e,
435                            "Failed to create Firestore document".to_string(),
436                            Some(key.to_string()),
437                        )),
438                    }
439                }
440            }
441        } else {
442            let document_id = key;
443            let document_path = format!("{}/{}", self.collection_name, document_id);
444            let document_with_name = self.kv_document_to_firestore_with_name(key, &kv_doc);
445            let current_document = match &options.condition {
446                PutCondition::Version(version) => {
447                    let expected = decode_version(key, version)?;
448                    if expected.expired {
449                        return Ok(false);
450                    }
451                    Some(alien_gcp_clients::gcp::firestore::Precondition {
452                        condition: alien_gcp_clients::gcp::firestore::PreconditionType::UpdateTime(
453                            expected.backend_version,
454                        ),
455                    })
456                }
457                PutCondition::None => None,
458                PutCondition::Absent => unreachable!("handled above"),
459            };
460
461            match self
462                .client
463                .patch_document(
464                    self.database_id.clone(),
465                    document_path,
466                    document_with_name,
467                    None,
468                    None,
469                    current_document,
470                )
471                .await
472            {
473                Ok(_) => Ok(true),
474                Err(error)
475                    if matches!(options.condition, PutCondition::Version(_))
476                        && matches!(
477                            error.error.as_ref(),
478                            Some(alien_client_core::ErrorData::RemoteResourceConflict { .. })
479                                | Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. })
480                        ) =>
481                {
482                    Ok(false)
483                }
484                Err(error) => Err(crate::error::map_cloud_client_error(
485                    error,
486                    "Failed to patch Firestore document".to_string(),
487                    Some(key.to_string()),
488                )),
489            }
490        }
491    }
492
493    async fn delete(&self, key: &str, if_version: Option<&str>) -> Result<bool> {
494        validate_key(key)?;
495
496        let document_id = key;
497        let document_path = format!("{}/{}", self.collection_name, document_id);
498
499        let current_document = if let Some(version) = if_version {
500            let expected = decode_version(key, version)?;
501            if expected.expired {
502                return Ok(false);
503            }
504            Some(alien_gcp_clients::gcp::firestore::Precondition {
505                condition: alien_gcp_clients::gcp::firestore::PreconditionType::UpdateTime(
506                    expected.backend_version,
507                ),
508            })
509        } else {
510            None
511        };
512
513        match self
514            .client
515            .delete_document(self.database_id.clone(), document_path, current_document)
516            .await
517        {
518            Ok(()) => Ok(true),
519            Err(error)
520                if if_version.is_some()
521                    && matches!(
522                        error.error.as_ref(),
523                        Some(alien_client_core::ErrorData::RemoteResourceConflict { .. })
524                            | Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. })
525                    ) =>
526            {
527                Ok(false)
528            }
529            Err(error) => Err(crate::error::map_cloud_client_error(
530                error,
531                "Failed to delete Firestore document".to_string(),
532                Some(key.to_string()),
533            )),
534        }
535    }
536
537    async fn exists(&self, key: &str) -> Result<bool> {
538        validate_key(key)?;
539
540        let document_id = key;
541        let document_path = format!("{}/{}", self.collection_name, document_id);
542
543        match self
544            .client
545            .get_document(self.database_id.clone(), document_path, None, None, None)
546            .await
547        {
548            Ok(doc) => {
549                let kv_doc = self.firestore_to_kv_document(&doc)?;
550
551                // Check TTL expiry (logical expiry contract)
552                Ok(!self.is_expired(kv_doc.expires_at))
553            }
554            Err(e) => {
555                match &e.error {
556                    Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. }) => {
557                        Ok(false) // Document doesn't exist
558                    }
559                    _ => Err(crate::error::map_cloud_client_error(
560                        e,
561                        "Failed to get Firestore document".to_string(),
562                        Some(key.to_string()),
563                    )),
564                }
565            }
566        }
567    }
568
569    async fn scan_prefix(
570        &self,
571        prefix: &str,
572        limit: Option<usize>,
573        cursor: Option<String>,
574    ) -> Result<ScanResult> {
575        validate_key(prefix)?;
576
577        let limit = limit.unwrap_or(1000);
578        let cursor_state = cursor
579            .as_deref()
580            .map(|cursor| self.decode_cursor(prefix, cursor))
581            .transpose()?;
582        if limit == 0 {
583            return Ok(ScanResult {
584                items: Vec::new(),
585                next_cursor: cursor,
586            });
587        }
588
589        let collection_selector = CollectionSelector::builder()
590            .collection_id(self.collection_name.clone())
591            .build();
592
593        let document_name_field = FieldReference::builder()
594            .field_path("__name__".to_string())
595            .build();
596        let lower = self.document_name(prefix);
597        let upper = self.document_name(&format!("{}~", prefix));
598        let mut structured_query = StructuredQuery::builder()
599            .from(vec![collection_selector])
600            .order_by(vec![Order::builder()
601                .field(document_name_field.clone())
602                .direction(Direction::Ascending)
603                .build()])
604            .r#where(Filter::CompositeFilter(
605                CompositeFilter::builder()
606                    .op(CompositeFilterOperator::And)
607                    .filters(vec![
608                        Filter::FieldFilter(
609                            FieldFilter::builder()
610                                .field(document_name_field.clone())
611                                .op(FieldFilterOperator::GreaterThanOrEqual)
612                                .value(Value::ReferenceValue(lower))
613                                .build(),
614                        ),
615                        Filter::FieldFilter(
616                            FieldFilter::builder()
617                                .field(document_name_field)
618                                .op(FieldFilterOperator::LessThan)
619                                .value(Value::ReferenceValue(upper))
620                                .build(),
621                        ),
622                    ])
623                    .build(),
624            ))
625            .limit(i32::try_from(limit).unwrap_or(i32::MAX))
626            .build();
627
628        if let Some(state) = cursor_state {
629            structured_query.start_at = Some(
630                Cursor::builder()
631                    .values(vec![Value::ReferenceValue(state.last_document_name)])
632                    .before(false)
633                    .build(),
634            );
635        }
636
637        let query_request = RunQueryRequest::builder()
638            .parent(format!(
639                "projects/{}/databases/{}/documents",
640                self.project_id, self.database_id
641            ))
642            .query_type(QueryType::StructuredQuery(structured_query))
643            .build();
644
645        let query_responses = self
646            .client
647            .run_query(self.database_id.clone(), query_request)
648            .await
649            .map_err(|e| {
650                crate::error::map_cloud_client_error(
651                    e,
652                    "Failed to run Firestore query".to_string(),
653                    Some(prefix.to_string()),
654                )
655            })?;
656
657        let documents = query_responses
658            .iter()
659            .filter_map(|response| response.document.as_ref())
660            .collect::<Vec<_>>();
661        let last_document_name = documents.last().and_then(|document| document.name.clone());
662        let mut items = Vec::with_capacity(documents.len());
663        for document in &documents {
664            let Some(name) = &document.name else {
665                continue;
666            };
667            let Some(key) = name.rsplit('/').next() else {
668                continue;
669            };
670            let kv_doc = self.firestore_to_kv_document(document)?;
671            if self.is_expired(kv_doc.expires_at) {
672                continue;
673            }
674            let value = base64::engine::general_purpose::STANDARD
675                .decode(&kv_doc.value)
676                .into_alien_error()
677                .context(ErrorData::UnexpectedResponseFormat {
678                    provider: "gcp".to_string(),
679                    binding_name: "firestore".to_string(),
680                    field: "value".to_string(),
681                    response_json: serde_json::to_string(document).unwrap_or_default(),
682                })?;
683            let update_time = document.update_time.clone().ok_or_else(|| {
684                AlienError::new(ErrorData::UnexpectedResponseFormat {
685                    provider: "gcp".to_string(),
686                    binding_name: "firestore".to_string(),
687                    field: "updateTime".to_string(),
688                    response_json: serde_json::to_string(document).unwrap_or_default(),
689                })
690            })?;
691            items.push(KvEntry {
692                key: key.to_string(),
693                value,
694                version: encode_version(
695                    key,
696                    update_time,
697                    kv_doc
698                        .expires_at
699                        .map(|expires_at| expires_at.timestamp_millis()),
700                )?,
701            });
702        }
703
704        let next_cursor = if documents.len() == limit {
705            last_document_name
706                .map(|last_document_name| {
707                    Self::encode_cursor(&CursorState {
708                        version: 1,
709                        prefix: prefix.to_string(),
710                        last_document_name,
711                    })
712                })
713                .transpose()?
714        } else {
715            None
716        };
717
718        Ok(ScanResult { items, next_cursor })
719    }
720}