firebase_rs_sdk/firestore/api/
document.rs

1use std::collections::BTreeMap;
2
3use crate::firestore::api::operations::{self, SetOptions};
4use crate::firestore::api::query::{
5    ConvertedQuery, LimitType, Query, QuerySnapshot, TypedQuerySnapshot,
6};
7use crate::firestore::api::snapshot::{DocumentSnapshot, TypedDocumentSnapshot};
8use crate::firestore::error::{internal_error, FirestoreResult};
9use std::sync::Arc;
10
11use crate::firestore::remote::datastore::{
12    Datastore, HttpDatastore, InMemoryDatastore, TokenProviderArc,
13};
14use crate::firestore::value::FirestoreValue;
15
16use super::{
17    ConvertedCollectionReference, ConvertedDocumentReference, Firestore, FirestoreDataConverter,
18};
19
20#[derive(Clone)]
21pub struct FirestoreClient {
22    firestore: Firestore,
23    datastore: Arc<dyn Datastore>,
24}
25
26impl FirestoreClient {
27    /// Creates a client backed by the supplied datastore implementation.
28    pub fn new(firestore: Firestore, datastore: Arc<dyn Datastore>) -> Self {
29        Self {
30            firestore,
31            datastore,
32        }
33    }
34
35    /// Returns a client that stores documents in memory only.
36    ///
37    /// Useful for tests or demos where persistence/network access is not
38    /// required.
39    pub fn with_in_memory(firestore: Firestore) -> Self {
40        Self::new(firestore, Arc::new(InMemoryDatastore::new()))
41    }
42
43    /// Builds a client that talks to Firestore over the REST endpoints using
44    /// anonymous credentials.
45    pub fn with_http_datastore(firestore: Firestore) -> FirestoreResult<Self> {
46        let datastore = HttpDatastore::from_database_id(firestore.database_id().clone())?;
47        Ok(Self::new(firestore, Arc::new(datastore)))
48    }
49
50    /// Builds an HTTP-backed client that attaches the provided Auth/App Check
51    /// providers to every request.
52    ///
53    /// Pass `None` for `app_check_provider` when App Check is not configured.
54    pub fn with_http_datastore_authenticated(
55        firestore: Firestore,
56        auth_provider: TokenProviderArc,
57        app_check_provider: Option<TokenProviderArc>,
58    ) -> FirestoreResult<Self> {
59        let mut builder = HttpDatastore::builder(firestore.database_id().clone())
60            .with_auth_provider(auth_provider);
61
62        if let Some(provider) = app_check_provider {
63            builder = builder.with_app_check_provider(provider);
64        }
65
66        let datastore = builder.build()?;
67        Ok(Self::new(firestore, Arc::new(datastore)))
68    }
69
70    /// Fetches the document located at `path`.
71    ///
72    /// Returns a snapshot that may or may not contain data depending on whether
73    /// the document exists.
74    pub fn get_doc(&self, path: &str) -> FirestoreResult<DocumentSnapshot> {
75        let key = operations::validate_document_path(path)?;
76        self.datastore.get_document(&key)
77    }
78
79    /// Writes the provided map of fields into the document at `path`.
80    ///
81    /// `options.merge == true` mirrors the JS API but is currently unsupported
82    /// for the HTTP datastore.
83    pub fn set_doc(
84        &self,
85        path: &str,
86        data: BTreeMap<String, FirestoreValue>,
87        options: Option<SetOptions>,
88    ) -> FirestoreResult<()> {
89        let key = operations::validate_document_path(path)?;
90        let encoded = operations::encode_document_data(data)?;
91        let merge = options.unwrap_or_default().merge;
92        self.datastore.set_document(&key, encoded, merge)
93    }
94
95    /// Adds a new document to the collection located at `collection_path` and
96    /// returns the resulting snapshot.
97    pub fn add_doc(
98        &self,
99        collection_path: &str,
100        data: BTreeMap<String, FirestoreValue>,
101    ) -> FirestoreResult<DocumentSnapshot> {
102        let collection = self.firestore.collection(collection_path)?;
103        let doc_ref = collection.doc(None)?;
104        self.set_doc(doc_ref.path().canonical_string().as_str(), data, None)?;
105        self.get_doc(doc_ref.path().canonical_string().as_str())
106    }
107
108    /// Reads a document using the converter attached to a typed reference.
109    pub fn get_doc_with_converter<C>(
110        &self,
111        reference: &ConvertedDocumentReference<C>,
112    ) -> FirestoreResult<TypedDocumentSnapshot<C>>
113    where
114        C: FirestoreDataConverter,
115    {
116        let path = reference.path().canonical_string();
117        let snapshot = self.get_doc(path.as_str())?;
118        let converter = reference.converter();
119        Ok(snapshot.into_typed(converter))
120    }
121
122    /// Executes the provided query and returns its results.
123    pub fn get_docs(&self, query: &Query) -> FirestoreResult<QuerySnapshot> {
124        self.ensure_same_database(query.firestore())?;
125        let definition = query.definition();
126        let mut documents = self.datastore.run_query(&definition)?;
127        if definition.limit_type() == LimitType::Last {
128            documents.reverse();
129        }
130        Ok(QuerySnapshot::new(query.clone(), documents))
131    }
132
133    /// Executes a converted query, producing typed snapshots.
134    pub fn get_docs_with_converter<C>(
135        &self,
136        query: &ConvertedQuery<C>,
137    ) -> FirestoreResult<TypedQuerySnapshot<C>>
138    where
139        C: FirestoreDataConverter,
140    {
141        let snapshot = self.get_docs(query.raw())?;
142        Ok(TypedQuerySnapshot::new(snapshot, query.converter()))
143    }
144
145    /// Writes a typed model to the location referenced by `reference`.
146    pub fn set_doc_with_converter<C>(
147        &self,
148        reference: &ConvertedDocumentReference<C>,
149        data: C::Model,
150        options: Option<SetOptions>,
151    ) -> FirestoreResult<()>
152    where
153        C: FirestoreDataConverter,
154    {
155        let converter = reference.converter();
156        let map = converter.to_map(&data)?;
157        let path = reference.path().canonical_string();
158        self.set_doc(path.as_str(), map, options)
159    }
160
161    /// Creates a document with auto-generated ID using the provided converter.
162    pub fn add_doc_with_converter<C>(
163        &self,
164        collection: &ConvertedCollectionReference<C>,
165        data: C::Model,
166    ) -> FirestoreResult<TypedDocumentSnapshot<C>>
167    where
168        C: FirestoreDataConverter,
169    {
170        let doc_ref = collection.doc(None)?;
171        let converter = doc_ref.converter();
172        let map = converter.to_map(&data)?;
173        let path = doc_ref.path().canonical_string();
174        self.set_doc(path.as_str(), map, None)?;
175        let snapshot = self.get_doc(path.as_str())?;
176        Ok(snapshot.into_typed(converter))
177    }
178
179    fn ensure_same_database(&self, firestore: &Firestore) -> FirestoreResult<()> {
180        if self.firestore.database_id() != firestore.database_id() {
181            return Err(internal_error(
182                "Query targets a different Firestore instance than this client",
183            ));
184        }
185        Ok(())
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::app::api::initialize_app;
193    use crate::app::{FirebaseAppSettings, FirebaseOptions};
194    use crate::firestore::api::get_firestore;
195    use crate::firestore::model::FieldPath;
196    use crate::firestore::value::MapValue;
197    use crate::firestore::value::ValueKind;
198
199    fn unique_settings() -> FirebaseAppSettings {
200        use std::sync::atomic::{AtomicUsize, Ordering};
201        static COUNTER: AtomicUsize = AtomicUsize::new(0);
202        FirebaseAppSettings {
203            name: Some(format!(
204                "firestore-doc-{}",
205                COUNTER.fetch_add(1, Ordering::SeqCst)
206            )),
207            ..Default::default()
208        }
209    }
210
211    fn build_client() -> FirestoreClient {
212        let options = FirebaseOptions {
213            project_id: Some("project".into()),
214            ..Default::default()
215        };
216        let app = initialize_app(options, Some(unique_settings())).unwrap();
217        let firestore = get_firestore(Some(app)).unwrap();
218        FirestoreClient::with_in_memory(Firestore::from_arc(firestore))
219    }
220
221    #[test]
222    fn set_and_get_document() {
223        let client = build_client();
224        let mut data = BTreeMap::new();
225        data.insert("name".to_string(), FirestoreValue::from_string("Ada"));
226        client
227            .set_doc("cities/sf", data.clone(), None)
228            .expect("set doc");
229        let snapshot = client.get_doc("cities/sf").expect("get doc");
230        assert!(snapshot.exists());
231        assert_eq!(
232            snapshot.data().unwrap().get("name"),
233            Some(&FirestoreValue::from_string("Ada"))
234        );
235    }
236
237    #[test]
238    fn query_returns_collection_documents() {
239        let client = build_client();
240        client
241            .set_doc(
242                "cities/sf",
243                BTreeMap::from([("name".into(), FirestoreValue::from_string("San Francisco"))]),
244                None,
245            )
246            .unwrap();
247        client
248            .set_doc(
249                "cities/la",
250                BTreeMap::from([("name".into(), FirestoreValue::from_string("Los Angeles"))]),
251                None,
252            )
253            .unwrap();
254
255        let collection = client.firestore.collection("cities").unwrap();
256        let query = collection.query();
257        let snapshot = client.get_docs(&query).expect("query");
258
259        assert_eq!(snapshot.len(), 2);
260        let ids: Vec<_> = snapshot
261            .documents()
262            .iter()
263            .map(|doc| doc.id().to_string())
264            .collect();
265        assert_eq!(ids, vec!["la", "sf"]);
266    }
267
268    #[derive(Clone, Debug, PartialEq, Eq)]
269    struct Person {
270        first: String,
271        last: String,
272    }
273
274    #[derive(Clone, Debug)]
275    struct PersonConverter;
276
277    impl super::FirestoreDataConverter for PersonConverter {
278        type Model = Person;
279
280        fn to_map(&self, value: &Self::Model) -> FirestoreResult<BTreeMap<String, FirestoreValue>> {
281            let mut map = BTreeMap::new();
282            map.insert(
283                "first".to_string(),
284                FirestoreValue::from_string(&value.first),
285            );
286            map.insert("last".to_string(), FirestoreValue::from_string(&value.last));
287            Ok(map)
288        }
289
290        fn from_map(&self, value: &MapValue) -> FirestoreResult<Self::Model> {
291            let first = match value.fields().get("first").and_then(|v| match v.kind() {
292                ValueKind::String(s) => Some(s.clone()),
293                _ => None,
294            }) {
295                Some(name) => name,
296                None => {
297                    return Err(crate::firestore::error::invalid_argument(
298                        "missing first field",
299                    ))
300                }
301            };
302            let last = match value.fields().get("last").and_then(|v| match v.kind() {
303                ValueKind::String(s) => Some(s.clone()),
304                _ => None,
305            }) {
306                Some(name) => name,
307                None => {
308                    return Err(crate::firestore::error::invalid_argument(
309                        "missing last field",
310                    ))
311                }
312            };
313            Ok(Person { first, last })
314        }
315    }
316
317    #[test]
318    fn typed_set_and_get_document() {
319        let client = build_client();
320        let collection = client.firestore.collection("people").unwrap();
321        let converted = collection.with_converter(PersonConverter);
322        let doc_ref = converted.doc(Some("ada")).unwrap();
323
324        let person = Person {
325            first: "Ada".into(),
326            last: "Lovelace".into(),
327        };
328
329        client
330            .set_doc_with_converter(&doc_ref, person.clone(), None)
331            .expect("typed set");
332
333        let snapshot = client.get_doc_with_converter(&doc_ref).expect("typed get");
334        assert!(snapshot.exists());
335        assert!(snapshot.from_cache());
336        assert!(!snapshot.has_pending_writes());
337
338        let decoded = snapshot.data().expect("converter result").unwrap();
339        assert_eq!(decoded, person);
340    }
341
342    #[test]
343    fn typed_query_returns_converted_results() {
344        let client = build_client();
345        let collection = client.firestore.collection("people").unwrap();
346        let converted = collection.with_converter(PersonConverter);
347
348        let doc_ref = converted.doc(Some("ada")).unwrap();
349        let ada = Person {
350            first: "Ada".into(),
351            last: "Lovelace".into(),
352        };
353        client
354            .set_doc_with_converter(&doc_ref, ada.clone(), None)
355            .expect("set typed doc");
356
357        let query = converted.query();
358        let snapshot = client
359            .get_docs_with_converter(&query)
360            .expect("converted query");
361
362        let docs = snapshot.documents();
363        assert_eq!(docs.len(), 1);
364        let decoded = docs[0].data().expect("converter data").unwrap();
365        assert_eq!(decoded, ada);
366    }
367
368    #[test]
369    fn query_with_filters_and_limit() {
370        use crate::firestore::api::query::{FilterOperator, OrderDirection};
371
372        let client = build_client();
373        let collection = client.firestore.collection("cities").unwrap();
374
375        let mut sf = BTreeMap::new();
376        sf.insert("name".into(), FirestoreValue::from_string("San Francisco"));
377        sf.insert("state".into(), FirestoreValue::from_string("California"));
378        sf.insert("population".into(), FirestoreValue::from_integer(860_000));
379        client.set_doc("cities/sf", sf, None).expect("insert sf");
380
381        let mut la = BTreeMap::new();
382        la.insert("name".into(), FirestoreValue::from_string("Los Angeles"));
383        la.insert("state".into(), FirestoreValue::from_string("California"));
384        la.insert("population".into(), FirestoreValue::from_integer(3_980_000));
385        client.set_doc("cities/la", la, None).expect("insert la");
386
387        let query = collection
388            .query()
389            .where_field(
390                FieldPath::from_dot_separated("state").unwrap(),
391                FilterOperator::Equal,
392                FirestoreValue::from_string("California"),
393            )
394            .unwrap()
395            .order_by(
396                FieldPath::from_dot_separated("population").unwrap(),
397                OrderDirection::Descending,
398            )
399            .unwrap()
400            .limit(1)
401            .unwrap();
402
403        let snapshot = client.get_docs(&query).expect("filtered query");
404        assert_eq!(snapshot.len(), 1);
405        let doc = &snapshot.documents()[0];
406        assert_eq!(doc.id(), "la");
407    }
408}