From bef8e31d2ec1fbef866324ea106429fda3b1fbf6 Mon Sep 17 00:00:00 2001
From: Hengfei Yang <hengfei.yang@gmail.com>
Date: Mon, 3 Apr 2023 21:59:13 +0800
Subject: [PATCH] feat: add prefix for s3 bucket (#492)

---
 src/infra/config.rs     |  5 +++++
 src/infra/storage/s3.rs | 43 +++++++++++++++++++++++++++++++++++------
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/infra/config.rs b/src/infra/config.rs
index 7ce94155aa..e68dd10ca1 100644
--- a/src/infra/config.rs
+++ b/src/infra/config.rs
@@ -291,6 +291,8 @@ pub struct S3 {
     pub secret_key: String,
     #[env_config(name = "ZO_S3_BUCKET_NAME", default = "")]
     pub bucket_name: String,
+    #[env_config(name = "ZO_S3_BUCKET_PREFIX", default = "")]
+    pub bucket_prefix: String,
 }
 
 pub fn init() -> Config {
@@ -453,6 +455,9 @@ fn check_s3_config(cfg: &mut Config) -> Result<(), anyhow::Error> {
             cfg.s3.provider = "aws".to_string();
         }
     }
+    if !cfg.s3.bucket_prefix.ends_with('/') {
+        cfg.s3.bucket_prefix = format!("{}/", cfg.s3.bucket_prefix);
+    }
     cfg.s3.provider = cfg.s3.provider.to_lowercase();
     Ok(())
 }
diff --git a/src/infra/storage/s3.rs b/src/infra/storage/s3.rs
index c37cdf000f..897488fe60 100644
--- a/src/infra/storage/s3.rs
+++ b/src/infra/storage/s3.rs
@@ -36,6 +36,14 @@ pub struct S3 {}
 #[async_trait]
 impl FileStorage for S3 {
     async fn list(&self, prefix: &str) -> Result<Vec<String>, anyhow::Error> {
+        let prefix = if !CONFIG.s3.bucket_prefix.is_empty()
+            && !prefix.starts_with(&CONFIG.s3.bucket_prefix)
+        {
+            format!("{}{}", CONFIG.s3.bucket_prefix, prefix)
+        } else {
+            prefix.to_string()
+        };
+        let prefix = prefix.as_str();
         let mut files = Vec::new();
         let client = S3CLIENT.get().await.clone().unwrap();
         let mut start_after: String = "".to_string();
@@ -90,11 +98,17 @@ impl FileStorage for S3 {
     }
 
     async fn get(&self, file: &str) -> Result<bytes::Bytes, anyhow::Error> {
+        let key =
+            if !CONFIG.s3.bucket_prefix.is_empty() && !file.starts_with(&CONFIG.s3.bucket_prefix) {
+                format!("{}{}", CONFIG.s3.bucket_prefix, file)
+            } else {
+                file.to_string()
+            };
         let client = S3CLIENT.get().await.clone().unwrap();
         let object = match client
             .get_object()
             .bucket(&CONFIG.s3.bucket_name)
-            .key(file)
+            .key(key)
             .send()
             .await
         {
@@ -118,12 +132,18 @@ impl FileStorage for S3 {
     }
 
     async fn put(&self, file: &str, data: bytes::Bytes) -> Result<(), anyhow::Error> {
+        let key =
+            if !CONFIG.s3.bucket_prefix.is_empty() && !file.starts_with(&CONFIG.s3.bucket_prefix) {
+                format!("{}{}", CONFIG.s3.bucket_prefix, file)
+            } else {
+                file.to_string()
+            };
         let client = S3CLIENT.get().await.clone().unwrap();
         let data_size = data.len();
         match client
             .put_object()
             .bucket(&CONFIG.s3.bucket_name)
-            .key(file)
+            .key(key)
             .body(data.into())
             .send()
             .await
@@ -177,9 +197,14 @@ impl FileStorage for S3 {
                             files[start..(start + step).min(files_len)]
                                 .iter()
                                 .map(|file| {
-                                    ObjectIdentifier::builder()
-                                        .set_key(Some(file.to_string()))
-                                        .build()
+                                    let key = if !CONFIG.s3.bucket_prefix.is_empty()
+                                        && !file.starts_with(&CONFIG.s3.bucket_prefix)
+                                    {
+                                        format!("{}{}", CONFIG.s3.bucket_prefix, file)
+                                    } else {
+                                        file.to_string()
+                                    };
+                                    ObjectIdentifier::builder().set_key(Some(key)).build()
                                 })
                                 .collect::<Vec<_>>(),
                         ))
@@ -205,11 +230,17 @@ impl FileStorage for S3 {
 
 impl S3 {
     async fn del_for_gcs(&self, file: &str) -> Result<(), anyhow::Error> {
+        let key =
+            if !CONFIG.s3.bucket_prefix.is_empty() && !file.starts_with(&CONFIG.s3.bucket_prefix) {
+                format!("{}{}", CONFIG.s3.bucket_prefix, file)
+            } else {
+                file.to_string()
+            };
         let client = S3CLIENT.get().await.clone().unwrap();
         let result = client
             .delete_object()
             .bucket(&CONFIG.s3.bucket_name)
-            .key(file)
+            .key(key)
             .send()
             .await;
         match result {