From 5743d55943b639e4a3d44611bec5e26dd37d84be Mon Sep 17 00:00:00 2001
From: Xiaoying Wang <xiaoying_wang@sfu.ca>
Date: Sun, 21 Apr 2024 19:18:10 -0700
Subject: [PATCH] add test

---
 Justfile                                      |  9 ++++++
 connectorx-python/Cargo.toml                  |  5 +--
 .../connectorx/tests/test_postgres.py         | 15 ++++++---
 connectorx-python/examples/test_aaa.rs        | 31 +++++++++++++++++++
 4 files changed, 54 insertions(+), 6 deletions(-)
 create mode 100644 connectorx-python/examples/test_aaa.rs

diff --git a/Justfile b/Justfile
index 07ac50601d..f51188dad7 100644
--- a/Justfile
+++ b/Justfile
@@ -63,6 +63,15 @@ seed-db-more:
     cat scripts/oracle.sql | sqlplus $ORACLE_URL_SCRIPT
     mysql --protocol tcp -h$MARIADB_HOST -P$MARIADB_PORT -u$MARIADB_USER -p$MARIADB_PASSWORD $MARIADB_DB < scripts/mysql.sql
 
+aaa2 conn="POSTGRES_URL":
+    cd connectorx-python && PYO3_PYTHON=$HOME/.pyenv/versions/3.8.12/bin/python3.8 PYTHONPATH=$HOME/.pyenv/versions/conn/lib/python3.8/site-packages LD_LIBRARY_PATH=$HOME/.pyenv/versions/3.8.12/lib/ cargo run --no-default-features --features aaa --example test_aaa
+
+aaa conn="POSTGRES_URL":
+    cd connectorx-python && PYO3_PYTHON=$HOME/.pyenv/versions/3.12.2/bin/python3.12 PYTHONPATH=$HOME/.pyenv/versions/conn/lib/python3.12/site-packages LD_LIBRARY_PATH=$HOME/.pyenv/versions/3.12.2/lib/ cargo run --no-default-features --features aaa --example test_aaa
+
+aaa-d conn="POSTGRES_URL":
+    cd connectorx-python && RUST_BACKTRACE=1 PYO3_PYTHON=$HOME/.pyenv/versions/3.12.2/bin/python3.12 PYTHONPATH=$HOME/.pyenv/versions/conn/lib/python3.12/site-packages LD_LIBRARY_PATH=$HOME/.pyenv/versions/3.12.2/lib/ rust-lldb target/debug/examples/test_aaa
+
 # benches 
 flame-tpch conn="POSTGRES_URL":
     cd connectorx-python && PYO3_PYTHON=$HOME/.pyenv/versions/3.8.6/bin/python3.8 PYTHONPATH=$HOME/.pyenv/versions/conn/lib/python3.8/site-packages LD_LIBRARY_PATH=$HOME/.pyenv/versions/3.8.6/lib/ cargo run --no-default-features --features executable --features fptr --features nbstr --features dsts --features srcs --release --example flame_tpch {{conn}}
diff --git a/connectorx-python/Cargo.toml b/connectorx-python/Cargo.toml
index f67db1effa..0accc088f3 100644
--- a/connectorx-python/Cargo.toml
+++ b/connectorx-python/Cargo.toml
@@ -55,12 +55,13 @@ pprof = {version = "0.5", features = ["flamegraph", "criterion", "protobuf"]}
 rayon = "1"
 
 [lib]
-crate-type = ["cdylib"]
-name = "connectorx"
+crate-type = ["cdylib", "rlib"]
+name = "connectorx_python"
 
 [features]
 branch = ["connectorx/branch"]
 default = ["extension", "fptr", "nbstr", "dsts", "srcs", "federation", "fed_exec"]
+aaa = ["executable", "fptr", "nbstr", "dsts", "srcs", "federation", "fed_exec"]
 dsts = ["connectorx/dst_arrow", "connectorx/dst_arrow2"]
 executable = ["pyo3/auto-initialize"]
 extension = ["pyo3/extension-module"]
diff --git a/connectorx-python/connectorx/tests/test_postgres.py b/connectorx-python/connectorx/tests/test_postgres.py
index c72f59f1aa..b87a623240 100644
--- a/connectorx-python/connectorx/tests/test_postgres.py
+++ b/connectorx-python/connectorx/tests/test_postgres.py
@@ -43,10 +43,17 @@ def test_postgres_on_non_select(postgres_url: str) -> None:
     df = read_sql(postgres_url, query)
 
 def test_postgres_aaa(postgres_url: str) -> None:
-    query = "SELECT test_int, test_str FROM test_table"
-    # query = "SELECT test_bytea FROM test_types"
-    # query = "SELECT test_boolarray FROM test_types"
-    df = read_sql(postgres_url, query, partition_on="test_int", partition_num=2)
+    # query = "SELECT test_int, test_str FROM test_table"
+    # # query = "SELECT test_bytea FROM test_types"
+    # # query = "SELECT test_boolarray FROM test_types"
+    # df = read_sql(postgres_url, query, partition_on="test_int", partition_num=2)
+
+    queries = [
+        "SELECT test_str FROM test_table WHERE test_int < 3",
+        "SELECT test_str FROM test_table WHERE test_int >= 3",
+    ]
+
+    df = read_sql(postgres_url, query=queries)
     print(df)
 
 def test_postgres_aggregation(postgres_url: str) -> None:
diff --git a/connectorx-python/examples/test_aaa.rs b/connectorx-python/examples/test_aaa.rs
new file mode 100644
index 0000000000..2e7b95125b
--- /dev/null
+++ b/connectorx-python/examples/test_aaa.rs
@@ -0,0 +1,31 @@
+extern crate connectorx;
+extern crate env_logger;
+extern crate log;
+use connectorx_python::cx_read_sql::read_sql;
+use log::debug;
+use pyo3::Python;
+use std::env;
+
+fn main() {
+    env_logger::init();
+    debug!("[!]  debug");
+    let conn = env::var("POSTGRES_URL").unwrap();
+
+    let res = Python::with_gil(|py| {
+        read_sql(
+            py,
+            &conn,
+            "pandas",
+            None,
+            Some(vec![
+                String::from("select test_int from test_table where test_int > 3"),
+                // String::from("select test_int, test_str from test_table where test_int <= 3"),
+            ]),
+            None,
+        )
+        .unwrap();
+    });
+
+    println!("finish!");
+    println!("res: {:?}", res);
+}