Skip to content

Commit 7e50089

Browse files
authored
[CLN] Pass Chroma API key via env var for chroma-load-start (#4469)
## Description of changes This change updates chroma-load-start to read the APi key from the environment. It is a bit more secure than passing the API key via the process arguments, which is potentially readable by other processes on the system. ## Test plan Tested locally via tilt.
1 parent ec1f94c commit 7e50089

File tree

2 files changed

+14
-4
lines changed

2 files changed

+14
-4
lines changed

rust/load/src/bin/chroma-load-start.rs

+5-3
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,6 @@ struct Args {
2424
#[arg(long)]
2525
database: String,
2626
#[arg(long)]
27-
api_key: String,
28-
#[arg(long)]
2927
constant_throughput: Option<f64>,
3028
#[arg(long)]
3129
sinusoid_throughput: Option<String>,
@@ -90,6 +88,10 @@ impl Args {
9088
#[tokio::main]
9189
async fn main() {
9290
let args = Args::parse();
91+
92+
// Read API key from environment variable.
93+
let api_key = std::env::var("CHROMA_API_KEY").expect("CHROMA_API_KEY is not set");
94+
9395
let client = reqwest::Client::new();
9496
let throughput = args.throughput();
9597
let mut workload = Workload::ByName(args.workload);
@@ -108,7 +110,7 @@ async fn main() {
108110
data_set: args.data_set,
109111
connection: Connection {
110112
url: args.url,
111-
api_key: args.api_key,
113+
api_key,
112114
database: args.database,
113115
},
114116
throughput,

rust/load/src/lib.rs

+9-1
Original file line numberDiff line numberDiff line change
@@ -1026,8 +1026,16 @@ pub struct LoadHarness {
10261026

10271027
impl LoadHarness {
10281028
/// The status of the load harness.
1029+
/// This returns the list of running workloads with secrets redacted.
10291030
pub fn status(&self) -> Vec<RunningWorkload> {
1030-
self.running.clone()
1031+
self.running
1032+
.iter()
1033+
.map(|w| {
1034+
let mut w = w.clone();
1035+
w.connection.api_key = "REDACTED".to_string();
1036+
w
1037+
})
1038+
.collect()
10311039
}
10321040

10331041
/// Start a workload on the load harness.

0 commit comments

Comments
 (0)