diff --git a/pyproject.toml b/pyproject.toml index b1c0971..b3cf18e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,6 +79,7 @@ release = [ test = [ "crate[sqlalchemy]", "pandas<2.1", + "polars[pyarrow]<0.21", "pytest<8", "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)