feature: impl BucketStorageBuilder

thorseraq · thorseraq · commit ff8bf7447765 · 2023-06-16T14:56:44.000+08:00
diff --git a/.env-template b/.env-template
@@ -20,3 +20,10 @@ SIGN_KEY =
 #AFFINE_CLOUD_LOG=debug,mio=off,hyper=off,rustls=off,tantivy=off,sqlx::query=off,jwst_rpc=trace,jwst_rpc::context=off,affine_cloud=trace
 #AFFINE_CLOUD_LOG=tower_http=debug,jwst_rpc=trace,jwst_rpc::context=off,affine_cloud=trace
 #KECK_LOG=debug,mio=off,hyper=off,rustls=off,tantivy=off,sqlx::query=off,jwst_rpc=trace,jwst_rpc::context=off,keck=trace
+
+# For Bucket storage, BUCKET_ACCESS_TOKEN, BUCKET_SECRET_TOKEN, BUCKET_ENDPOINT is necessary
+#BUCKET_ACCESS_TOKEN=
+#BUCKET_SECRET_TOKEN=
+#BUCKET_ENDPOINT=
+#BUCKET_NAME=
+#BUCKET_ROOT=
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/libs/jwst-storage/Cargo.toml b/libs/jwst-storage/Cargo.toml
@@ -33,6 +33,7 @@ opendal = { version = "0.37.0", default-features = false, features = [
     "rustls",
     "services-s3",
 ] }
+dotenvy = "0.15.7"
 
 # ======= workspace dependencies =======
 jwst = { workspace = true }
diff --git a/libs/jwst-storage/src/storage/blobs/bucket_local_db.rs b/libs/jwst-storage/src/storage/blobs/bucket_local_db.rs
@@ -5,20 +5,19 @@ use crate::JwstStorageError;
 use bytes::Bytes;
 use futures::Stream;
 use jwst::{BlobMetadata, BlobStorage, BucketBlobStorage, JwstResult};
+use jwst_storage_migration::Migrator;
+use opendal::services::S3;
 use opendal::Operator;
 use sea_orm::{DatabaseConnection, EntityTrait};
+use sea_orm_migration::MigratorTrait;
 use std::collections::HashMap;
 use std::sync::Arc;
 
-#[allow(unused)]
 pub(super) type BucketBlobModel = <BucketBlobs as EntityTrait>::Model;
-#[allow(unused)]
 type BucketBlobActiveModel = entities::bucket_blobs::ActiveModel;
-#[allow(unused)]
 type BucketBlobColumn = <BucketBlobs as EntityTrait>::Column;
 
 #[derive(Clone)]
-#[allow(unused)]
 pub struct BlobBucketDBStorage {
     bucket: Arc<Bucket>,
     pub(super) pool: DatabaseConnection,
@@ -31,17 +30,53 @@ impl AsRef<DatabaseConnection> for BlobBucketDBStorage {
     }
 }
 
-#[allow(unused)]
 impl BlobBucketDBStorage {
+    #[allow(unused)]
     pub async fn init_with_pool(
         pool: DatabaseConnection,
         bucket: Arc<Bucket>,
+        bucket_storage: Option<BucketStorage>,
     ) -> JwstStorageResult<Self> {
-        todo!()
+        Migrator::up(&pool, None).await?;
+        Ok(Self {
+            bucket,
+            pool,
+            bucket_storage: bucket_storage.unwrap_or(BucketStorage::new()?),
+        })
     }
 
-    pub async fn init_pool(database: &str) -> JwstStorageResult<Self> {
-        todo!()
+    #[allow(unused)]
+    pub async fn init_pool(
+        database: &str,
+        bucket_storage: Option<BucketStorage>,
+    ) -> JwstStorageResult<Self> {
+        let is_sqlite = is_sqlite(database);
+        let pool = create_connection(database, is_sqlite).await?;
+
+        Self::init_with_pool(pool, get_bucket(is_sqlite), bucket_storage).await
+    }
+
+    #[allow(unused)]
+    async fn all(&self, workspace: &str) -> Result<Vec<BucketBlobModel>, DbErr> {
+        BucketBlobs::find()
+            .filter(BucketBlobColumn::Workspace.eq(workspace))
+            .all(&self.pool)
+            .await
+    }
+
+    #[allow(unused)]
+    async fn count(&self, workspace: &str) -> Result<u64, DbErr> {
+        BucketBlobs::find()
+            .filter(BucketBlobColumn::Workspace.eq(workspace))
+            .count(&self.pool)
+            .await
+    }
+
+    async fn exists(&self, workspace: &str, hash: &str) -> Result<bool, DbErr> {
+        BucketBlobs::find_by_id((workspace.into(), hash.into()))
+            .count(&self.pool)
+            .await
+            .map(|c| c > 0)
     }
 
     pub(super) async fn metadata(
@@ -60,13 +95,6 @@ impl BlobBucketDBStorage {
             .and_then(|r| r.ok_or(JwstBlobError::BlobNotFound(hash.into())))
     }
 
-    async fn exists(&self, workspace: &str, hash: &str) -> Result<bool, DbErr> {
-        BucketBlobs::find_by_id((workspace.into(), hash.into()))
-            .count(&self.pool)
-            .await
-            .map(|c| c > 0)
-    }
-
     pub(super) async fn get_blobs_size(&self, workspace: &str) -> Result<Option<i64>, DbErr> {
         BucketBlobs::find()
             .filter(BucketBlobColumn::Workspace.eq(workspace))
@@ -115,7 +143,29 @@ pub struct BucketStorage {
     pub(super) op: Operator,
 }
 
-// TODO Builder for BucketStorage;
+impl BucketStorage {
+    #[allow(unused)]
+    pub fn new() -> JwstStorageResult<Self> {
+        let access_key = dotenvy::var("BUCKET_ACCESS_TOKEN")?;
+        let secret_access_key = dotenvy::var("BUCKET_SECRET_TOKEN")?;
+        let endpoint = dotenvy::var("BUCKET_ENDPOINT")?;
+        let bucket = dotenvy::var("BUCKET_NAME");
+        let root = dotenvy::var("BUCKET_ROOT");
+
+        let mut builder = S3::default();
+
+        builder.bucket(bucket.unwrap_or("default_bucket".to_string()).as_str());
+        builder.root(root.unwrap_or("default_root".to_string()).as_str());
+        builder.endpoint(endpoint.as_str());
+        builder.access_key_id(access_key.as_str());
+        builder.secret_access_key(secret_access_key.as_str());
+
+        Ok(Self {
+            op: Operator::new(builder)?.finish(),
+        })
+    }
+}
+
 // TODO add retry layer
 
 #[allow(unused_variables)]
@@ -250,3 +300,23 @@ impl BlobStorage<JwstStorageError> for BlobBucketDBStorage {
         return Ok(size.unwrap_or(0));
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::storage::blobs::utils::BucketStorageBuilder;
+
+    #[tokio::test]
+    #[ignore = "need to config bucket auth"]
+    async fn test_init_bucket_storage() {
+        let bucket_storage = BucketStorageBuilder::new()
+            .bucket("bucket")
+            .root("root")
+            .build()
+            .unwrap();
+
+        BlobBucketDBStorage::init_pool("sqlite::memory:", Some(bucket_storage))
+            .await
+            .unwrap();
+    }
+}
diff --git a/libs/jwst-storage/src/storage/blobs/mod.rs b/libs/jwst-storage/src/storage/blobs/mod.rs
@@ -14,6 +14,9 @@ use thiserror::Error;
 use tokio::task::JoinError;
 use utils::{ImageParams, InternalBlobMetadata};
 
+pub use utils::BucketStorageBuilder;
+pub use bucket_local_db::BlobBucketDBStorage;
+
 #[derive(Debug, Error)]
 pub enum JwstBlobError {
     #[error("blob not found: {0}")]
diff --git a/libs/jwst-storage/src/storage/blobs/utils.rs b/libs/jwst-storage/src/storage/blobs/utils.rs
@@ -1,3 +1,5 @@
+use crate::storage::blobs::bucket_local_db::BucketStorage;
+use crate::{JwstStorageError, JwstStorageResult};
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use futures::{
@@ -6,6 +8,8 @@ use futures::{
 };
 use image::{load_from_memory, ImageOutputFormat, ImageResult};
 use jwst::{Base64Engine, BlobMetadata, URL_SAFE_ENGINE};
+use opendal::services::S3;
+use opendal::Operator;
 use sea_orm::FromQueryResult;
 use sha2::{Digest, Sha256};
 use std::{collections::HashMap, io::Cursor};
@@ -139,3 +143,101 @@ impl From<InternalBlobMetadata> for BlobMetadata {
         }
     }
 }
+
+impl TryFrom<HashMap<String, String>> for BucketStorage {
+    type Error = JwstStorageError;
+
+    fn try_from(map: HashMap<String, String>) -> Result<Self, Self::Error> {
+        let mut builder = BucketStorageBuilder::new();
+        let access_token = map.get("BUCKET_ACCESS_TOKEN");
+        let secret_access_key = map.get("BUCKET_SECRET_TOKEN");
+        let endpoint = map.get("BUCKET_ENDPOINT");
+        let bucket = map.get("BUCKET_NAME");
+        let root = map.get("BUCKET_ROOT");
+
+        if let Some(access_token) = access_token {
+            builder = builder.access_key(access_token);
+        }
+        if let Some(secret_access_key) = secret_access_key {
+            builder = builder.secret_access_key(secret_access_key);
+        }
+        if let Some(endpoint) = endpoint {
+            builder = builder.endpoint(endpoint);
+        }
+        if let Some(bucket) = bucket {
+            builder = builder.bucket(bucket);
+        }
+        if let Some(root) = root {
+            builder = builder.root(root);
+        }
+
+        builder.build()
+    }
+}
+
+pub struct BucketStorageBuilder {
+    access_key: String,
+    secret_access_key: String,
+    endpoint: String,
+    bucket: String,
+    root: String,
+}
+
+impl BucketStorageBuilder {
+    pub fn new() -> Self {
+        let access_key =
+            dotenvy::var("BUCKET_ACCESS_TOKEN").unwrap_or("default_access_key".to_string());
+        let secret_access_key =
+            dotenvy::var("BUCKET_SECRET_TOKEN").unwrap_or("default_secret_key".to_string());
+        let endpoint = dotenvy::var("BUCKET_ENDPOINT").unwrap_or("default_endpoint".to_string());
+        let bucket = dotenvy::var("BUCKET_NAME").unwrap_or("__default_bucket__".to_string());
+        let root = dotenvy::var("BUCKET_ROOT").unwrap_or("__default_root__".to_string());
+
+        Self {
+            access_key,
+            secret_access_key,
+            endpoint,
+            bucket,
+            root,
+        }
+    }
+
+    pub fn access_key(mut self, access_key: &str) -> Self {
+        self.access_key = access_key.to_string();
+        self
+    }
+
+    pub fn secret_access_key(mut self, secret_access_key: &str) -> Self {
+        self.secret_access_key = secret_access_key.to_string();
+        self
+    }
+
+    pub fn endpoint(mut self, endpoint: &str) -> Self {
+        self.endpoint = endpoint.to_string();
+        self
+    }
+
+    pub fn bucket(mut self, bucket: &str) -> Self {
+        self.bucket = bucket.to_string();
+        self
+    }
+
+    pub fn root(mut self, root: &str) -> Self {
+        self.root = root.to_string();
+        self
+    }
+
+    pub fn build(self) -> JwstStorageResult<BucketStorage> {
+        let mut builder = S3::default();
+
+        builder.bucket(self.bucket.as_str());
+        builder.root(self.root.as_str());
+        builder.endpoint(self.endpoint.as_str());
+        builder.access_key_id(self.access_key.as_str());
+        builder.secret_access_key(self.secret_access_key.as_str());
+
+        Ok(BucketStorage {
+            op: Operator::new(builder)?.finish(),
+        })
+    }
+}
diff --git a/libs/jwst-storage/src/types.rs b/libs/jwst-storage/src/types.rs
@@ -19,8 +19,9 @@ pub enum JwstStorageError {
     #[error("failed to process blob")]
     JwstBlob(#[from] crate::storage::blobs::JwstBlobError),
     #[error("bucket error")]
-    JwstBucketError(#[from] opendal::Error),
-
+    JwstBucket(#[from] opendal::Error),
+    #[error("env variables read error")]
+    DotEnvy(#[from] dotenvy::Error),
 }
 
 pub type JwstStorageResult<T> = Result<T, JwstStorageError>;
