[ENH] Add CMEK to collection schema (#5975)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - N/A
- New functionality
  - Add CMEK to collection schema, and with builder methods
  - Propagate CMEK from schema to downstream

## Test plan

_How are these changes tested?_

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: Sicheng-Pan

clients/new-js/packages/chromadb/src/api/types.gen.ts

@@ -414,6 +414,12 @@ export type RawWhereFields = {
  * This represents the server-side schema structure used for index management
  */
 export type Schema = {
+    /**
+     * Customer-managed encryption key for collection data
+     */
+    cmek?: {
+        [key: string]: unknown;
+    } | null;
     /**
      * Default index configurations for each value type
      */

rust/chroma/src/types.rs

@@ -20,6 +20,7 @@ pub use chroma_types::BoolInvertedIndexConfig;
 pub use chroma_types::BoolInvertedIndexType;
 pub use chroma_types::BoolValueType;
 pub use chroma_types::BooleanOperator;
+pub use chroma_types::Cmek;
 pub use chroma_types::Collection;
 pub use chroma_types::CollectionConfiguration;
 pub use chroma_types::CollectionConversionError;

rust/frontend/src/impls/service_based_frontend.rs

@@ -182,20 +182,6 @@ impl ServiceBasedFrontend {
             .collection)
     }
 
-    async fn get_collection_dimension(
-        &mut self,
-        collection_id: CollectionUuid,
-    ) -> Result<Option<u32>, GetCollectionError> {
-        Ok(self
-            .collections_with_segments_provider
-            .get_collection_with_segments(collection_id)
-            .await
-            .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?
-            .collection
-            .dimension
-            .map(|dim| dim as u32))
-    }
-
     async fn set_collection_dimension(
         &mut self,
         collection_id: CollectionUuid,
@@ -219,10 +205,11 @@ impl ServiceBasedFrontend {
         option_embeddings: Option<&Vec<Embedding>>,
         update_if_not_present: bool,
         read_length: F,
-    ) -> Result<(), ValidationError>
+    ) -> Result<Collection, ValidationError>
     where
         F: Fn(&Embedding) -> Option<usize>,
     {
+        let collection = self.get_cached_collection(collection_id).await?;
         if let Some(embeddings) = option_embeddings {
             let emb_dims = embeddings
                 .iter()
@@ -237,28 +224,23 @@ impl ServiceBasedFrontend {
                 low as u32
             } else {
                 // No embedding to check, return
-                return Ok(());
+                return Ok(collection);
             };
-            match self.get_collection_dimension(collection_id).await {
-                Ok(Some(expected_dim)) => {
+            match collection.dimension.map(|dim| dim as u32) {
+                Some(expected_dim) => {
                     if expected_dim != emb_dim {
                         return Err(ValidationError::DimensionMismatch(expected_dim, emb_dim));
                     }
-
-                    Ok(())
                 }
-                Ok(None) => {
+                None => {
                     if update_if_not_present {
                         self.set_collection_dimension(collection_id, emb_dim)
                             .await?;
                     }
-                    Ok(())
                 }
-                Err(err) => Err(err.into()),
-            }
-        } else {
-            Ok(())
+            };
         }
+        Ok(collection)
     }
 
     pub async fn reset(&mut self) -> Result<ResetResponse, ResetError> {
@@ -830,31 +812,31 @@ impl ServiceBasedFrontend {
             ..
         }: AddCollectionRecordsRequest,
     ) -> Result<AddCollectionRecordsResponse, AddCollectionRecordsError> {
-        self.validate_embedding(
-            collection_id,
-            Some(&embeddings),
-            true,
-            |embedding: &Vec<f32>| Some(embedding.len()),
-        )
-        .await
-        .map_err(|err| err.boxed())?;
+        let collection = self
+            .validate_embedding(
+                collection_id,
+                Some(&embeddings),
+                true,
+                |embedding: &Vec<f32>| Some(embedding.len()),
+            )
+            .await
+            .map_err(|err| err.boxed())?;
 
         let embeddings = Some(embeddings.into_iter().map(Some).collect());
 
         let (records, log_size_bytes) =
             to_records(ids, embeddings, documents, uris, metadatas, Operation::Add)
                 .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?;
 
-        // TODO: Extract CMEK from collection metadata
-        // For now, pass None until collection-level CMEK storage is implemented
-        let cmek = None;
-
         let retries = Arc::new(AtomicUsize::new(0));
         let add_to_retry = || {
             let mut self_clone = self.clone();
             let records_clone = records.clone();
             let tenant_id_clone = tenant_id.clone();
-            let cmek_clone = cmek.clone();
+            let cmek_clone = collection
+                .schema
+                .as_ref()
+                .and_then(|schema| schema.cmek.clone());
             async move {
                 self_clone
                     .retryable_push_logs(&tenant_id_clone, collection_id, records_clone, cmek_clone)
@@ -922,11 +904,12 @@ impl ServiceBasedFrontend {
             ..
         }: UpdateCollectionRecordsRequest,
     ) -> Result<UpdateCollectionRecordsResponse, UpdateCollectionRecordsError> {
-        self.validate_embedding(collection_id, embeddings.as_ref(), true, |embedding| {
-            embedding.as_ref().map(|emb| emb.len())
-        })
-        .await
-        .map_err(|err| err.boxed())?;
+        let collection = self
+            .validate_embedding(collection_id, embeddings.as_ref(), true, |embedding| {
+                embedding.as_ref().map(|emb| emb.len())
+            })
+            .await
+            .map_err(|err| err.boxed())?;
 
         let (records, log_size_bytes) = to_records(
             ids,
@@ -938,16 +921,15 @@ impl ServiceBasedFrontend {
         )
         .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?;
 
-        // TODO: Extract CMEK from collection metadata
-        // For now, pass None until collection-level CMEK storage is implemented
-        let cmek = None;
-
         let retries = Arc::new(AtomicUsize::new(0));
         let add_to_retry = || {
             let mut self_clone = self.clone();
             let records_clone = records.clone();
             let tenant_id_clone = tenant_id.clone();
-            let cmek_clone = cmek.clone();
+            let cmek_clone = collection
+                .schema
+                .as_ref()
+                .and_then(|schema| schema.cmek.clone());
             async move {
                 self_clone
                     .retryable_push_logs(&tenant_id_clone, collection_id, records_clone, cmek_clone)
@@ -1015,14 +997,15 @@ impl ServiceBasedFrontend {
             ..
         }: UpsertCollectionRecordsRequest,
     ) -> Result<UpsertCollectionRecordsResponse, UpsertCollectionRecordsError> {
-        self.validate_embedding(
-            collection_id,
-            Some(&embeddings),
-            true,
-            |embedding: &Vec<f32>| Some(embedding.len()),
-        )
-        .await
-        .map_err(|err| err.boxed())?;
+        let collection = self
+            .validate_embedding(
+                collection_id,
+                Some(&embeddings),
+                true,
+                |embedding: &Vec<f32>| Some(embedding.len()),
+            )
+            .await
+            .map_err(|err| err.boxed())?;
 
         let embeddings = Some(embeddings.into_iter().map(Some).collect());
 
@@ -1036,16 +1019,15 @@ impl ServiceBasedFrontend {
         )
         .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?;
 
-        // TODO: Extract CMEK from collection metadata
-        // For now, pass None until collection-level CMEK storage is implemented
-        let cmek = None;
-
         let retries = Arc::new(AtomicUsize::new(0));
         let add_to_retry = || {
             let mut self_clone = self.clone();
             let records_clone = records.clone();
             let tenant_id_clone = tenant_id.clone();
-            let cmek_clone = cmek.clone();
+            let cmek_clone = collection
+                .schema
+                .as_ref()
+                .and_then(|schema| schema.cmek.clone());
             async move {
                 self_clone
                     .retryable_push_logs(&tenant_id_clone, collection_id, records_clone, cmek_clone)
@@ -1229,10 +1211,13 @@ impl ServiceBasedFrontend {
 
             let log_size_bytes = records.iter().map(OperationRecord::size_bytes).sum();
 
-            // TODO: Extract CMEK from collection metadata
-            // For now, pass None until collection-level CMEK storage is implemented
-            self.log_client
-                .push_logs(&tenant_id, collection_id, records, None)
+            let cmek = self
+                .get_cached_collection(collection_id)
+                .await
+                .map_err(|err| DeleteCollectionRecordsError::Internal(err.boxed()))?
+                .schema
+                .and_then(|schema| schema.cmek.clone());
+            self.retryable_push_logs(&tenant_id, collection_id, records, cmek)
                 .await
                 .map_err(|err| {
                     if err.code() == ErrorCodes::Unavailable {
@@ -1987,10 +1972,8 @@ impl ServiceBasedFrontend {
         collection_id: CollectionUuid,
         _attached_function_id: chroma_types::AttachedFunctionUuid,
     ) -> Result<(), chroma_types::AttachFunctionError> {
-        let embedding_dim = self
-            .get_collection_dimension(collection_id)
-            .await?
-            .unwrap_or(1);
+        let collection = self.get_cached_collection(collection_id).await?;
+        let embedding_dim = collection.dimension.unwrap_or(1);
         let fake_embedding = vec![0.0; embedding_dim as usize];
         // TODO(tanujnay112): Make this either a configurable or better yet a separate
         // RPC to the logs service.
@@ -2006,10 +1989,9 @@ impl ServiceBasedFrontend {
             };
             num_fake_logs
         ];
-        // TODO: Extract CMEK from collection metadata
-        // For now, pass None until collection-level CMEK storage is implemented
-        self.log_client
-            .push_logs(&tenant, collection_id, logs, None)
+
+        let cmek = collection.schema.and_then(|schema| schema.cmek.clone());
+        self.retryable_push_logs(&tenant, collection_id, logs, cmek)
             .await?;
         Ok(())
     }