Faster preloading with query

abdolence · Oct 7, 2023 · d4db0f3 · d4db0f3
1 parent 60d77ab
commit d4db0f3
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 6 deletions.
diff --git a/src/cache/backends/persistent_backend.rs b/src/cache/backends/persistent_backend.rs
@@ -98,16 +98,27 @@ impl FirestorePersistentCacheBackend {
                     debug!("Preloading {}", collection_path.as_str());
                     let params = if let Some(parent) = &config.parent {
                         db.fluent()
-                            .list()
-                            .from(&config.collection_name)
+                            .select()
+                            .from(config.collection_name.as_str())
                             .parent(parent)
                     } else {
-                        db.fluent().list().from(&config.collection_name)
+                        db.fluent().select().from(config.collection_name.as_str())
                     };
 
-                    let stream = params.page_size(1000).stream_all().await?;
+                    let stream = params.stream_query().await?;
 
                     stream
+                        .enumerate()
+                        .map(|(index, docs)| {
+                            if index > 0 && index % 10000 == 0 {
+                                debug!(
+                                    "Preloading collection `{}`: {} entries loaded",
+                                    collection_path.as_str(),
+                                    index
+                                );
+                            }
+                            docs
+                        })
                         .ready_chunks(100)
                         .for_each(|docs| async move {
                             if let Err(err) = self.write_batch_docs(collection_path, docs) {

diff --git a/src/cache/configuration.rs b/src/cache/configuration.rs
@@ -43,6 +43,7 @@ pub struct FirestoreCacheCollectionConfiguration {
     pub parent: Option<String>,
     pub listener_target: FirestoreListenerTarget,
     pub collection_load_mode: FirestoreCacheCollectionLoadMode,
+    pub indices: Vec<FirestoreCacheIndexConfiguration>,
 }
 
 impl FirestoreCacheCollectionConfiguration {
@@ -60,6 +61,7 @@ impl FirestoreCacheCollectionConfiguration {
             parent: None,
             listener_target,
             collection_load_mode,
+            indices: Vec::new(),
         }
     }
 
@@ -73,6 +75,13 @@ impl FirestoreCacheCollectionConfiguration {
             ..self
         }
     }
+
+    #[inline]
+    pub fn with_index(self, index: FirestoreCacheIndexConfiguration) -> Self {
+        let mut indices = self.indices;
+        indices.push(index);
+        Self { indices, ..self }
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -81,3 +90,31 @@ pub enum FirestoreCacheCollectionLoadMode {
     PreloadAllIfEmpty,
     PreloadNone,
 }
+
+#[derive(Debug, Clone)]
+pub struct FirestoreCacheIndexConfiguration {
+    pub fields: Vec<String>,
+    pub unique: bool,
+}
+
+impl FirestoreCacheIndexConfiguration {
+    #[inline]
+    pub fn new<I>(fields: I) -> Self
+    where
+        I: IntoIterator,
+        I::Item: AsRef<str>,
+    {
+        Self {
+            fields: fields.into_iter().map(|s| s.as_ref().to_string()).collect(),
+            unique: false,
+        }
+    }
+
+    #[inline]
+    pub fn unique(self, value: bool) -> Self {
+        Self {
+            unique: value,
+            ..self
+        }
+    }
+}
diff --git a/src/db/list.rs b/src/db/list.rs
@@ -371,9 +371,10 @@ impl FirestoreDb {
                     );
                     span.in_scope(|| {
                         debug!(
-                            "Listing documents in {:?} took {}ms",
+                            "Listing documents in {:?} took {}ms. Returned: {}",
                             params.collection_id,
-                            listing_duration.num_milliseconds()
+                            listing_duration.num_milliseconds(),
+                            result.documents.len()
                         );
                     });