From ddf2ad23e1c1f38c54530ead364a5ca62b9c8b9b Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@panodata.org>
Date: Mon, 29 Jan 2024 03:38:09 +0100
Subject: [PATCH] Add software tests using Polars

Polars does not support writing to databases yet, nor does it support
SQLAlchemy's AsyncEngine. So, this exercises `pl.read_database()` only,
together with the `crate` and `postgresql+psycopg_relaxed` dialects,
using urllib3 resp. Psycopg3.
---
 pyproject.toml                    |  1 +
 tests/test_cratedb_polars_read.py | 43 +++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 tests/test_cratedb_polars_read.py

diff --git a/pyproject.toml b/pyproject.toml
index d099366..194c844 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -79,6 +79,7 @@ release = [
 test = [
   "crate[sqlalchemy]",
   "pandas<2.3",
+  "polars[pyarrow]<0.21",
   "pytest<9",
   "pytest-asyncio<1",
   "pytest-cov<5",
diff --git a/tests/test_cratedb_polars_read.py b/tests/test_cratedb_polars_read.py
new file mode 100644
index 0000000..6f6dd9b
--- /dev/null
+++ b/tests/test_cratedb_polars_read.py
@@ -0,0 +1,43 @@
+import sys
+
+import polars as pl
+import pytest
+import sqlalchemy as sa
+from polars.testing import assert_frame_equal
+
+REFERENCE_FRAME = pl.from_records([{"mountain": "Mont Blanc", "height": 4808}])
+SQL_SELECT_STATEMENT = "SELECT mountain, height FROM sys.summits ORDER BY height DESC LIMIT 1;"
+
+
+if sys.version_info < (3, 8):
+    pytest.skip("Does not work on Python 3.7", allow_module_level=True)
+
+
+def test_crate_read_sql(cratedb_http_host, cratedb_http_port):
+    engine = sa.create_engine(
+        url=f"crate://{cratedb_http_host}:{cratedb_http_port}",
+        echo=True,
+    )
+    conn = engine.connect()
+    df = pl.read_database(
+        query=SQL_SELECT_STATEMENT,
+        connection=conn,
+        schema_overrides={"value": pl.Utf8},
+    )
+    assert_frame_equal(df, REFERENCE_FRAME)
+
+
+def test_psycopg_read_sql(cratedb_psql_host, cratedb_psql_port):
+    engine = sa.create_engine(
+        url=f"postgresql+psycopg_relaxed://crate@{cratedb_psql_host}:{cratedb_psql_port}",
+        isolation_level="AUTOCOMMIT",
+        use_native_hstore=False,
+        echo=True,
+    )
+    conn = engine.connect()
+    df = pl.read_database(
+        query=SQL_SELECT_STATEMENT,
+        connection=conn,
+        schema_overrides={"value": pl.Utf8},
+    )
+    assert_frame_equal(df, REFERENCE_FRAME)