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#[derive(Debug, Clone, Serialize, Deserialize)]
20struct KvDocument {
21 value: String, created_at: DateTime<Utc>,
23 expires_at: Option<DateTime<Utc>>, }
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
34pub 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 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 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 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 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 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 if self.is_expired(kv_doc.expires_at) {
344 return Ok(None); }
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 match &e.error {
361 Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. }) => {
362 Ok(None) }
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 match &e.error {
406 Some(alien_client_core::ErrorData::RemoteResourceConflict { .. }) => {
407 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 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) }
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}