Skip to main content

alien_bindings/providers/kv/
gcp_firestore.rs

1use crate::error::{ErrorData, Result};
2use crate::traits::{Binding, Kv, 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::{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<Vec<u8>>> {
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                Ok(Some(value))
357            }
358            Err(e) => {
359                // Check if this is a "not found" error
360                match &e.error {
361                    Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. }) => {
362                        Ok(None) // Document doesn't exist
363                    }
364                    _ => Err(crate::error::map_cloud_client_error(
365                        e,
366                        "Failed to get Firestore document".to_string(),
367                        Some(key.to_string()),
368                    )),
369                }
370            }
371        }
372    }
373
374    async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
375        validate_key(key)?;
376        validate_value(&value)?;
377
378        let options = options.unwrap_or_default();
379
380        let encoded_value = base64::engine::general_purpose::STANDARD.encode(&value);
381        let kv_doc = KvDocument {
382            value: encoded_value,
383            created_at: Utc::now(),
384            expires_at: options.ttl.map(|d| Utc::now() + d),
385        };
386
387        let document = self.kv_document_to_firestore(key, &kv_doc);
388
389        if options.if_not_exists {
390            let document_id = key.to_string();
391            match self
392                .client
393                .create_document(
394                    self.database_id.clone(),
395                    self.collection_name.clone(),
396                    Some(document_id),
397                    document.clone(),
398                    None,
399                )
400                .await
401            {
402                Ok(_) => Ok(true),
403                Err(e) => {
404                    // Check if this is a conflict (document already exists)
405                    match &e.error {
406                        Some(alien_client_core::ErrorData::RemoteResourceConflict { .. }) => {
407                            // Expired documents count as ABSENT, matching the
408                            // local provider's atomic takeover: Firestore's
409                            // TTL deletion can lag the logical expiry by
410                            // hours, and without a takeover a conditional
411                            // put (e.g. a command lease after the previous
412                            // holder died) stays blocked until then.
413                            self.try_take_over_expired(key, &document).await
414                        }
415                        _ => Err(crate::error::map_cloud_client_error(
416                            e,
417                            "Failed to create Firestore document".to_string(),
418                            Some(key.to_string()),
419                        )),
420                    }
421                }
422            }
423        } else {
424            let document_id = key;
425            let document_path = format!("{}/{}", self.collection_name, document_id);
426            let document_with_name = self.kv_document_to_firestore_with_name(key, &kv_doc);
427
428            self.client
429                .patch_document(
430                    self.database_id.clone(),
431                    document_path,
432                    document_with_name,
433                    None,
434                    None,
435                    None,
436                )
437                .await
438                .map_err(|e| {
439                    crate::error::map_cloud_client_error(
440                        e,
441                        "Failed to patch Firestore document".to_string(),
442                        Some(key.to_string()),
443                    )
444                })?;
445
446            Ok(true)
447        }
448    }
449
450    async fn delete(&self, key: &str) -> Result<()> {
451        validate_key(key)?;
452
453        let document_id = key;
454        let document_path = format!("{}/{}", self.collection_name, document_id);
455
456        self.client
457            .delete_document(self.database_id.clone(), document_path, None)
458            .await
459            .map_err(|e| {
460                crate::error::map_cloud_client_error(
461                    e,
462                    "Failed to delete Firestore document".to_string(),
463                    Some(key.to_string()),
464                )
465            })?;
466
467        Ok(())
468    }
469
470    async fn exists(&self, key: &str) -> Result<bool> {
471        validate_key(key)?;
472
473        let document_id = key;
474        let document_path = format!("{}/{}", self.collection_name, document_id);
475
476        match self
477            .client
478            .get_document(self.database_id.clone(), document_path, None, None, None)
479            .await
480        {
481            Ok(doc) => {
482                let kv_doc = self.firestore_to_kv_document(&doc)?;
483
484                // Check TTL expiry (logical expiry contract)
485                Ok(!self.is_expired(kv_doc.expires_at))
486            }
487            Err(e) => {
488                match &e.error {
489                    Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. }) => {
490                        Ok(false) // Document doesn't exist
491                    }
492                    _ => Err(crate::error::map_cloud_client_error(
493                        e,
494                        "Failed to get Firestore document".to_string(),
495                        Some(key.to_string()),
496                    )),
497                }
498            }
499        }
500    }
501
502    async fn scan_prefix(
503        &self,
504        prefix: &str,
505        limit: Option<usize>,
506        cursor: Option<String>,
507    ) -> Result<ScanResult> {
508        validate_key(prefix)?;
509
510        let limit = limit.unwrap_or(1000);
511        let cursor_state = cursor
512            .as_deref()
513            .map(|cursor| self.decode_cursor(prefix, cursor))
514            .transpose()?;
515        if limit == 0 {
516            return Ok(ScanResult {
517                items: Vec::new(),
518                next_cursor: cursor,
519            });
520        }
521
522        let collection_selector = CollectionSelector::builder()
523            .collection_id(self.collection_name.clone())
524            .build();
525
526        let document_name_field = FieldReference::builder()
527            .field_path("__name__".to_string())
528            .build();
529        let lower = self.document_name(prefix);
530        let upper = self.document_name(&format!("{}~", prefix));
531        let mut structured_query = StructuredQuery::builder()
532            .from(vec![collection_selector])
533            .order_by(vec![Order::builder()
534                .field(document_name_field.clone())
535                .direction(Direction::Ascending)
536                .build()])
537            .r#where(Filter::CompositeFilter(
538                CompositeFilter::builder()
539                    .op(CompositeFilterOperator::And)
540                    .filters(vec![
541                        Filter::FieldFilter(
542                            FieldFilter::builder()
543                                .field(document_name_field.clone())
544                                .op(FieldFilterOperator::GreaterThanOrEqual)
545                                .value(Value::ReferenceValue(lower))
546                                .build(),
547                        ),
548                        Filter::FieldFilter(
549                            FieldFilter::builder()
550                                .field(document_name_field)
551                                .op(FieldFilterOperator::LessThan)
552                                .value(Value::ReferenceValue(upper))
553                                .build(),
554                        ),
555                    ])
556                    .build(),
557            ))
558            .limit(i32::try_from(limit).unwrap_or(i32::MAX))
559            .build();
560
561        if let Some(state) = cursor_state {
562            structured_query.start_at = Some(
563                Cursor::builder()
564                    .values(vec![Value::ReferenceValue(state.last_document_name)])
565                    .before(false)
566                    .build(),
567            );
568        }
569
570        let query_request = RunQueryRequest::builder()
571            .parent(format!(
572                "projects/{}/databases/{}/documents",
573                self.project_id, self.database_id
574            ))
575            .query_type(QueryType::StructuredQuery(structured_query))
576            .build();
577
578        let query_responses = self
579            .client
580            .run_query(self.database_id.clone(), query_request)
581            .await
582            .map_err(|e| {
583                crate::error::map_cloud_client_error(
584                    e,
585                    "Failed to run Firestore query".to_string(),
586                    Some(prefix.to_string()),
587                )
588            })?;
589
590        let documents = query_responses
591            .iter()
592            .filter_map(|response| response.document.as_ref())
593            .collect::<Vec<_>>();
594        let last_document_name = documents.last().and_then(|document| document.name.clone());
595        let mut items = Vec::with_capacity(documents.len());
596        for document in &documents {
597            let Some(name) = &document.name else {
598                continue;
599            };
600            let Some(key) = name.rsplit('/').next() else {
601                continue;
602            };
603            let kv_doc = self.firestore_to_kv_document(document)?;
604            if self.is_expired(kv_doc.expires_at) {
605                continue;
606            }
607            let value = base64::engine::general_purpose::STANDARD
608                .decode(&kv_doc.value)
609                .into_alien_error()
610                .context(ErrorData::UnexpectedResponseFormat {
611                    provider: "gcp".to_string(),
612                    binding_name: "firestore".to_string(),
613                    field: "value".to_string(),
614                    response_json: serde_json::to_string(document).unwrap_or_default(),
615                })?;
616            items.push((key.to_string(), value));
617        }
618
619        let next_cursor = if documents.len() == limit {
620            last_document_name
621                .map(|last_document_name| {
622                    Self::encode_cursor(&CursorState {
623                        version: 1,
624                        prefix: prefix.to_string(),
625                        last_document_name,
626                    })
627                })
628                .transpose()?
629        } else {
630            None
631        };
632
633        Ok(ScanResult { items, next_cursor })
634    }
635}