prog

chroma-core · Sep 24, 2024 · f1748e2 · f1748e2
1 parent ef540dd
commit f1748e2
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 2 deletions.
diff --git a/rust/worker/src/execution/orchestration/count.rs b/rust/worker/src/execution/orchestration/count.rs
@@ -150,8 +150,19 @@ impl CountQueryOrchestrator {
             }
         };
 
+        // If the collection version does not match the request version then we terminate with an error
+        if collection.version as u32 != self.collection_version {
+            terminate_with_error(
+                self.result_channel.take(),
+                Box::new(CountQueryOrchestratorError::CollectionVersionMismatch),
+                ctx,
+            );
+            return;
+        }
+
         self.record_segment = Some(record_segment);
         self.collection = Some(collection);
+        self.pull_logs(ctx).await;
     }
 
     // shared
@@ -319,7 +330,6 @@ impl Component for CountQueryOrchestrator {
 
     async fn on_start(&mut self, ctx: &crate::system::ComponentContext<Self>) -> () {
         self.start(ctx).await;
-        self.pull_logs(ctx).await;
     }
 }
 

diff --git a/rust/worker/src/execution/orchestration/get_vectors.rs b/rust/worker/src/execution/orchestration/get_vectors.rs
@@ -252,6 +252,16 @@ impl Component for GetVectorsOrchestrator {
             }
         };
 
+        // If the collection version does not match the request version then we terminate with an error
+        if collection.version as u32 != self.collection_version {
+            terminate_with_error(
+                self.result_channel.take(),
+                Box::new(GetVectorsError::CollectionVersionMismatch),
+                ctx,
+            );
+            return;
+        }
+
         let record_segment =
             match get_record_segment_by_collection_id(self.sysdb.clone(), collection_id).await {
                 Ok(segment) => segment,

diff --git a/rust/worker/src/execution/orchestration/hnsw.rs b/rust/worker/src/execution/orchestration/hnsw.rs
@@ -583,6 +583,16 @@ impl Component for HnswQueryOrchestrator {
             }
         };
 
+        // If the collection version does not match the request version then we terminate with an error
+        if collection.version as u32 != self.collection_version {
+            terminate_with_error(
+                self.result_channel.take(),
+                Box::new(HnswSegmentQueryError::CollectionVersionMismatch),
+                ctx,
+            );
+            return;
+        }
+
         // If segment is uninitialized and dimension is not set then we assume
         // that this is a query before any add so return empty response.
         if hnsw_segment.file_path.len() <= 0 && collection.dimension.is_none() {

diff --git a/rust/worker/src/execution/orchestration/metadata.rs b/rust/worker/src/execution/orchestration/metadata.rs
@@ -194,8 +194,19 @@ impl MetadataQueryOrchestrator {
             }
         };
 
+        // If the collection version does not match the version in the request, return an error
+        if collection.version as u32 != self.collection_version {
+            terminate_with_error(
+                self.result_channel.take(),
+                Box::new(MetadataQueryOrchestratorError::CollectionVersionMismatch),
+                ctx,
+            );
+            return;
+        }
+
         self.record_segment = Some(record_segment);
         self.collection = Some(collection);
+        self.pull_logs(ctx).await;
     }
 
     async fn pull_logs(&mut self, ctx: &ComponentContext<Self>) {
@@ -399,7 +410,6 @@ impl Component for MetadataQueryOrchestrator {
 
     async fn on_start(&mut self, ctx: &crate::system::ComponentContext<Self>) -> () {
         self.start(ctx).await;
-        self.pull_logs(ctx).await;
     }
 }