[ENH] Add GCS client to storage (#5869)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - N/A
- New functionality
  - Add GCS support with `object_store` crate

## 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

Cargo.lock

@@ -7,6 +7,9 @@ members = ["rust/benchmark", "rust/blockstore", "rust/cache", "rust/chroma", "ru
 anyhow = "1.0"
 arrow = "55.1"
 async-trait = "0.1"
+aws-config = { version = "1.5", features = ["behavior-version-latest"] }
+aws-sdk-s3 = "1.63"
+aws-smithy-types = "1.2"
 axum = { version = "0.8", features = ["macros"] }
 bytes = "1.10"
 chrono = { version = "0.4", features = ["serde"] }
@@ -16,6 +19,7 @@ flatbuffers = "25.2.10"
 futures = "0.3"
 futures-core = "0.3"
 http-body-util = "0.1.3"
+object_store = { version = "0.12", features = ["gcp"] }
 lazy_static = { version = "1.4" }
 lexical-core = "1.0"
 murmur3 = "0.5.2"

Cargo.toml

@@ -7,18 +7,18 @@ edition = "2021"
 path = "src/lib.rs"
 
 [dependencies]
-bytes = "1.8.0"
-aws-sdk-s3 = "1.63"
-aws-smithy-types = "1.2"
-aws-config = { version = "1.5", features = ["behavior-version-latest"] }
-object_store = { version = "0.11", features = ["aws"] }
-
 async-trait = { workspace = true }
+aws-sdk-s3 = { workspace = true }
+aws-smithy-types = { workspace = true }
+aws-config = { workspace = true }
+bytes = { workspace = true }
 futures = { workspace = true }
+object_store = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true }
+serde_json = { workspace = true }
 tempfile = { workspace = true }
-thiserror.workspace = true
+thiserror = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
 opentelemetry = { workspace = true }

rust/storage/Cargo.toml

@@ -4,12 +4,17 @@ use serde::{Deserialize, Serialize};
 /// The configuration for the chosen storage.
 /// # Options
 /// - S3: The configuration for the s3 storage.
+/// - Object: The configuration for the object storage.
+/// - Local: The configuration for local filesystem storage.
+/// - AdmissionControlledS3: S3 with rate limiting and request coalescing.
 /// # Notes
 /// See config.rs in the root of the worker crate for an example of how to use
 /// config files to configure the worker.
 pub enum StorageConfig {
     #[serde(alias = "s3")]
     S3(S3StorageConfig),
+    #[serde(alias = "object")]
+    Object(ObjectStorageConfig),
     #[serde(alias = "local")]
     Local(LocalStorageConfig),
     #[serde(alias = "admissioncontrolleds3")]
@@ -156,3 +161,80 @@ impl Default for RateLimitingConfig {
         RateLimitingConfig::CountBasedPolicy(CountBasedPolicyConfig::default())
     }
 }
+
+#[derive(Deserialize, Debug, Clone, Serialize)]
+pub enum ObjectStorageProvider {
+    /// GCS uses Application Default Credentials (ADC) automatically
+    GCS,
+}
+
+#[derive(Deserialize, Debug, Clone, Serialize)]
+/// The configuration for the ObjectStorage type
+/// # Fields
+/// - bucket: The name of the bucket to use.
+/// - connect_timeout_ms: Connection timeout in milliseconds.
+/// - download_part_size_bytes: Size of each part for parallel range downloads.
+/// - provider: Which backend to use for storage.
+/// - request_retry_count: Number of retry attempts for failed requests.
+/// - request_timeout_ms: Request timeout in milliseconds.
+/// - upload_part_size_bytes: Size of each part in multipart uploads.
+pub struct ObjectStorageConfig {
+    #[serde(default = "ObjectStorageConfig::default_bucket")]
+    pub bucket: String,
+    #[serde(default = "ObjectStorageConfig::default_connect_timeout_ms")]
+    pub connect_timeout_ms: u64,
+    #[serde(default = "ObjectStorageConfig::default_download_part_size_bytes")]
+    pub download_part_size_bytes: u64,
+    #[serde(default = "ObjectStorageConfig::default_provider")]
+    pub provider: ObjectStorageProvider,
+    #[serde(default = "ObjectStorageConfig::default_request_retry_count")]
+    pub request_retry_count: usize,
+    #[serde(default = "ObjectStorageConfig::default_request_timeout_ms")]
+    pub request_timeout_ms: u64,
+    #[serde(default = "ObjectStorageConfig::default_upload_part_size_bytes")]
+    pub upload_part_size_bytes: u64,
+}
+
+impl ObjectStorageConfig {
+    fn default_bucket() -> String {
+        "chroma-storage".to_string()
+    }
+
+    fn default_connect_timeout_ms() -> u64 {
+        5000
+    }
+
+    fn default_download_part_size_bytes() -> u64 {
+        8 * 1024 * 1024 // 8 MB
+    }
+
+    fn default_provider() -> ObjectStorageProvider {
+        ObjectStorageProvider::GCS
+    }
+
+    fn default_request_retry_count() -> usize {
+        3
+    }
+
+    fn default_request_timeout_ms() -> u64 {
+        60000
+    }
+
+    fn default_upload_part_size_bytes() -> u64 {
+        512 * 1024 * 1024 // 512 MB
+    }
+}
+
+impl Default for ObjectStorageConfig {
+    fn default() -> Self {
+        ObjectStorageConfig {
+            bucket: Self::default_bucket(),
+            connect_timeout_ms: Self::default_connect_timeout_ms(),
+            download_part_size_bytes: Self::default_download_part_size_bytes(),
+            provider: Self::default_provider(),
+            request_retry_count: Self::default_request_retry_count(),
+            request_timeout_ms: Self::default_request_timeout_ms(),
+            upload_part_size_bytes: Self::default_upload_part_size_bytes(),
+        }
+    }
+}

rust/storage/src/config.rs

@@ -9,6 +9,7 @@ use chroma_error::{ChromaError, ErrorCodes};
 pub mod admissioncontrolleds3;
 pub mod config;
 pub mod local;
+pub mod object_storage;
 pub mod s3;
 pub mod stream;
 use local::LocalStorage;
@@ -173,7 +174,7 @@ impl ChromaError for StorageError {
             StorageError::PermissionDenied { .. } => ErrorCodes::PermissionDenied,
             StorageError::Unauthenticated { .. } => ErrorCodes::Unauthenticated,
             StorageError::UnknownConfigurationKey { .. } => ErrorCodes::InvalidArgument,
-            StorageError::Backoff { .. } => ErrorCodes::ResourceExhausted,
+            StorageError::Backoff => ErrorCodes::ResourceExhausted,
             StorageError::CallbackError { .. } => ErrorCodes::Internal,
         }
     }
@@ -193,8 +194,6 @@ pub enum PathError {
     },
 }
 
-// END BORROWED CODE
-
 #[derive(Error, Debug)]
 pub enum StorageConfigError {
     #[error("Invalid storage config")]
@@ -207,6 +206,7 @@ pub enum StorageConfigError {
 #[allow(clippy::large_enum_variant)]
 pub enum Storage {
     S3(s3::S3Storage),
+    Object(object_storage::ObjectStorage),
     Local(local::LocalStorage),
     AdmissionControlledS3(admissioncontrolleds3::AdmissionControlledS3Storage),
 }
@@ -215,6 +215,7 @@ impl std::fmt::Debug for Storage {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Storage::S3(_) => f.debug_tuple("S3").finish(),
+            Storage::Object(_) => f.debug_tuple("Object").finish(),
             Storage::Local(_) => f.debug_tuple("Local").finish(),
             Storage::AdmissionControlledS3(_) => f.debug_tuple("AdmissionControlledS3").finish(),
         }
@@ -235,6 +236,7 @@ impl Storage {
     pub fn bucket_name(&self) -> Option<&str> {
         match self {
             Storage::S3(s3) => Some(&s3.bucket),
+            Storage::Object(obj) => Some(&obj.bucket),
             Storage::AdmissionControlledS3(ac_s3) => Some(&ac_s3.storage.bucket),
             Storage::Local(_) => None,
         }
@@ -243,6 +245,10 @@ impl Storage {
     pub async fn get(&self, key: &str, options: GetOptions) -> Result<Arc<Vec<u8>>, StorageError> {
         match self {
             Storage::S3(s3) => s3.get(key, options).await,
+            Storage::Object(obj) => obj
+                .get(key, options)
+                .await
+                .map(|(bytes, _)| Vec::from(bytes).into()),
             Storage::Local(local) => local.get(key).await,
             Storage::AdmissionControlledS3(admission_controlled_storage) => {
                 admission_controlled_storage.get(key, options).await
@@ -267,6 +273,11 @@ impl Storage {
                 let fetch_result = fetch_fn(Ok(res.0)).await?;
                 Ok((fetch_result, res.1))
             }
+            Storage::Object(obj) => {
+                let (bytes, etag) = obj.get(key, options).await?;
+                let fetch_result = fetch_fn(Ok(Vec::from(bytes).into())).await?;
+                Ok((fetch_result, Some(etag)))
+            }
             Storage::Local(local) => {
                 let res = local.get_with_e_tag(key).await?;
                 let fetch_result = fetch_fn(Ok(res.0)).await?;
@@ -318,6 +329,7 @@ impl Storage {
     {
         match self {
             Storage::S3(_) => self.fetch_batch_generic(keys, options, fetch_fn).await,
+            Storage::Object(_) => self.fetch_batch_generic(keys, options, fetch_fn).await,
             Storage::Local(_) => self.fetch_batch_generic(keys, options, fetch_fn).await,
             Storage::AdmissionControlledS3(admission_controlled_storage) => {
                 admission_controlled_storage
@@ -334,6 +346,10 @@ impl Storage {
     ) -> Result<(Arc<Vec<u8>>, Option<ETag>), StorageError> {
         match self {
             Storage::S3(s3) => s3.get_with_e_tag(key).await,
+            Storage::Object(obj) => {
+                let (bytes, etag) = obj.get(key, options).await?;
+                Ok((Vec::from(bytes).into(), Some(etag)))
+            }
             Storage::Local(local) => local.get_with_e_tag(key).await,
             Storage::AdmissionControlledS3(admission_controlled_storage) => {
                 admission_controlled_storage
@@ -350,6 +366,7 @@ impl Storage {
     pub async fn confirm_same(&self, key: &str, e_tag: &ETag) -> Result<bool, StorageError> {
         match self {
             Storage::S3(s3) => s3.confirm_same(key, e_tag).await,
+            Storage::Object(obj) => obj.confirm_same(key, e_tag).await,
             Storage::Local(local) => local.confirm_same(key, e_tag).await,
             Storage::AdmissionControlledS3(as3) => as3.confirm_same(key, e_tag).await,
         }
@@ -363,6 +380,7 @@ impl Storage {
     ) -> Result<Option<ETag>, StorageError> {
         match self {
             Storage::S3(s3) => s3.put_file(key, path, options).await,
+            Storage::Object(obj) => obj.put_file(key, path, options).await.map(Some),
             Storage::Local(local) => local.put_file(key, path, options).await,
             Storage::AdmissionControlledS3(as3) => as3.put_file(key, path, options).await,
         }
@@ -376,6 +394,7 @@ impl Storage {
     ) -> Result<Option<ETag>, StorageError> {
         match self {
             Storage::S3(s3) => s3.put_bytes(key, bytes, options).await,
+            Storage::Object(obj) => obj.put(key, bytes.into(), options).await.map(Some),
             Storage::Local(local) => local.put_bytes(key, &bytes, options).await,
             Storage::AdmissionControlledS3(as3) => as3.put_bytes(key, bytes, options).await,
         }
@@ -384,12 +403,20 @@ impl Storage {
     pub async fn delete(&self, key: &str, options: DeleteOptions) -> Result<(), StorageError> {
         match self {
             Storage::S3(s3) => s3.delete(key, options).await,
-            Storage::Local(local) => {
+            Storage::Object(obj) => {
                 if options.if_match.is_some() {
                     return Err(StorageError::Message {
                         message: "if match not supported for object store backend".to_string(),
                     });
                 }
+                obj.delete(key).await
+            }
+            Storage::Local(local) => {
+                if options.if_match.is_some() {
+                    return Err(StorageError::Message {
+                        message: "if match not supported for local backend".to_string(),
+                    });
+                }
                 local.delete(key).await
             }
             Storage::AdmissionControlledS3(ac) => ac.delete(key, options).await,
@@ -402,6 +429,7 @@ impl Storage {
     ) -> Result<crate::s3::DeletedObjects, StorageError> {
         match self {
             Storage::S3(s3) => s3.delete_many(keys).await,
+            Storage::Object(obj) => obj.delete_many(keys).await,
             Storage::Local(local) => local.delete_many(keys).await,
             Storage::AdmissionControlledS3(ac) => ac.delete_many(keys).await,
         }
@@ -410,6 +438,7 @@ impl Storage {
     pub async fn rename(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> {
         match self {
             Storage::S3(s3) => s3.rename(src_key, dst_key).await,
+            Storage::Object(obj) => obj.rename(src_key, dst_key).await,
             Storage::Local(local) => local.rename(src_key, dst_key).await,
             Storage::AdmissionControlledS3(_) => Err(StorageError::NotImplemented),
         }
@@ -418,6 +447,7 @@ impl Storage {
     pub async fn copy(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> {
         match self {
             Storage::S3(s3) => s3.copy(src_key, dst_key).await,
+            Storage::Object(obj) => obj.copy(src_key, dst_key).await,
             Storage::Local(local) => local.copy(src_key, dst_key).await,
             Storage::AdmissionControlledS3(ac) => ac.copy(src_key, dst_key).await,
         }
@@ -431,6 +461,7 @@ impl Storage {
         match self {
             Storage::Local(local) => local.list_prefix(prefix).await,
             Storage::S3(s3) => s3.list_prefix(prefix).await,
+            Storage::Object(obj) => obj.list_prefix(prefix).await,
             Storage::AdmissionControlledS3(acs3) => acs3.list_prefix(prefix, options).await,
         }
     }
@@ -446,6 +477,9 @@ impl Configurable<StorageConfig> for Storage {
             StorageConfig::S3(_) => Ok(Storage::S3(
                 s3::S3Storage::try_from_config(config, registry).await?,
             )),
+            StorageConfig::Object(_) => Ok(Storage::Object(
+                object_storage::ObjectStorage::try_from_config(config, registry).await?,
+            )),
             StorageConfig::Local(_) => Ok(Storage::Local(
                 local::LocalStorage::try_from_config(config, registry).await?,
             )),