From 71ed33b7312dbbb867386916043b7252441a0c1f Mon Sep 17 00:00:00 2001
From: Robert Escriva <robert@trychroma.com>
Date: Mon, 16 Dec 2024 14:46:41 -0800
Subject: [PATCH] [ENH]  Support delayed workloads in chroma-load-start.

---
 rust/load/src/bin/chroma-load-start.rs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/rust/load/src/bin/chroma-load-start.rs b/rust/load/src/bin/chroma-load-start.rs
index d2d715611b7..2ddfac9f882 100644
--- a/rust/load/src/bin/chroma-load-start.rs
+++ b/rust/load/src/bin/chroma-load-start.rs
@@ -14,6 +14,8 @@ struct Args {
     #[arg(long)]
     expires: String,
     #[arg(long)]
+    delay: Option<String>,
+    #[arg(long)]
     data_set: String,
     #[arg(long)]
     workload: String,
@@ -84,11 +86,20 @@ async fn main() {
     let args = Args::parse();
     let client = reqwest::Client::new();
     let throughput = args.throughput();
+    let mut workload = Workload::ByName(args.workload);
+    if let Some(delay) = args.delay {
+        let delay = humanize_expires(&delay).expect("delay must be humanizable");
+        let delay = delay.parse().expect("delay must be a date time");
+        workload = Workload::Delay {
+            after: delay,
+            wrap: Box::new(workload),
+        };
+    }
     let req = StartRequest {
         name: args.name,
         expires: humanize_expires(&args.expires).unwrap_or(args.expires),
         data_set: args.data_set,
-        workload: Workload::ByName(args.workload),
+        workload,
         throughput,
     };
     match client