Tests: Add software tests for CrateDB/pandas

Author: amotl

.github/workflows/tests.yml

@@ -32,6 +32,7 @@ jobs:
       cratedb:
         image: crate/crate:nightly
         ports:
+          - 4200:4200
           - 5432:5432
 
     env:

pyproject.toml

@@ -77,6 +77,8 @@ release = [
   "twine<5",
 ]
 test = [
+  "crate[sqlalchemy]",
+  "pandas<2.1",
   "pytest<8",
   "pytest-asyncio<1",
   "pytest-cov<5",
@@ -122,7 +124,7 @@ show_missing = true
 
 [tool.ruff]
 line-length = 120
-
+extend-ignore = ["PD901"]
 select = [
   # Bandit
   "S",

tests/conftest.py

@@ -0,0 +1,46 @@
+import pytest
+
+
+def pytest_addoption(parser):
+    """
+    Register custom options to support invocation by cr8.
+    Use cr8 to invoke a pytest test suite, see `run.sh`.
+
+    Example::
+
+        pytest -vvv --http-host 127.0.0.1 --http-port 4200 --psql-host 127.0.0.1 --psql-port 5432
+
+    https://github.com/mfussenegger/cr8
+
+    TODO: Refactor to `cratedb-toolkit` or `pytest-cratedb` in the long run.
+    """
+    parser.addoption("--http-url", action="store", default="localhost:4200")
+    parser.addoption("--http-host", action="store", default="localhost")
+    parser.addoption("--http-port", action="store", default="4200")
+    parser.addoption("--psql-host", action="store", default="localhost")
+    parser.addoption("--psql-port", action="store", default="5432")
+
+
+@pytest.fixture
+def cratedb_http_url(pytestconfig):
+    return pytestconfig.getoption("--http-url")
+
+
+@pytest.fixture
+def cratedb_http_host(pytestconfig):
+    return pytestconfig.getoption("--http-host")
+
+
+@pytest.fixture
+def cratedb_http_port(pytestconfig):
+    return pytestconfig.getoption("--http-port")
+
+
+@pytest.fixture
+def cratedb_psql_host(pytestconfig):
+    return pytestconfig.getoption("--psql-host")
+
+
+@pytest.fixture
+def cratedb_psql_port(pytestconfig):
+    return pytestconfig.getoption("--psql-port")

tests/test_cratedb_pandas_read.py

@@ -0,0 +1,68 @@
+import pandas as pd
+import pytest
+import sqlalchemy as sa
+from pandas._testing import assert_frame_equal
+from sqlalchemy.ext.asyncio import create_async_engine
+
+REFERENCE_FRAME = pd.DataFrame.from_records([{"mountain": "Mont Blanc", "height": 4808}])
+SQL_SELECT_STATEMENT = "SELECT mountain, height FROM sys.summits ORDER BY height DESC LIMIT 1;"
+
+
+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 = pd.read_sql(sql=sa.text(SQL_SELECT_STATEMENT), con=conn)
+    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 = pd.read_sql(sql=sa.text(SQL_SELECT_STATEMENT), con=conn)
+    assert_frame_equal(df, REFERENCE_FRAME)
+
+
+@pytest.mark.asyncio
+async def test_psycopg_async_read_sql(cratedb_psql_host, cratedb_psql_port):
+    engine = create_async_engine(
+        url=f"postgresql+psycopg_relaxed://crate@{cratedb_psql_host}:{cratedb_psql_port}",
+        isolation_level="AUTOCOMMIT",
+        use_native_hstore=False,
+        echo=True,
+    )
+
+    async with engine.begin() as conn:
+        df = await conn.run_sync(read_sql_sync, sa.text(SQL_SELECT_STATEMENT))
+        assert_frame_equal(df, REFERENCE_FRAME)
+
+
+@pytest.mark.asyncio
+async def test_asyncpg_read_sql(cratedb_psql_host, cratedb_psql_port):
+    engine = create_async_engine(
+        url=f"postgresql+asyncpg_relaxed://crate@{cratedb_psql_host}:{cratedb_psql_port}",
+        isolation_level="AUTOCOMMIT",
+        echo=True,
+    )
+
+    async with engine.begin() as conn:
+        df = await conn.run_sync(read_sql_sync, sa.text(SQL_SELECT_STATEMENT))
+        assert_frame_equal(df, REFERENCE_FRAME)
+
+
+def read_sql_sync(conn, stmt):
+    """
+    Making pd.read_sql connection the first argument to make it compatible
+    with conn.run_sync(), see https://stackoverflow.com/a/70861276.
+    """
+    return pd.read_sql(stmt, conn)

tests/test_cratedb_pandas_write.py

@@ -0,0 +1,104 @@
+import pandas as pd
+import pytest
+import sqlalchemy as sa
+from pandas._testing import assert_frame_equal
+from sqlalchemy.ext.asyncio import create_async_engine
+
+INPUT_FRAME = pd.DataFrame.from_records([{"id": 1, "name": "foo", "value": 42.42}])
+SQL_SELECT_STATEMENT = "SELECT * FROM doc.foo;"
+SQL_REFRESH_STATEMENT = "REFRESH TABLE doc.foo;"
+
+
+pytest.skip("Does not work on Python 3.7", allow_module_level=True)
+
+
+def test_crate_to_sql(cratedb_http_host, cratedb_http_port):
+    # Connect to database.
+    engine = sa.create_engine(
+        url=f"crate://{cratedb_http_host}:{cratedb_http_port}",
+        echo=True,
+    )
+    con = engine.connect()
+
+    # Insert data using pandas.
+    df = INPUT_FRAME
+    retval = df.to_sql(name="foo", con=con, if_exists="replace", index=False)
+    assert retval == -1
+
+    # Synchronize table content.
+    con.execute(sa.text(SQL_REFRESH_STATEMENT))
+
+    # Read back and verify data using pandas.
+    df = pd.read_sql(sql=sa.text(SQL_SELECT_STATEMENT), con=con)
+    assert_frame_equal(df, INPUT_FRAME)
+
+
+@pytest.mark.skip(reason="Needs COLLATE and pg_table_is_visible")
+def test_psycopg_to_sql(cratedb_psql_host, cratedb_psql_port):
+    # Connect to database.
+    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()
+
+    # Insert data using pandas.
+    df = INPUT_FRAME
+    retval = df.to_sql(name="foo", con=conn, if_exists="replace", index=False)
+    assert retval == -1
+
+    # Synchronize table content.
+    conn.execute(sa.text(SQL_REFRESH_STATEMENT))
+
+    # Read back and verify data using pandas.
+    df = pd.read_sql(sql=sa.text(SQL_SELECT_STATEMENT), con=conn)
+    assert_frame_equal(df, INPUT_FRAME)
+
+
+@pytest.mark.skip(reason="Needs COLLATE and pg_table_is_visible")
+@pytest.mark.asyncio
+async def test_psycopg_async_to_sql(cratedb_psql_host, cratedb_psql_port):
+    # Connect to database.
+    engine = create_async_engine(
+        url=f"postgresql+psycopg_relaxed://crate@{cratedb_psql_host}:{cratedb_psql_port}",
+        isolation_level="AUTOCOMMIT",
+        use_native_hstore=False,
+        echo=True,
+    )
+
+    # Insert data using pandas.
+    async with engine.begin() as conn:
+        df = INPUT_FRAME
+        retval = await conn.run_sync(to_sql_sync, df=df, name="foo", if_exists="replace", index=False)
+        assert retval == -1
+
+    # TODO: Read back dataframe and compare with original.
+
+
+@pytest.mark.skip(reason="Needs COLLATE and pg_table_is_visible")
+@pytest.mark.asyncio
+async def test_asyncpg_to_sql(cratedb_psql_host, cratedb_psql_port):
+    # Connect to database.
+    engine = create_async_engine(
+        url=f"postgresql+asyncpg_relaxed://crate@{cratedb_psql_host}:{cratedb_psql_port}",
+        isolation_level="AUTOCOMMIT",
+        echo=True,
+    )
+
+    # Insert data using pandas.
+    async with engine.begin() as conn:
+        df = INPUT_FRAME
+        retval = await conn.run_sync(to_sql_sync, df=df, name="foo", if_exists="replace", index=False)
+        assert retval == -1
+
+    # TODO: Read back dataframe and compare with original.
+
+
+def to_sql_sync(conn, df, name, **kwargs):
+    """
+    Making df.to_sql connection the first argument to make it compatible
+    with conn.run_sync(), see https://stackoverflow.com/a/70861276.
+    """
+    return df.to_sql(name=name, con=conn, **kwargs)