From 7bd4f3d1f93cf49c3cac659021ac8b8cb30b28dd Mon Sep 17 00:00:00 2001
From: Dmitriy Pleshevskiy <dmitriy@ideascup.me>
Date: Fri, 19 Feb 2021 22:17:01 +0300
Subject: [PATCH] fix: merge query with params

---
 migra-cli/src/database.rs           | 42 ++++++++++++++++++++++++-----
 migra-cli/src/databases/mod.rs      |  5 +++-
 migra-cli/src/databases/postgres.rs | 24 ++++-------------
 3 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/migra-cli/src/database.rs b/migra-cli/src/database.rs
index 2692119..d79af1a 100644
--- a/migra-cli/src/database.rs
+++ b/migra-cli/src/database.rs
@@ -4,6 +4,8 @@ pub trait ToSql {
     fn to_sql(&self) -> String;
 }
 
+pub type ToSqlParams<'a> = &'a [&'a dyn ToSql];
+
 impl ToSql for &str {
     fn to_sql(&self) -> String {
         format!("'{}'", self)
@@ -21,11 +23,39 @@ pub trait OpenDatabaseConnection: Sized {
 pub trait DatabaseConnection {
     fn batch_execute(&mut self, query: &str) -> StdResult<()>;
 
-    fn execute<'b>(&mut self, query: &str, params: &'b [&'b dyn ToSql]) -> StdResult<u64>;
+    fn execute<'b>(&mut self, query: &str, params: ToSqlParams<'b>) -> StdResult<u64>;
 
-    fn query<'b>(
-        &mut self,
-        query: &str,
-        params: &'b [&'b dyn ToSql],
-    ) -> StdResult<Vec<Vec<String>>>;
+    fn query<'b>(&mut self, query: &str, params: ToSqlParams<'b>) -> StdResult<Vec<Vec<String>>>;
+}
+
+pub(crate) fn merge_query_with_params(query: &str, params: ToSqlParams) -> String {
+    params
+        .iter()
+        .enumerate()
+        .fold(query.to_string(), |acc, (i, p)| {
+            str::replace(&acc, &format!("${}", i + 1), &p.to_sql())
+        })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn merge_query_with_params_successfully() {
+        assert_eq!(
+            merge_query_with_params("SELECT $1", &[&"foo"]),
+            "SELECT 'foo'"
+        );
+
+        assert_eq!(
+            merge_query_with_params("SELECT $1, $2", &[&"foo", &"bar"]),
+            "SELECT 'foo', 'bar'"
+        );
+
+        assert_eq!(
+            merge_query_with_params("SELECT $1, $1", &[&"foo"]),
+            "SELECT 'foo', 'foo'"
+        );
+    }
 }
diff --git a/migra-cli/src/databases/mod.rs b/migra-cli/src/databases/mod.rs
index 2cea946..d4b41b5 100644
--- a/migra-cli/src/databases/mod.rs
+++ b/migra-cli/src/databases/mod.rs
@@ -17,7 +17,10 @@ impl DatabaseConnectionManager {
         }
     }
 
-    pub fn connect_with_string(&self, connection_string: &str) -> StdResult<Box<dyn DatabaseConnection>> {
+    pub fn connect_with_string(
+        &self,
+        connection_string: &str,
+    ) -> StdResult<Box<dyn DatabaseConnection>> {
         let conn = match self.config.client()? {
             SupportedDatabaseClient::Postgres => PostgresConnection::open(&connection_string)?,
         };
diff --git a/migra-cli/src/databases/postgres.rs b/migra-cli/src/databases/postgres.rs
index b744138..c5e610b 100644
--- a/migra-cli/src/databases/postgres.rs
+++ b/migra-cli/src/databases/postgres.rs
@@ -1,4 +1,4 @@
-use crate::database::{DatabaseConnection, OpenDatabaseConnection, ToSql};
+use crate::database::*;
 use crate::error::StdResult;
 use postgres::{Client, NoTls};
 
@@ -19,29 +19,15 @@ impl DatabaseConnection for PostgresConnection {
         Ok(())
     }
 
-    fn execute<'b>(&mut self, query: &str, params: &'b [&'b dyn ToSql]) -> StdResult<u64> {
-        let stmt = params
-            .iter()
-            .enumerate()
-            .fold(query.to_string(), |acc, (i, p)| {
-                str::replace(&acc, &format!("${}", i), &p.to_sql())
-            });
+    fn execute<'b>(&mut self, query: &str, params: ToSqlParams<'b>) -> StdResult<u64> {
+        let stmt = merge_query_with_params(query, params);
 
         let res = self.client.execute(stmt.as_str(), &[])?;
         Ok(res)
     }
 
-    fn query<'b>(
-        &mut self,
-        query: &str,
-        params: &'b [&'b dyn ToSql],
-    ) -> StdResult<Vec<Vec<String>>> {
-        let stmt = params
-            .iter()
-            .enumerate()
-            .fold(query.to_string(), |acc, (i, p)| {
-                str::replace(&acc, &format!("${}", i), &p.to_sql())
-            });
+    fn query<'b>(&mut self, query: &str, params: ToSqlParams<'b>) -> StdResult<Vec<Vec<String>>> {
+        let stmt = merge_query_with_params(query, params);
 
         let res = self.client.query(stmt.as_str(), &[])?;