Skip to content

Commit

Permalink
[ENH] Allow config of storage from env (#2093)
Browse files Browse the repository at this point in the history
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- Allows configuring storage credentials from ENV by allowing case
insensitivity.
 - New functionality
	 - None

## Test plan
*How are these changes tested?*
A test was modified to use env and assert the config is as expected
- [x/ Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
None
  • Loading branch information
HammadB authored May 1, 2024
1 parent eb18e5c commit e6c35a0
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 5 deletions.
16 changes: 12 additions & 4 deletions rust/worker/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,8 @@ mod tests {
let _ = jail.set_env("CHROMA_QUERY_SERVICE__MY_PORT", 50051);
let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__MY_IP", "192.0.0.1");
let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__MY_PORT", 50051);
let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__BUCKET", "buckets!");
let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__CREDENTIALS", "AWS");
let _ = jail.create_file(
"chroma_config.yaml",
r#"
Expand Down Expand Up @@ -437,10 +439,6 @@ mod tests {
Grpc:
host: "localhost"
port: 50051
storage:
S3:
bucket: "chroma"
credentials: Minio
log:
Grpc:
host: "localhost"
Expand All @@ -458,6 +456,16 @@ mod tests {
let config = RootConfig::load();
assert_eq!(config.query_service.my_ip, "192.0.0.1");
assert_eq!(config.query_service.my_port, 50051);
match &config.compaction_service.storage {
crate::storage::config::StorageConfig::S3(s) => {
assert_eq!(s.bucket, "buckets!");
assert_eq!(
s.credentials,
crate::storage::config::S3CredentialsConfig::AWS
);
}
_ => panic!("Invalid storage config"),
}
Ok(())
});
}
Expand Down
5 changes: 4 additions & 1 deletion rust/worker/src/storage/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,14 @@ use std::path::Path;
/// See config.rs in the root of the worker crate for an example of how to use
/// config files to configure the worker.
pub(crate) enum StorageConfig {
// case-insensitive
#[serde(alias = "s3")]
S3(S3StorageConfig),
#[serde(alias = "local")]
Local(LocalStorageConfig),
}

#[derive(Deserialize)]
#[derive(Deserialize, PartialEq, Debug)]
pub(crate) enum S3CredentialsConfig {
Minio,
AWS,
Expand Down
Binary file added rust/worker/test.arrow
Binary file not shown.

0 comments on commit e6c35a0

Please sign in to comment.