Merge pull request #61 from CoLearn-Dev/storage_macro_rdbc

RDBC support for storage macro extension

Author: nociza

.github/workflows/check.yml

@@ -30,6 +30,11 @@ jobs:
         run: |
           brew install docker
           colima start
+      - name: Start container (MySQL)
+        run: |
+          docker run --name mysql -e MYSQL_ROOT_PASSWORD=password -e MYSQL_DATABASE=test_db -p 3306:3306 -d mysql:8.0
+          MYSQL_HOST="127.0.0.1"
+          echo "MYSQL_DATABASE_URL=mysql://root:password@${MYSQL_HOST}:3306/test_db" >> $GITHUB_ENV
       - name: Start container (mq)
         if: ${{ matrix.mq != 'standalone' }}
         run: docker run -d -p 5672:5672 -p 15672:15672 -p 16379:6379 ${{ matrix.docker_image }}

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "colink"
-version = "0.3.7"
+version = "0.3.8"
 edition = "2021"
 description = "CoLink Rust SDK"
 license = "MIT"
@@ -23,6 +23,7 @@ lapin = "2.1"
 prost = "0.10"
 rand = { version = "0.8", features = ["std_rng"] }
 rcgen = { version = "0.10", optional = true }
+rdbc2 = { version = "0.2.2", optional = true }
 redis = { version = "0.22", features = ["tokio-comp"] }
 reqwest = { version = "0.11", default-features = false, features = ["rustls-tls-native-roots"], optional = true }
 secp256k1 = { version = "0.25", features = ["rand-std"] }
@@ -49,4 +50,4 @@ variable_transfer = ["extensions", "remote_storage", "hyper", "jsonwebtoken", "r
 registry = []
 policy_module = []
 instant_server = ["reqwest"]
-storage_macro = ["async-recursion"]
+storage_macro = ["async-recursion", "rdbc2"]

src/extensions/instant_server.rs

@@ -81,7 +81,7 @@ impl InstantServer {
             .join(instant_server_id.clone());
         std::fs::create_dir_all(&working_dir).unwrap();
         let mut user_init_config_file =
-            std::fs::File::create(&Path::new(&working_dir).join("user_init_config.toml")).unwrap();
+            std::fs::File::create(Path::new(&working_dir).join("user_init_config.toml")).unwrap();
         user_init_config_file
             .write_all(user_init_config.as_bytes())
             .unwrap();

src/extensions/storage_macro.rs

@@ -1,8 +1,10 @@
+use crate::StorageEntry;
+
 mod append;
 mod chunk;
+mod dbc;
 mod fs;
 mod redis;
-use crate::StorageEntry;
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
@@ -64,6 +66,7 @@ impl crate::application::CoLink {
         match macro_type.as_str() {
             "chunk" => self._read_entry_chunk(&string_before).await,
             "redis" => self._read_entry_redis(&string_before, &string_after).await,
+            "dbc" => self._read_entry_dbc(&string_before, &string_after).await,
             "fs" => self._read_entry_fs(&string_before, &string_after).await,
             _ => Err(format!(
                 "invalid storage macro, found {} in key name {}",
@@ -136,7 +139,7 @@ impl crate::application::CoLink {
                     "invalid storage macro, found {} in prefix {}",
                     macro_type, key_name_prefix
                 )
-                .into())
+                .into());
             }
         };
         let mut res: Vec<StorageEntry> = Vec::new();

src/extensions/storage_macro/dbc.rs

@@ -0,0 +1,52 @@
+use async_recursion::async_recursion;
+use rdbc2;
+
+type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
+
+impl crate::application::CoLink {
+    #[async_recursion]
+    async fn _search_and_generate_query(
+        &self,
+        string_before_dbc: &str,
+        string_after_dbc: &str,
+    ) -> Result<(String, Vec<String>), Error> {
+        let split_key_path: Vec<&str> = string_after_dbc.split(':').collect();
+        for i in (0..split_key_path.len()).rev() {
+            let current_key_path =
+                format!("{}:{}", string_before_dbc, split_key_path[0..=i].join(":"));
+            let payload = self.read_entry(current_key_path.as_str()).await;
+            if let Ok(payload) = payload {
+                let query_string = String::from_utf8(payload)?;
+                let count = query_string.matches('?').count();
+                if count != split_key_path.len() - (i + 1) {
+                    return Err("Number of parameters does not match specified query string")?;
+                }
+                let params = split_key_path[(i + 1)..]
+                    .iter()
+                    .map(|x| x.to_string())
+                    .collect::<Vec<String>>();
+                return Ok((query_string, params));
+            }
+        }
+        Err("no query string found.")?
+    }
+
+    #[async_recursion]
+    pub(crate) async fn _read_entry_dbc(
+        &self,
+        string_before_dbc: &str,
+        string_after_dbc: &str,
+    ) -> Result<Vec<u8>, Error> {
+        let url_key = format!("{}:url", string_before_dbc);
+        let url = self.read_entry(url_key.as_str()).await?;
+        let url_string = String::from_utf8(url)?;
+        let (query_string, params) = self
+            ._search_and_generate_query(string_before_dbc, string_after_dbc)
+            .await?;
+        let params: Vec<&str> = params.iter().map(AsRef::as_ref).collect();
+        let mut database = rdbc2::dbc::Database::new(url_string.as_str())?;
+        let result = database.execute_query_with_params(query_string.as_str(), &params)?;
+        let seralized_result = serde_json::to_vec(&result)?;
+        Ok(seralized_result)
+    }
+}

tests/test_storage_macro_dbc.rs

@@ -0,0 +1,99 @@
+use common::*;
+mod common;
+
+fn _get_mysql_connection_url() -> String {
+    if std::env::var("MYSQL_DATABASE_URL").is_ok() {
+        std::env::var("MYSQL_DATABASE_URL").unwrap()
+    } else {
+        panic!("Please set the environment variable MYSQL_DATABASE_URL to run this test.")
+    }
+}
+
+#[tokio::test]
+async fn test_storage_macro_dbc_mysql(
+) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+    let (_ir, _is, cl) = set_up_test_env_single_user().await?;
+
+    cl.create_entry(
+        "storage_macro_test:db:url",
+        _get_mysql_connection_url().as_bytes(),
+    )
+    .await?;
+
+    cl.create_entry(
+        "storage_macro_test:db:create_db",
+        b"CREATE DATABASE IF NOT EXISTS test_db" as &[u8],
+    )
+    .await?;
+
+    // Create a table and insert dummy data
+    cl.create_entry(
+        "storage_macro_test:db:create_table",
+        b"CREATE TABLE IF NOT EXISTS users (name VARCHAR(255), age INT)" as &[u8],
+    )
+    .await?;
+
+    cl.create_entry(
+        "storage_macro_test:db:insert_data",
+        b"INSERT INTO users VALUES ('Alice', 20)" as &[u8],
+    )
+    .await?;
+
+    cl.create_entry(
+        "storage_macro_test:db:query_users",
+        b"SELECT * FROM users WHERE name = ? AND age = ?" as &[u8],
+    )
+    .await?;
+
+    cl.create_entry(
+        "storage_macro_test:db:cleanup",
+        b"DROP TABLE IF EXISTS users" as &[u8],
+    )
+    .await?;
+
+    cl.read_entry("storage_macro_test:db:$dbc:create_db")
+        .await?;
+    cl.read_entry("storage_macro_test:db:$dbc:create_table")
+        .await?;
+    cl.read_entry("storage_macro_test:db:$dbc:insert_data")
+        .await?;
+    let query_result = cl
+        .read_entry("storage_macro_test:db:$dbc:query_users:'Alice':'20'")
+        .await?;
+
+    let stringified = String::from_utf8(query_result.clone())?;
+    println!("{}", stringified);
+    assert_eq!(
+        stringified,
+        r#"{"rows":[{"values":[{"Bytes":"QWxpY2U"},{"Int":20}],"columns":[{"name":"name","column_type":"VARCHAR"},{"name":"age","column_type":"INT"}]}],"affected_row_count":0}"#
+    );
+
+    let deserialized: rdbc2::dbc::QueryResult = serde_json::from_slice(&query_result)?;
+    assert_eq!(deserialized.rows.len(), 1);
+
+    // Test query string parsing order
+    cl.create_entry(
+        "storage_macro_test:db:query_users2",
+        b"SELECT * FROM users WHERE name = ? AND age = ?" as &[u8],
+    )
+    .await?;
+
+    cl.create_entry(
+        "storage_macro_test:db:query_users2:additional",
+        b"SELECT * FROM users WHERE name = ?" as &[u8],
+    )
+    .await?;
+
+    let result = cl
+        .read_entry("storage_macro_test:db:$dbc:query_users2:additional:'Alice'")
+        .await?;
+
+    assert_eq!(
+        String::from_utf8(result)?,
+        r#"{"rows":[{"values":[{"Bytes":"QWxpY2U"},{"Int":20}],"columns":[{"name":"name","column_type":"VARCHAR"},{"name":"age","column_type":"INT"}]}],"affected_row_count":0}"#
+    );
+
+    cl.read_entry("storage_macro_test:db:$dbc:cleanup").await?;
+
+    Ok(())
+}