feat(websocket): support GCS client (#290)

* feat(websocket): Support snapshot for GCS

* fix(websocket): Remove unnecessary mut

* feat(websocket): Make GcsClient enable to take a arbitrary value

* feat(websocket): support download

* feat(websocket): define GCS client

* fix(websocket): resolve conflicts

* organize imports of domain files

---------

Co-authored-by: pyshx <[email protected]>

Author: kirisaki

websocket/Cargo.lock

@@ -1,6 +1,6 @@
 [workspace]
 members = [
-    "crates/*",
+  "crates/*",
 ]
 
 resolver = "2"
@@ -38,16 +38,17 @@ strip = true
 flow-websocket-domain = {path = "crates/domain"}
 
 async-trait = "0.1.80"
-serde = { version = "1.0", features = ["derive"] }
-serde_json = {version = "1.0.117", features = ["arbitrary_precision"]}
-chrono = { version = "0.4", features = ["serde"] }
-redis = { version = "0.25.4", features = ["aio", "tokio-comp"] }
+chrono = {version = "0.4", features = ["serde"]}
+google-cloud-storage = "0.18"
+redis = {version = "0.25.4", features = ["aio", "tokio-comp"]}
 rslock = "0.3.0"
+serde = {version = "1.0", features = ["derive"]}
+serde_json = {version = "1.0.117", features = ["arbitrary_precision"]}
 tokio = {version = "1.38.0", features = ["full", "time"]}
 uuid = {version = "1.8.0", features = [
   "v4",
   "fast-rng",
   "macro-diagnostics",
   "serde",
 ]}
-yrs = "0.18"
+yrs = "0.18"

websocket/Cargo.toml

@@ -1,10 +1,10 @@
-use serde::{Deserialize, Serialize};
-use crate::snapshot::{ProjectSnapshot, ObjectTenant, ObjectDelete};
-use crate::repository::ProjectSnapshotRepository;
 use std::error::Error;
 
-use crate::utils::generate_id;
+use serde::{Deserialize, Serialize};
 
+use crate::repository::ProjectSnapshotRepository;
+use crate::snapshot::{ObjectDelete, ObjectTenant, ProjectSnapshot};
+use crate::utils::generate_id;
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct Project {

websocket/crates/domain/src/project.rs

@@ -1,6 +1,7 @@
+use std::error::Error;
+
 use crate::project::{Project, ProjectEditingSession};
 use crate::snapshot::ProjectSnapshot;
-use std::error::Error;
 
 #[async_trait::async_trait]
 pub trait ProjectRepository {

websocket/crates/domain/src/repository.rs

@@ -14,6 +14,7 @@ version.workspace = true
 flow-websocket-domain.workspace = true
 
 async-trait.workspace = true
+google-cloud-storage.workspace = true
 redis.workspace = true
 rslock.workspace = true
 serde.workspace = true

websocket/crates/infra/Cargo.toml

@@ -0,0 +1,60 @@
+use google_cloud_storage::client::{Client, ClientConfig};
+use google_cloud_storage::http::objects::download::Range;
+use google_cloud_storage::http::objects::get::GetObjectRequest;
+use google_cloud_storage::http::objects::upload::{Media, UploadObjectRequest, UploadType};
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone)]
+pub struct GcsClient {
+    client: Client,
+    bucket: String,
+}
+
+impl GcsClient {
+    pub async fn new(bucket: String) -> Result<Self, Box<dyn std::error::Error>> {
+        let config = ClientConfig::default().with_auth().await?;
+        let client = Client::new(config);
+        Ok(GcsClient { client, bucket })
+    }
+
+    pub async fn upload<T: Serialize>(
+        &self,
+        path: String,
+        data: &T,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let upload_type = UploadType::Simple(Media::new(path));
+        let bytes = serde_json::to_string(data)?;
+        let _uploaded = self
+            .client
+            .upload_object(
+                &UploadObjectRequest {
+                    bucket: self.bucket.clone(),
+                    ..Default::default()
+                },
+                bytes,
+                &upload_type,
+            )
+            .await?;
+        Ok(())
+    }
+
+    pub async fn download<T: for<'de> Deserialize<'de>>(
+        &self,
+        path: String,
+    ) -> Result<T, Box<dyn std::error::Error>> {
+        let bytes = self
+            .client
+            .download_object(
+                &GetObjectRequest {
+                    bucket: self.bucket.clone(),
+                    object: path,
+                    ..Default::default()
+                },
+                &Range::default(),
+            )
+            .await?;
+        let src = String::from_utf8(bytes)?;
+        let data = serde_json::from_str(&src)?;
+        Ok(data)
+    }
+}

websocket/crates/infra/src/persistence/gcs/gcs_client.rs

@@ -0,0 +1 @@
+pub mod gcs_client;

websocket/crates/infra/src/persistence/gcs/mod.rs

@@ -1,2 +1,3 @@
+pub mod gcs;
+pub mod project_repository;
 pub mod redis;
-pub mod project_repository;

websocket/crates/infra/src/persistence/mod.rs

@@ -3,11 +3,15 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 
-use flow_websocket_domain::project::{Project, ProjectEditingSession};
-use flow_websocket_domain::repository::{ProjectRepository, ProjectEditingSessionRepository};
-
+use crate::persistence::gcs::gcs_client::GcsClient;
 use crate::persistence::redis::redis_client::RedisClient;
 
+use flow_websocket_domain::project::{Project, ProjectEditingSession};
+use flow_websocket_domain::repository::{
+    ProjectEditingSessionRepository, ProjectRepository, ProjectSnapshotRepository,
+};
+use flow_websocket_domain::snapshot::ProjectSnapshot;
+
 pub struct ProjectRedisRepository {
     redis_client: Arc<RedisClient>,
 }
@@ -34,7 +38,10 @@ impl ProjectEditingSessionRepository for ProjectRedisRepository {
         Ok(())
     }
 
-    async fn get_active_session(&self, project_id: &str) -> Result<Option<ProjectEditingSession>, Box<dyn Error>> {
+    async fn get_active_session(
+        &self,
+        project_id: &str,
+    ) -> Result<Option<ProjectEditingSession>, Box<dyn Error>> {
         let key = format!("project:{}:active_session", project_id);
         self.redis_client.get(&key).await
     }
@@ -45,3 +52,40 @@ impl ProjectEditingSessionRepository for ProjectRedisRepository {
         Ok(())
     }
 }
+
+/// A `ProjectGcsRepository` is a thin wrapper of `GcsClient`.
+pub struct ProjectGcsRepository {
+    client: GcsClient,
+}
+
+impl ProjectGcsRepository {
+    /// Returns the `ProjectGcsRepository`.
+    fn new(client: GcsClient) -> Self {
+        Self { client }
+    }
+}
+
+#[async_trait]
+impl ProjectSnapshotRepository for ProjectGcsRepository {
+    /// Create a snapshot.
+    async fn create_snapshot(&self, snapshot: ProjectSnapshot) -> Result<(), Box<dyn Error>> {
+        let path = format!("snapshot/{}", snapshot.id);
+        self.client.upload(path, &snapshot).await?;
+        Ok(())
+    }
+
+    /// Get the latest snapshot.
+    async fn get_latest_snapshot(
+        &self,
+        project_id: &str,
+    ) -> Result<Option<ProjectSnapshot>, Box<dyn Error>> {
+        let path = format!("snapshot/{}:latest_snapshot", project_id);
+        self.client.download(path).await
+    }
+
+    /// Get the state of the latest snapshot.
+    async fn get_latest_snapshot_state(&self, project_id: &str) -> Result<Vec<u8>, Box<dyn Error>> {
+        let path = format!("snapshot/{}:latest_snapshot_state", project_id);
+        self.client.download(path).await
+    }
+}