test: Expand passthrough test cases

Author: tobyhede

packages/cipherstash-proxy-integration/src/common.rs

@@ -36,6 +36,9 @@ pub async fn clear() {
 
     let sql = "TRUNCATE encrypted";
     client.simple_query(sql).await.unwrap();
+
+    let sql = "TRUNCATE plaintext";
+    client.simple_query(sql).await.unwrap();
 }
 
 pub async fn reset_schema() {

packages/cipherstash-proxy-integration/src/passthrough.rs

@@ -1,13 +1,15 @@
 #[cfg(test)]
 mod tests {
-    use crate::common::{connect_with_tls, id, random_string, PROXY};
+    use crate::common::{clear, connect_with_tls, id, random_string, PROXY};
     use rand::Rng;
     use std::error::Error;
 
     #[tokio::test]
     async fn passthrough_statement() {
         let client = connect_with_tls(PROXY).await;
 
+        clear().await;
+
         let id = id();
         let encrypted_text = "hello@cipherstash.com";
 
@@ -29,6 +31,8 @@ mod tests {
     async fn passthrough_invalid_statement() {
         let client = connect_with_tls(PROXY).await;
 
+        clear().await;
+
         let sql = "SELECT * FROM blahvtha";
         let result = client.query(sql, &[]).await;
 
@@ -79,4 +83,145 @@ mod tests {
             tokio::time::sleep(std::time::Duration::from_millis(sleep_duration)).await;
         }
     }
+
+    #[tokio::test]
+    async fn passthrough_insert_from_select() {
+        let client = connect_with_tls(PROXY).await;
+
+        clear().await;
+
+        // Setup data
+        let id_1 = id();
+        let plaintext = "hello@cipherstash.com";
+
+        let sql = "INSERT INTO plaintext (id, plaintext) VALUES ($1, $2)";
+        client.query(sql, &[&id_1, &plaintext]).await.unwrap();
+
+        let id_2 = id_1 + 1;
+
+        // Insert value is selected from the record we just created
+        let select =
+            "SELECT id + 1, plaintext FROM plaintext WHERE id = ANY( ARRAY[ $1::Int8, $2::Int8 ] )";
+        let sql = format!("INSERT INTO plaintext (id, plaintext) {select} ON CONFLICT DO NOTHING");
+        client.query(&sql, &[&id_1, &id_2]).await.unwrap();
+
+        let sql = "SELECT id, plaintext FROM plaintext WHERE id = $1";
+        let rows = client.query(sql, &[&id_2]).await.unwrap();
+
+        assert_eq!(rows.len(), 1);
+
+        for row in rows {
+            let result: String = row.get("plaintext");
+            assert_eq!(plaintext, result);
+        }
+    }
+
+    #[tokio::test]
+    async fn passthrough_insert_with_value_from_select() {
+        let client = connect_with_tls(PROXY).await;
+
+        clear().await;
+
+        // Setup data
+        let id_1 = id();
+        let plaintext = "hello@cipherstash.com";
+
+        let sql = "INSERT INTO plaintext (id, plaintext) VALUES ($1, $2)";
+        client.query(sql, &[&id_1, &plaintext]).await.unwrap();
+
+        let id_2 = id();
+
+        // Insert value is selected from the record we just created
+        let select = "SELECT plaintext FROM plaintext WHERE id = $2";
+        let sql = format!("INSERT INTO plaintext (id, plaintext) VALUES ($1, ({select}))");
+        client.query(&sql, &[&id_2, &id_1]).await.unwrap();
+
+        let sql = "SELECT id, plaintext FROM plaintext WHERE id = $1";
+        let rows = client.query(sql, &[&id_2]).await.unwrap();
+
+        assert_eq!(rows.len(), 1);
+
+        for row in rows {
+            let result: String = row.get("plaintext");
+            assert_eq!(plaintext, result);
+        }
+    }
+
+    #[tokio::test]
+    async fn passthrough_insert_with_returning() {
+        let client = connect_with_tls(PROXY).await;
+
+        clear().await;
+
+        // Setup data
+        let id = id();
+        let plaintext = "hello@cipherstash.com";
+
+        let sql = "INSERT INTO plaintext (id, plaintext) VALUES ($1, $2) RETURNING *";
+        let rows = client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        assert_eq!(rows.len(), 1);
+
+        for row in rows {
+            let result: String = row.get("plaintext");
+            assert_eq!(plaintext, result);
+        }
+    }
+
+    #[tokio::test]
+    async fn passthrough_select_with_cardinality() {
+        let client = connect_with_tls(PROXY).await;
+
+        clear().await;
+
+        // Setup data
+        let id_1 = id();
+        let plaintext = "hello@cipherstash.com";
+
+        let sql = "INSERT INTO plaintext (id, plaintext) VALUES ($1, $2)";
+        client.query(sql, &[&id_1, &plaintext]).await.unwrap();
+
+        let id_2 = id();
+
+        let sql = "INSERT INTO plaintext (id) VALUES ($1)";
+        client.query(sql, &[&id_2]).await.unwrap();
+
+        let sql = "SELECT ARRAY_REMOVE(ARRAY_AGG(id), NULL), plaintext
+                         FROM plaintext
+                         WHERE CARDINALITY(ARRAY[1,2]) <> 0
+                         GROUP BY plaintext";
+
+        let rows = client.query(sql, &[]).await.unwrap();
+
+        assert_eq!(rows.len(), 2);
+    }
+
+    #[tokio::test]
+    async fn passthrough_delete_with_select() {
+        let client = connect_with_tls(PROXY).await;
+
+        clear().await;
+
+        // Setup data
+        let id_1 = id();
+        let plaintext = "one@cipherstash.com";
+
+        let sql = "INSERT INTO plaintext (id, plaintext) VALUES ($1, $2)";
+        client.query(sql, &[&id_1, &plaintext]).await.unwrap();
+
+        let id_2 = id();
+        let plaintext = "two@cipherstash.com";
+
+        let sql = "INSERT INTO plaintext (id, plaintext) VALUES ($1, $2)";
+        client.query(sql, &[&id_2, &plaintext]).await.unwrap();
+
+        let sql = "DELETE FROM plaintext
+                         WHERE id IN (SELECT id FROM plaintext WHERE plaintext = $1)";
+        client.query(sql, &[&plaintext]).await.unwrap();
+
+        let sql = "SELECT * FROM plaintext WHERE plaintext = $1";
+        let rows = client.query(sql, &[&plaintext]).await.unwrap();
+
+        assert_eq!(rows.len(), 0);
+    }
 }