Make wrapper type for batching

TabbyML · May 30, 2024 · 8caf3a3 · 8caf3a3
1 parent 8f176bf
commit 8caf3a3
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git a/crates/tabby-scheduler/src/code/cache.rs b/crates/tabby-scheduler/src/code/cache.rs
@@ -65,7 +65,20 @@ pub struct CacheStore {
     code: CodeIntelligence,
 }
 
-pub const INDEX_VERSION: &str = "1";
+const INDEX_ALGORITHM_VERSION: &str = "1";
+
+#[derive(Default)]
+pub struct IndexBatch {
+    batch: Batch<String, String>,
+}
+
+impl IndexBatch {
+    pub fn set_indexed(&mut self, file_id: String) {
+        self.batch
+            .set(&file_id, &INDEX_ALGORITHM_VERSION.into())
+            .expect("Failed to write to batch");
+    }
+}
 
 impl CacheStore {
     pub fn new(path: PathBuf) -> Self {
@@ -89,7 +102,10 @@ impl CacheStore {
             .index_bucket()
             .get(&key)
             .expect("Failed to read index bucket");
-        (key, indexed.is_some_and(|indexed| indexed == INDEX_VERSION))
+        (
+            key,
+            indexed.is_some_and(|indexed| indexed == INDEX_ALGORITHM_VERSION),
+        )
     }
 
     pub fn clear_indexed(&self) {
@@ -98,9 +114,9 @@ impl CacheStore {
             .expect("Failed to clear indexed files bucket");
     }
 
-    pub fn apply_indexed(&self, batch: Batch<String, String>) {
+    pub fn apply_indexed(&self, batch: IndexBatch) {
         self.index_bucket()
-            .batch(batch)
+            .batch(batch.batch)
             .expect("Failed to commit batched index update")
     }
 

diff --git a/crates/tabby-scheduler/src/code/index.rs b/crates/tabby-scheduler/src/code/index.rs
@@ -1,13 +1,12 @@
 use ignore::Walk;
-use kv::Batch;
 use tabby_common::{
     api::code::CodeSearchDocument, config::RepositoryConfig, index::CodeSearchSchema, path,
 };
 use tantivy::{Index, IndexWriter, Term};
 use tracing::warn;
 
 use super::{
-    cache::{CacheStore, INDEX_VERSION},
+    cache::{CacheStore, IndexBatch},
     intelligence::{CodeIntelligence, SourceCode},
 };
 use crate::tantivy_utils;
@@ -43,7 +42,7 @@ fn add_changed_documents(cache: &mut CacheStore, repository: &RepositoryConfig,
         .expect("Failed to create index writer");
 
     let intelligence = CodeIntelligence::default();
-    let mut indexed_files_batch = Batch::new();
+    let mut indexed_files_batch = IndexBatch::default();
     for file in Walk::new(repository.dir()) {
         let file = match file {
             Ok(file) => file,
@@ -87,9 +86,7 @@ fn add_changed_documents(cache: &mut CacheStore, repository: &RepositoryConfig,
                 .expect("Failed to add document");
         }
 
-        indexed_files_batch
-            .set(&file_id, &INDEX_VERSION.to_string())
-            .expect("Failed to mark file as indexed");
+        indexed_files_batch.set_indexed(file_id);
     }
 
     // Commit updating indexed documents