Skip to main content

arche_firestore/
lib.rs

1use anyhow::Context;
2use firestore::{
3    FirestoreDb, FirestoreDbOptions, FirestoreDocument, FirestoreValue, FirestoreWritePrecondition,
4    errors::FirestoreError,
5};
6use gcloud_sdk::{BoxSource, Source, Token, TokenSourceType};
7use serde::{Serialize, de::DeserializeOwned};
8use std::sync::Arc;
9
10pub fn deserialize_stored_document<T>(document: &FirestoreDocument) -> Result<T, FirestoreError>
11where
12    T: DeserializeOwned,
13{
14    // Firestore's typed decoder injects `_firestore_*` document metadata into the
15    // Serde map. Domain records are intentionally strict, so decode only persisted fields.
16    let fields = FirestoreValue::from_map(
17        document
18            .fields
19            .iter()
20            .map(|(name, value)| (name, FirestoreValue::from(value.clone()))),
21    );
22    T::deserialize(fields).map_err(|error| match error {
23        FirestoreError::DeserializeError(error) => {
24            FirestoreError::DeserializeError(error.with_document_path(document.name.clone()))
25        }
26        error => error,
27    })
28}
29
30pub async fn get_stored_obj_at_if_exists<T>(
31    db: &FirestoreDb,
32    parent: &str,
33    collection: &str,
34    document_id: &str,
35) -> Result<Option<T>, FirestoreError>
36where
37    T: DeserializeOwned,
38{
39    db.fluent()
40        .select()
41        .by_id_in(collection)
42        .parent(parent)
43        .one(document_id)
44        .await?
45        .map(|document| deserialize_stored_document(&document))
46        .transpose()
47}
48
49#[derive(Clone)]
50pub struct Db(Arc<FirestoreDb>);
51
52pub enum Credentials {
53    ApplicationDefault,
54    AccessToken(String),
55}
56
57pub struct DatabaseOptions {
58    pub project_id: Option<String>,
59    pub database_id: String,
60    pub credentials: Credentials,
61    pub connect_timeout: std::time::Duration,
62}
63impl Default for DatabaseOptions {
64    fn default() -> Self {
65        Self {
66            project_id: None,
67            database_id: "(default)".into(),
68            credentials: Credentials::ApplicationDefault,
69            connect_timeout: std::time::Duration::from_secs(30),
70        }
71    }
72}
73pub struct DatabaseConfig(DatabaseOptions);
74impl DatabaseOptions {
75    pub fn from_env() -> Self {
76        Self {
77            project_id: std::env::var("GOOGLE_CLOUD_PROJECT").ok(),
78            database_id: std::env::var("FIRESTORE_DATABASE_ID")
79                .unwrap_or_else(|_| "(default)".into()),
80            ..Self::default()
81        }
82    }
83    pub fn validate(self) -> anyhow::Result<DatabaseConfig> {
84        if let Some(project) = &self.project_id {
85            anyhow::ensure!(
86                !project.is_empty()
87                    && project
88                        .bytes()
89                        .all(|c| c.is_ascii_alphanumeric() || c == b'-'),
90                "invalid Firestore project id"
91            );
92        }
93        anyhow::ensure!(
94            self.database_id == "(default)"
95                || (!self.database_id.is_empty()
96                    && self
97                        .database_id
98                        .bytes()
99                        .all(|c| c.is_ascii_alphanumeric() || c == b'-')),
100            "invalid Firestore database id"
101        );
102        anyhow::ensure!(
103            !self.connect_timeout.is_zero(),
104            "database connection timeout must be positive"
105        );
106        if let Credentials::AccessToken(token) = &self.credentials {
107            anyhow::ensure!(!token.trim().is_empty(), "access token is empty");
108        }
109        Ok(DatabaseConfig(self))
110    }
111}
112impl DatabaseConfig {
113    pub async fn connect(self) -> anyhow::Result<Db> {
114        let timeout = self.0.connect_timeout;
115        tokio::time::timeout(timeout, async {
116            let options = match self.0.project_id {
117                Some(project) => FirestoreDbOptions::new(project),
118                None => FirestoreDbOptions::for_default_project_id()
119                    .await
120                    .context("GCP project id not detected; set GOOGLE_CLOUD_PROJECT")?,
121            }
122            .with_database_id(self.0.database_id);
123            let db = match self.0.credentials {
124                Credentials::ApplicationDefault => FirestoreDb::with_options(options).await?,
125                Credentials::AccessToken(access_token) => {
126                    let source: BoxSource = Box::new(StaticAccessTokenSource { access_token });
127                    FirestoreDb::with_options_token_source(
128                        options,
129                        gcloud_sdk::GCP_DEFAULT_SCOPES.clone(),
130                        TokenSourceType::ExternalSource(source),
131                    )
132                    .await?
133                }
134            };
135            Ok::<_, anyhow::Error>(Db::from_firestore(db))
136        })
137        .await
138        .context("Firestore connection deadline exceeded")?
139    }
140}
141impl Db {
142    pub fn from_firestore(db: FirestoreDb) -> Self {
143        Self(Arc::new(db))
144    }
145    pub fn shared(&self) -> Arc<FirestoreDb> {
146        self.0.clone()
147    }
148    pub fn inner(&self) -> &FirestoreDb {
149        &self.0
150    }
151}
152
153struct StaticAccessTokenSource {
154    access_token: String,
155}
156
157#[async_trait::async_trait]
158impl Source for StaticAccessTokenSource {
159    async fn token(&self) -> gcloud_sdk::error::Result<Token> {
160        Ok(Token::new(
161            "Bearer".to_string(),
162            secret_vault_value::SecretValue::from(self.access_token.clone()),
163            firestore::FirestoreInstant::now()
164                .checked_add(firestore::jiff::SignedDuration::from_secs(30 * 60))
165                .expect("valid token expiry"),
166        ))
167    }
168}
169
170pub async fn load_optional_typed<T: DeserializeOwned>(
171    db: &FirestoreDb,
172    collection: &str,
173    doc: &str,
174) -> anyhow::Result<Option<T>> {
175    load_optional_typed_at(db, db.get_documents_path(), collection, doc).await
176}
177
178pub async fn load_optional_typed_at<T: DeserializeOwned>(
179    db: &FirestoreDb,
180    parent: &str,
181    collection: &str,
182    doc: &str,
183) -> anyhow::Result<Option<T>> {
184    get_stored_obj_at_if_exists::<T>(db, parent, collection, doc)
185        .await
186        .with_context(|| format!("loading document {parent}/{collection}/{doc}"))
187}
188
189pub async fn create_typed<T>(
190    db: &FirestoreDb,
191    collection: &str,
192    doc: &str,
193    value: &T,
194) -> anyhow::Result<()>
195where
196    T: Serialize + DeserializeOwned + Sync + Send,
197{
198    create_typed_at(db, db.get_documents_path(), collection, doc, value).await
199}
200
201pub async fn create_typed_at<T>(
202    db: &FirestoreDb,
203    parent: &str,
204    collection: &str,
205    doc: &str,
206    value: &T,
207) -> anyhow::Result<()>
208where
209    T: Serialize + DeserializeOwned + Sync + Send,
210{
211    db.fluent()
212        .update()
213        .in_col(collection)
214        .precondition(FirestoreWritePrecondition::Exists(false))
215        .document_id(doc)
216        .parent(parent)
217        .object(value)
218        .execute::<()>()
219        .await?;
220
221    Ok(())
222}
223
224pub fn normalized_text(value: Option<&String>) -> Option<&str> {
225    value
226        .map(String::as_str)
227        .map(str::trim)
228        .filter(|value| !value.is_empty())
229}
230
231pub fn is_firestore_data_conflict(error: &anyhow::Error) -> bool {
232    error
233        .downcast_ref::<FirestoreError>()
234        .is_some_and(|error| matches!(error, FirestoreError::DataConflictError(_)))
235}
236
237pub fn should_retry_bootstrap_conflict(
238    error: &anyhow::Error,
239    doc_path: &str,
240    attempt: usize,
241) -> bool {
242    if !is_firestore_data_conflict(error) || attempt + 1 >= CONFIG_BOOTSTRAP_MAX_RETRIES {
243        return false;
244    }
245
246    tracing::warn!(
247        attempt = attempt + 1,
248        doc_path,
249        "Firestore bootstrap write conflicted; reloading and retrying"
250    );
251    true
252}
253
254pub const CONFIG_BOOTSTRAP_MAX_RETRIES: usize = 8;