feat: add dynamic bucket creation support

src/args.rs

@@ -14,6 +14,15 @@ pub struct Args {
     #[arg(long, env = "AWS_BUCKET", value_name = "NAME")]
     pub bucket_name: String,
 
+    #[arg(
+        long,
+        env = "SOS_BUCKET_CREATE", 
+        action = ArgAction::SetTrue,
+        default_value_t = Args::default_bucket_create(),
+    )]
+    #[serde(default = "Args::default_bucket_create")]
+    pub bucket_create: bool,
+
     #[command(flatten)]
     #[serde(default, flatten)]
     pub credentials: CredentialsArgs,
@@ -35,20 +44,28 @@ impl Args {
     pub fn print(&self) {
         let Self {
             bucket_name,
+            bucket_create,
             credentials,
             load_tester,
             load_tester_job,
             region,
         } = self;
 
         info!("bucket_name: {bucket_name}");
+        info!("bucket_create: {bucket_create}");
         credentials.print();
         load_tester.print();
         load_tester_job.print();
         region.print();
     }
 }
 
+impl Args {
+    const fn default_bucket_create() -> bool {
+        false
+    }
+}
+
 #[derive(Clone, Default, PartialEq, Parser, Serialize, Deserialize)]
 #[clap(rename_all = "kebab-case")]
 #[serde(rename_all = "camelCase")]

src/session.rs

@@ -15,7 +15,7 @@ use clap::Parser;
 use futures::{stream::FuturesUnordered, FutureExt, TryFutureExt, TryStreamExt};
 use indicatif::{ProgressBar, ProgressState, ProgressStyle};
 use rand::{rngs::SmallRng, RngCore, SeedableRng};
-use s3::{serde_types::InitiateMultipartUploadResponse, Bucket};
+use s3::{serde_types::InitiateMultipartUploadResponse, Bucket, BucketConfiguration};
 use tokio::{spawn, task::JoinHandle, time::sleep};
 use tracing::{error, info};
 
@@ -35,22 +35,40 @@ impl ObjectStorageSession {
 
         let Args {
             bucket_name,
+            bucket_create,
             credentials,
             load_tester,
             load_tester_job,
             region,
         } = args;
 
-        let bucket = Bucket::new(&bucket_name, region.into(), credentials.into())
-            .map_err(|error| anyhow!("failed to initialize object storage bucket client: {error}"))?
-            .with_path_style();
-
-        if !bucket
-            .exists()
-            .await
-            .map_err(|error| anyhow!("failed to check object storage bucket: {error}"))?
-        {
-            bail!("no such bucket: {bucket_name}")
+        let mut bucket = Bucket::new(
+            &bucket_name,
+            region.clone().into(),
+            credentials.clone().into(),
+        )
+        .map_err(|error| anyhow!("failed to initialize object storage bucket client: {error}"))?
+        .with_path_style();
+
+        if bucket.get_object("/").await.is_err() {
+            if bucket_create {
+                let config = BucketConfiguration::private();
+                let response = Bucket::create_with_path_style(
+                    &bucket_name,
+                    region.into(),
+                    credentials.into(),
+                    config,
+                )
+                .await
+                .map_err(|error| anyhow!("failed to create object storage bucket: {error}"))?;
+                if response.success() {
+                    bucket = response.bucket.with_path_style();
+                } else {
+                    bail!("failed to create bucket: {bucket_name}")
+                }
+            } else {
+                bail!("no such bucket: {bucket_name}")
+            }
         }
 
         Ok(Self {