substantial mysql tests and sanitize parameters

Author: nociza

src/dbc.rs

@@ -1,12 +1,11 @@
 use std::sync::Arc;
 
 use serde::{Deserialize, Serialize};
-
 use serde_json;
 
 mod mysql;
-mod sqlite;
 mod postgres;
+mod sqlite;
 
 pub type Error = Box<dyn std::error::Error + Send + Sync>;
 
@@ -47,15 +46,31 @@ impl Database {
         Ok(serde_json::to_vec(&result)?)
     }
 
-    pub fn execute_query_with_params(&mut self, query: &str, params: &[&str]) -> Result<QueryResult, Error> {
+    pub fn execute_query_with_params(
+        &mut self,
+        query: &str,
+        params: &[&str],
+    ) -> Result<QueryResult, Error> {
         let mut query = query.to_string();
         for param in params {
-            query = query.replace("?", param);
+            let quoted_param = if let Some(_) = param.strip_prefix('\'') {
+                // If the parameter already has single quotes, don't add them again.
+                // This avoids SQL injection vulnerabilities when the parameter contains a quote.
+                param.to_string()
+            } else {
+                // If the parameter doesn't have single quotes, add them.
+                format!("'{}'", param)
+            };
+            query = query.replacen("?", quoted_param.as_str(), 1);
         }
         self.execute_query(&query)
     }
 
-    pub fn execute_query_with_params_and_serialize(&mut self, query: &str, params: &[&str]) -> Result<String, Error> {
+    pub fn execute_query_with_params_and_serialize(
+        &mut self,
+        query: &str,
+        params: &[&str],
+    ) -> Result<String, Error> {
         let result = self.execute_query_with_params(query, params)?;
         Ok(serde_json::to_string(&result)?)
     }
@@ -89,6 +104,31 @@ pub struct Row {
     columns: Arc<[Column]>,
 }
 
+impl Row {
+    pub fn new(values: Vec<Value>, columns: Arc<[Column]>) -> Self {
+        Row { values, columns }
+    }
+
+    pub fn get_value(&self, index: usize) -> Option<&Value> {
+        self.values.get(index)
+    }
+
+    pub fn get_column(&self, index: usize) -> Option<&Column> {
+        self.columns.get(index)
+    }
+
+    pub fn get_value_by_name(&self, name: &str) -> Option<&Value> {
+        self.columns
+            .iter()
+            .position(|column| column.name == name)
+            .and_then(|index| self.values.get(index))
+    }
+
+    pub fn get_column_by_name(&self, name: &str) -> Option<&Column> {
+        self.columns.iter().find(|column| column.name == name)
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug)]
 pub struct QueryResult {
     pub rows: Vec<Row>,

tests/mysql_test.rs

@@ -1,7 +1,32 @@
+use mysql;
+
 use rdbc2::dbc;
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
+fn _prepare_mysql_database(url: String) -> Result<(), Error> {
+    let mut database = dbc::Database::new(url.as_str())?;
+
+    let query = "DROP DATABASE IF EXISTS test";
+    database.execute_query(query)?;
+    let query = "CREATE DATABASE IF NOT EXISTS test";
+    database.execute_query(query)?;
+    let query = "USE test";
+    database.execute_query(query)?;
+    let query = "CREATE TABLE IF NOT EXISTS test_table (id INT NOT NULL AUTO_INCREMENT, name VARCHAR(255) NOT NULL, PRIMARY KEY (id))";
+    database.execute_query(query)?;
+
+    Ok(())
+}
+
+fn _cleanup_mysql_database(url: String) -> Result<(), Error> {
+    let mut database = dbc::Database::new(url.as_str())?;
+    let query = "DROP DATABASE IF EXISTS test";
+    database.execute_query(query)?;
+
+    Ok(())
+}
+
 fn _get_mysql_connection_url() -> String {
     if std::env::var("MYSQL_DATABASE_URL").is_ok() {
         std::env::var("MYSQL_DATABASE_URL").unwrap()
@@ -13,32 +38,108 @@ fn _get_mysql_connection_url() -> String {
 #[tokio::test]
 async fn test_mysql_simple_query() -> Result<(), Error> {
     let url = _get_mysql_connection_url();
+    _prepare_mysql_database(url.clone())?;
+
     let mut database = dbc::Database::new(url.as_str())?;
-    let query = "SELECT 1";
-    let result = database.execute_query(query)?;
-    assert_eq!(result.rows.len(), 1);
+
+    // Use the test database
+    let use_query = "USE test";
+    database.execute_query(use_query)?;
+
+    // Insert two rows into the test_table
+    let insert_query = "INSERT INTO test_table (name) VALUES ('test1'), ('test2')";
+    database.execute_query(insert_query)?;
+
+    // Select all rows from test_table
+    let select_query = "SELECT * FROM test_table";
+    let result = database.execute_query(select_query)?;
+
+    assert_eq!(result.rows.len(), 2);
+
+    // Verify the data returned by the query
+    let first_row = &result.rows[0];
+    let second_row = &result.rows[1];
+    assert_eq!(
+        first_row.get_value_by_name("name"),
+        Some(&dbc::Value::Bytes("test1".to_owned().into_bytes()))
+    );
+    assert_eq!(
+        second_row.get_value_by_name("name"),
+        Some(&dbc::Value::Bytes("test2".to_owned().into_bytes()))
+    );
+
+    _cleanup_mysql_database(url.clone())?;
 
     Ok(())
 }
 
 #[tokio::test]
 async fn test_mysql_query_with_params() -> Result<(), Error> {
     let url = _get_mysql_connection_url();
+    _prepare_mysql_database(url.clone())?;
+
     let mut database = dbc::Database::new(url.as_str())?;
-    let query = "SELECT ? + ?";
-    let result = database.execute_query_with_params(query, &["1", "2"])?;
+
+    // Use the test database
+    let use_query = "USE test";
+    database.execute_query(use_query)?;
+
+    // Insert two rows into the test_table
+    let insert_query = "INSERT INTO test_table (name) VALUES ('test1'), ('test2')";
+    database.execute_query(insert_query)?;
+
+    // Select all rows from test_table
+    let select_query = "SELECT * FROM test_table WHERE name = ?";
+    let result = database.execute_query_with_params(select_query, &["test1"])?;
+
     assert_eq!(result.rows.len(), 1);
 
+    // Verify the data returned by the query
+    let first_row = &result.rows[0];
+    assert_eq!(
+        first_row.get_value_by_name("name"),
+        Some(&dbc::Value::Bytes("test1".to_owned().into_bytes()))
+    );
+
+    _cleanup_mysql_database(url.clone())?;
+
     Ok(())
 }
 
 #[tokio::test]
 async fn test_mysql_query_with_params_and_serialize() -> Result<(), Error> {
     let url = _get_mysql_connection_url();
+    _prepare_mysql_database(url.clone())?;
+
     let mut database = dbc::Database::new(url.as_str())?;
-    let query = "SELECT ? + ?";
-    let result = database.execute_query_with_params_and_serialize(query, &["1", "2"])?;
-    assert_eq!(result, r#"{"rows":[{"values":[{"Bytes":[50]}],"columns":[{"name":"1 + 1","column_type":"LONGLONG"}]}]}"#);
+
+    // Use the test database
+    let use_query = "USE test";
+    database.execute_query(use_query)?;
+
+    // Insert two rows into the test_table
+    let insert_query = "INSERT INTO test_table (name) VALUES ('test1'), ('test2')";
+    database.execute_query(insert_query)?;
+
+    // Update the test_table to set the name of the first row to "updated"
+    let update_query = "UPDATE test_table SET name = ? WHERE id = ?";
+    let result = database.execute_query_with_params(update_query, &["updated", "1"])?;
+
+    // Select all rows from test_table
+    let select_query = "SELECT * FROM test_table WHERE id = ?";
+    let result = database.execute_query_with_params_and_serialize(select_query, &["1"])?;
+    let expected_result = r#"{"rows":[{"values":[{"Bytes":[49]},{"Bytes":[117,112,100,97,116,101,100]}],"columns":[{"name":"id","column_type":"LONG"},{"name":"name","column_type":"STRING"}]}]}"#;
+    assert_eq!(result, expected_result);
+
+    // deserialize the result and verify the data
+    let deserialized_result: dbc::QueryResult = serde_json::from_str(&result)?;
+    let first_row = &deserialized_result.rows[0];
+    assert_eq!(
+        first_row.get_value_by_name("name"),
+        Some(&dbc::Value::Bytes("updated".to_owned().into_bytes()))
+    );
+
+    _cleanup_mysql_database(url.clone())?;
 
     Ok(())
-}
+}