test_adbc.py

import datetime
import sys
from pathlib import Path

import adbc_driver_manager.dbapi
import numpy as np
import pyarrow
import pytest

import adbc_driver_duckdb.dbapi

xfail = pytest.mark.xfail
driver_path = adbc_driver_duckdb.driver_path()


@pytest.fixture
def duck_conn():
    with adbc_driver_manager.dbapi.connect(driver=driver_path, entrypoint="duckdb_adbc_init") as conn:
        yield conn


def example_table():
    return pyarrow.table(
        [
            [1, 2, 3, 4],
            ["a", "b", None, "d"],
        ],
        names=["ints", "strs"],
    )


@xfail(sys.platform == "win32", reason="adbc-driver-manager.adbc_get_info() returns an empty dict on windows")
def test_connection_get_info(duck_conn):
    assert duck_conn.adbc_get_info() != {}


def test_connection_get_table_types(duck_conn):
    assert duck_conn.adbc_get_table_types() == []
    with duck_conn.cursor() as cursor:
        # Test Default Schema
        cursor.execute("CREATE TABLE tableschema (ints BIGINT)")
    assert duck_conn.adbc_get_table_types() == ["BASE TABLE"]


@xfail(
    sys.platform == "win32", reason="adbc-driver-manager.adbc_get_objects() returns an invalid schema dict on windows"
)
def test_connection_get_objects(duck_conn):
    with duck_conn.cursor() as cursor:
        cursor.execute("CREATE TABLE getobjects (ints BIGINT PRIMARY KEY)")
        depth_all = duck_conn.adbc_get_objects(depth="all").read_all()
    assert sorted_get_objects(depth_all.to_pylist()) is not None

    depth_tables = duck_conn.adbc_get_objects(depth="tables").read_all()
    assert sorted_get_objects(depth_tables.to_pylist()) is not None

    depth_db_schemas = duck_conn.adbc_get_objects(depth="db_schemas").read_all()
    assert sorted_get_objects(depth_db_schemas.to_pylist()) is not None

    depth_catalogs = duck_conn.adbc_get_objects(depth="catalogs").read_all()
    assert sorted_get_objects(depth_catalogs.to_pylist()) is not None

    # All result schemas should be the same
    assert depth_all.schema == depth_tables.schema
    assert depth_all.schema == depth_db_schemas.schema
    assert depth_all.schema == depth_catalogs.schema


@xfail(
    sys.platform == "win32", reason="adbc-driver-manager.adbc_get_objects() returns an invalid schema dict on windows"
)
def test_connection_get_objects_filters(duck_conn):
    with duck_conn.cursor() as cursor:
        cursor.execute("CREATE TABLE getobjects (ints BIGINT PRIMARY KEY)")

    no_filter = duck_conn.adbc_get_objects(depth="all").read_all()
    assert sorted_get_objects(no_filter.to_pylist()) is not None

    column_filter = duck_conn.adbc_get_objects(depth="all", column_name_filter="notexist").read_all()
    assert sorted_get_objects(column_filter.to_pylist()) is not None

    table_name_filter = duck_conn.adbc_get_objects(depth="all", table_name_filter="notexist").read_all()
    assert sorted_get_objects(table_name_filter.to_pylist()) is not None

    db_schema_filter = duck_conn.adbc_get_objects(depth="all", db_schema_filter="notexist").read_all()
    assert sorted_get_objects(db_schema_filter.to_pylist()) is not None

    catalog_filter = duck_conn.adbc_get_objects(depth="all", catalog_filter="notexist").read_all()
    assert catalog_filter.to_pylist() == []

    assert no_filter.schema == column_filter.schema
    assert no_filter.schema == table_name_filter.schema
    assert no_filter.schema == db_schema_filter.schema
    assert no_filter.schema == catalog_filter.schema


def test_commit(tmp_path):
    db = Path(tmp_path) / "tmp.db"
    if db.exists():
        db.unlink()
    table = example_table()
    db_kwargs = {"path": f"{db}"}
    # Start connection with auto-commit off
    with adbc_driver_manager.dbapi.connect(
        driver=driver_path,
        entrypoint="duckdb_adbc_init",
        db_kwargs=db_kwargs,
    ) as conn:
        assert not conn._autocommit
        with conn.cursor() as cur:
            cur.adbc_ingest("ingest", table, "create")

    # Check Data is not there
    with adbc_driver_manager.dbapi.connect(
        driver=driver_path,
        entrypoint="duckdb_adbc_init",
        db_kwargs=db_kwargs,
        autocommit=True,
    ) as conn:
        assert conn._autocommit
        with conn.cursor() as cur:
            # This errors because the table does not exist
            with pytest.raises(
                adbc_driver_manager._lib.InternalError,
                match=r"Table with name ingest does not exist!",
            ):
                cur.execute("SELECT count(*) from ingest")

            cur.adbc_ingest("ingest", table, "create")

    # This now works because we enabled autocommit
    with (
        adbc_driver_manager.dbapi.connect(
            driver=driver_path,
            entrypoint="duckdb_adbc_init",
            db_kwargs=db_kwargs,
        ) as conn,
        conn.cursor() as cur,
    ):
        cur.execute("SELECT count(*) from ingest")
        assert cur.fetch_arrow_table().to_pydict() == {"count_star()": [4]}


def test_connection_get_table_schema(duck_conn):
    with duck_conn.cursor() as cursor:
        # Test Default Schema
        cursor.execute("CREATE TABLE tableschema (ints BIGINT)")
        assert duck_conn.adbc_get_table_schema("tableschema") == pyarrow.schema(
            [
                ("ints", "int64"),
            ]
        )

        # Test Given Schema
        cursor.execute("CREATE SCHEMA test;")
        cursor.execute("CREATE TABLE test.tableschema (test_ints BIGINT)")
        assert duck_conn.adbc_get_table_schema("tableschema", db_schema_filter="test") == pyarrow.schema(
            [
                ("test_ints", "int64"),
            ]
        )
        assert duck_conn.adbc_get_table_schema("tableschema") == pyarrow.schema(
            [
                ("ints", "int64"),
            ]
        )

        # Test invalid catalog name
        with pytest.raises(
            adbc_driver_manager.InternalError,
            match=r'Catalog "bla" does not exist',
        ):
            duck_conn.adbc_get_table_schema("tableschema", catalog_filter="bla", db_schema_filter="test")

        # Catalog and DB Schema name
        assert duck_conn.adbc_get_table_schema(
            "tableschema", catalog_filter="memory", db_schema_filter="test"
        ) == pyarrow.schema(
            [
                ("test_ints", "int64"),
            ]
        )

        # DB Schema is inferred to be "main" if unspecified
        assert duck_conn.adbc_get_table_schema("tableschema", catalog_filter="memory") == pyarrow.schema(
            [
                ("ints", "int64"),
            ]
        )


def test_prepared_statement(duck_conn):
    with duck_conn.cursor() as cursor:
        cursor.adbc_prepare("SELECT 1")
        cursor.execute("SELECT 1")
        assert cursor.fetchone() == (1,)
        assert cursor.fetchone() is None


def test_statement_query(duck_conn):
    with duck_conn.cursor() as cursor:
        cursor.execute("SELECT 1")
        assert cursor.fetchone() == (1,)
        assert cursor.fetchone() is None

        cursor.execute("SELECT 1 AS foo")
        assert cursor.fetch_arrow_table().to_pylist() == [{"foo": 1}]


@xfail(sys.platform == "win32", reason="adbc-driver-manager returns an invalid table schema on windows")
def test_insertion(duck_conn):
    table = example_table()
    reader = table.to_reader()

    with duck_conn.cursor() as cursor:
        cursor.adbc_ingest("ingest", reader, "create")
        cursor.execute("SELECT * FROM ingest")
        assert cursor.fetch_arrow_table() == table

    with duck_conn.cursor() as cursor:
        cursor.adbc_ingest("ingest_table", table, "create")
        cursor.execute("SELECT * FROM ingest")
        assert cursor.fetch_arrow_table() == table

    # Test Append
    with duck_conn.cursor() as cursor:
        with pytest.raises(
            adbc_driver_manager.InternalError,
            match=r'Table with name "ingest_table" already exists!',
        ):
            cursor.adbc_ingest("ingest_table", table, "create")
        cursor.adbc_ingest("ingest_table", table, "append")
        cursor.execute("SELECT count(*) FROM ingest_table")
        assert cursor.fetch_arrow_table().to_pydict() == {"count_star()": [8]}


@xfail(sys.platform == "win32", reason="adbc-driver-manager returns an invalid table schema on windows")
def test_read(duck_conn):
    with duck_conn.cursor() as cursor:
        filename = Path(__file__).parent / ".." / "data" / "category.csv"
        cursor.execute(f"SELECT * FROM '{filename}'")
        assert cursor.fetch_arrow_table().to_pydict() == {
            "CATEGORY_ID": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
            "NAME": [
                "Action",
                "Animation",
                "Children",
                "Classics",
                "Comedy",
                "Documentary",
                "Drama",
                "Family",
                "Foreign",
                "Games",
                "Horror",
                "Music",
                "New",
                "Sci-Fi",
                "Sports",
                "Travel",
            ],
            "LAST_UPDATE": [
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
                datetime.datetime(2006, 2, 15, 4, 46, 27),
            ],
        }


def test_large_chunk(tmp_path):
    num_chunks = 3
    chunk_size = 10_000

    # Create data for each chunk
    chunks_col1 = [pyarrow.array(np.random.randint(0, 100, chunk_size)) for _ in range(num_chunks)]
    chunks_col2 = [pyarrow.array(np.random.rand(chunk_size)) for _ in range(num_chunks)]
    chunks_col3 = [
        pyarrow.array([f"str_{i}" for i in range(j * chunk_size, (j + 1) * chunk_size)]) for j in range(num_chunks)
    ]

    # Create chunked arrays
    col1 = pyarrow.chunked_array(chunks_col1)
    col2 = pyarrow.chunked_array(chunks_col2)
    col3 = pyarrow.chunked_array(chunks_col3)

    # Create the table
    table = pyarrow.table([col1, col2, col3], names=["ints", "floats", "strings"])

    db = Path(tmp_path) / "tmp.db"
    if db.exists():
        db.unlink()
    db_kwargs = {"path": f"{db}"}
    with (
        adbc_driver_manager.dbapi.connect(
            driver=driver_path,
            entrypoint="duckdb_adbc_init",
            db_kwargs=db_kwargs,
            autocommit=True,
        ) as conn,
        conn.cursor() as cur,
    ):
        cur.adbc_ingest("ingest", table, "create")
        cur.execute("SELECT count(*) from ingest")
        assert cur.fetch_arrow_table().to_pydict() == {"count_star()": [30_000]}


def test_dictionary_data(tmp_path):
    data = ["apple", "banana", "apple", "orange", "banana", "banana"]

    dict_type = pyarrow.dictionary(index_type=pyarrow.int32(), value_type=pyarrow.string())
    dict_array = pyarrow.array(data, type=dict_type)

    # Wrap in a table
    table = pyarrow.table({"fruits": dict_array})
    db = Path(tmp_path) / "tmp.db"
    if db.exists():
        db.unlink()
    db_kwargs = {"path": f"{db}"}
    with (
        adbc_driver_manager.dbapi.connect(
            driver=driver_path,
            entrypoint="duckdb_adbc_init",
            db_kwargs=db_kwargs,
            autocommit=True,
        ) as conn,
        conn.cursor() as cur,
    ):
        cur.adbc_ingest("ingest", table, "create")
        cur.execute("from ingest")
        assert cur.fetch_arrow_table().to_pydict() == {
            "fruits": ["apple", "banana", "apple", "orange", "banana", "banana"]
        }


def test_ree_data(tmp_path):
    run_ends = pyarrow.array([3, 5, 6], type=pyarrow.int32())  # positions: [0-2], [3-4], [5]
    values = pyarrow.array(["apple", "banana", "orange"], type=pyarrow.string())

    ree_array = pyarrow.RunEndEncodedArray.from_arrays(run_ends, values)

    table = pyarrow.table({"fruits": ree_array})

    db = Path(tmp_path) / "tmp.db"
    if db.exists():
        db.unlink()
    db_kwargs = {"path": f"{db}"}
    with (
        adbc_driver_manager.dbapi.connect(
            driver=driver_path,
            entrypoint="duckdb_adbc_init",
            db_kwargs=db_kwargs,
            autocommit=True,
        ) as conn,
        conn.cursor() as cur,
    ):
        cur.adbc_ingest("ingest", table, "create")
        cur.execute("from ingest")
        assert cur.fetch_arrow_table().to_pydict() == {
            "fruits": ["apple", "apple", "apple", "banana", "banana", "orange"]
        }


def sorted_get_objects(catalogs):
    res = []
    for catalog in sorted(catalogs, key=lambda cat: cat["catalog_name"]):
        new_catalog = {
            "catalog_name": catalog["catalog_name"],
            "catalog_db_schemas": [],
        }

        for db_schema in sorted(catalog["catalog_db_schemas"] or [], key=lambda sch: sch["db_schema_name"]):
            new_db_schema = {
                "db_schema_name": db_schema["db_schema_name"],
                "db_schema_tables": [],
            }

            for table in sorted(db_schema["db_schema_tables"] or [], key=lambda tab: tab["table_name"]):
                new_table = {
                    "table_name": table["table_name"],
                    "table_type": table["table_type"],
                    "table_columns": [],
                    "table_constraints": [],
                }

                for column in sorted(table["table_columns"] or [], key=lambda col: col["ordinal_position"]):
                    new_table["table_columns"].append(column)

                for constraint in sorted(table["table_constraints"] or [], key=lambda con: con["constraint_name"]):
                    new_table["table_constraints"].append(constraint)

                new_db_schema["db_schema_tables"].append(new_table)
            new_catalog["catalog_db_schemas"].append(new_db_schema)
        res.append(new_catalog)

    return res