duckdb-sqlalchemy 0.19.0__py3-none-any.whl

duckdb_sqlalchemy/tests/test_datatypes.py

@@ -0,0 +1,251 @@
+import decimal
+import json
+import warnings
+from typing import Any, Callable, Dict, Type, cast
+from uuid import uuid4
+
+from packaging.version import Version
+from pytest import importorskip, mark
+from pytest_snapshot.plugin import Snapshot
+from sqlalchemy import (
+    Column,
+    Integer,
+    Interval,
+    MetaData,
+    Sequence,
+    String,
+    Table,
+    inspect,
+    schema,
+    select,
+    text,
+)
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.engine import Engine, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import Session
+from sqlalchemy.sql import sqltypes
+from sqlalchemy.types import FLOAT, JSON
+
+from .._supports import duckdb_version, has_uhugeint_support
+from ..datatypes import Map, Struct, types
+
+
+@mark.parametrize("coltype", types)
+@mark.skipif(not has_uhugeint_support, reason="duckdb version too old")
+def test_unsigned_integer_type(
+    engine: Engine, session: Session, coltype: Type[Integer]
+) -> None:
+    Base = declarative_base()
+
+    tname = "table"
+    table = type(
+        "Table",
+        (Base,),
+        {
+            "__tablename__": tname,
+            "id": Column(Integer, primary_key=True, default=0),
+            "a": Column(coltype),
+        },
+    )
+    Base.metadata.create_all(engine)
+
+    has_table = cast(
+        Callable[[str], bool],
+        engine.has_table if hasattr(engine, "has_table") else inspect(engine).has_table,
+    )
+
+    assert has_table(tname)
+
+    session.add(table(a=1))
+    session.commit()
+
+    assert session.query(table).one()
+
+
+@mark.remote_data()
+def test_raw_json(engine: Engine) -> None:
+    importorskip("duckdb", "0.9.3.dev4040")
+
+    with engine.connect() as conn:
+        assert conn.execute(text("load json"))
+
+        assert conn.execute(text("select {'Hello': 'world'}::JSON")).fetchone() == (
+            {"Hello": "world"},
+        )
+
+
+@mark.remote_data()
+def test_custom_json_serializer() -> None:
+    def default(o: Any) -> Any:
+        if isinstance(o, decimal.Decimal):
+            return {"__tag": "decimal", "value": str(o)}
+
+    def object_hook(pairs: Dict[str, Any]) -> Any:
+        if pairs.get("__tag", None) == "decimal":
+            return decimal.Decimal(pairs["value"])
+        else:
+            return pairs
+
+    engine = create_engine(
+        "duckdb://",
+        json_serializer=json.JSONEncoder(default=default).encode,
+        json_deserializer=json.JSONDecoder(object_hook=object_hook).decode,
+    )
+
+    Base = declarative_base()
+
+    class Entry(Base):
+        __tablename__ = "test_json"
+        id = Column(Integer, Sequence("id_seq"), primary_key=True)
+        data = Column(JSON, nullable=False)
+
+    Base.metadata.create_all(engine)
+
+    with engine.connect() as conn:
+        session = Session(bind=conn)
+
+        data = {"hello": decimal.Decimal("42")}
+
+        session.add(Entry(data=data))  # type: ignore[call-arg]
+        session.commit()
+
+        (res,) = session.execute(select(Entry)).one()
+
+        assert res.data == data
+
+
+def test_json(engine: Engine, session: Session) -> None:
+    base = declarative_base()
+
+    class Entry(base):
+        __tablename__ = "test_json"
+
+        id = Column(Integer, primary_key=True, default=0)
+        meta = Column(JSON, nullable=False)
+
+    base.metadata.create_all(bind=engine)
+
+    session.add(Entry(meta={"hello": "world"}))  # type: ignore[call-arg]
+    session.commit()
+
+    result = session.query(Entry).one()
+
+    assert result.meta == {"hello": "world"}
+
+
+def test_uuid(engine: Engine, session: Session) -> None:
+    importorskip("duckdb", "0.7.1")
+    base = declarative_base()
+
+    class Entry(base):
+        __tablename__ = "test_uuid"
+
+        id = Column(UUID, primary_key=True, default=0)
+
+    base.metadata.create_all(bind=engine)
+
+    ident = uuid4()
+
+    session.add(Entry(id=ident))  # type: ignore[call-arg]
+    session.commit()
+
+    result = session.query(Entry).one()
+
+    assert result.id == ident
+
+
+def test_double_in_sqla_v2(engine: Engine) -> None:
+    with engine.begin() as con:
+        con.execute(text("CREATE TABLE t (x DOUBLE)"))
+        con.execute(text("INSERT INTO t VALUES (1.0), (2.0), (3.0)"))
+
+    md = MetaData()
+
+    t = Table("t", md, autoload_with=engine)
+
+    with engine.begin() as con:
+        con.execute(t.select())
+
+
+def test_all_types_reflection(engine: Engine) -> None:
+    importorskip("sqlalchemy", "1.4.0")
+    importorskip("duckdb", "0.5.1")
+
+    with warnings.catch_warnings() as capture, engine.connect() as conn:
+        conn.execute(text("create table t2 as select * from test_all_types()"))
+        table = Table("t2", MetaData(), autoload_with=conn)
+        for col in table.columns:
+            name = col.name
+            if name.endswith("_enum") and duckdb_version < Version("0.7.1"):
+                continue
+            if "array" in name or "struct" in name or "map" in name or "union" in name:
+                assert col.type == sqltypes.NULLTYPE, name
+            else:
+                assert col.type != sqltypes.NULLTYPE, name
+        assert not capture
+
+
+def test_nested_types(engine: Engine, session: Session) -> None:
+    importorskip("duckdb", "0.5.0")  # nested types require at least duckdb 0.5.0
+    base = declarative_base()
+
+    class Entry(base):
+        __tablename__ = "test_struct"
+
+        id = Column(Integer, primary_key=True, default=0)
+        struct = Column(Struct(fields={"name": String}))
+        map = Column(Map(String, Integer))
+        # union = Column(Union(fields={"name": String, "age": Integer}))
+
+    base.metadata.create_all(bind=engine)
+
+    struct_data = {"name": "Edgar"}
+    map_data = {"one": 1, "two": 2}
+
+    session.add(Entry(struct=struct_data, map=map_data))  # type: ignore[call-arg]
+    session.commit()
+
+    result = session.query(Entry).one()
+
+    assert result.struct == struct_data
+    assert result.map == map_data
+
+
+def test_double_nested_types(engine: Engine, session: Session) -> None:
+    """Test for https://github.com/leonardovida/duckdb-sqlalchemy/issues/1138"""
+    importorskip("duckdb", "0.5.0")  # nested types require at least duckdb 0.5.0
+    base = declarative_base()
+
+    class Entry(base):
+        __tablename__ = "test_struct"
+
+        id = Column(Integer, primary_key=True, default=0)
+        outer = Column(Struct({"inner": Struct({"val": Integer})}))
+
+    base.metadata.create_all(bind=engine)
+
+    outer = {"inner": {"val": 42}}
+
+    session.add(Entry(outer=outer))  # type: ignore[call-arg]
+    session.commit()
+
+    result = session.query(Entry).one()
+
+    assert result.outer == outer
+
+
+def test_interval(engine: Engine, snapshot: Snapshot) -> None:
+    test_table = Table("test_table", MetaData(), Column("duration", Interval))
+    create_sql = f"{str(schema.CreateTable(test_table).compile(engine)).strip()}\n"
+    snapshot.assert_match(create_sql, "schema.sql")
+
+
+def test_div_is_floordiv(engine: Engine) -> None:
+    test_table = Table(
+        "test_table", MetaData(), Column("value", FLOAT), Column("eur2usd_rate", FLOAT)
+    )
+
+    stmt = test_table.c.value / test_table.c.eur2usd_rate
+
+    assert str(stmt.compile(engine)) == "test_table.value / test_table.eur2usd_rate"

duckdb_sqlalchemy/tests/test_execution_options.py

@@ -0,0 +1,32 @@
+import importlib.util
+
+import pytest
+from sqlalchemy import Column, Integer, MetaData, String, Table, create_engine, select
+
+
+def test_bulk_insert_register_path() -> None:
+    if (
+        importlib.util.find_spec("pandas") is None
+        and importlib.util.find_spec("pyarrow") is None
+    ):
+        pytest.skip("pandas or pyarrow is required for bulk insert fast path")
+
+    engine = create_engine("duckdb:///:memory:")
+    md = MetaData()
+    t = Table(
+        "bulk_insert",
+        md,
+        Column("id", Integer),
+        Column("name", String),
+    )
+    md.create_all(engine)
+
+    rows = [{"id": 1, "name": "Ada"}, {"id": 2, "name": "Grace"}]
+
+    with engine.begin() as conn:
+        conn = conn.execution_options(duckdb_copy_threshold=1)
+        conn.execute(t.insert(), rows)
+
+    with engine.connect() as conn:
+        result = conn.execute(select(t.c.id, t.c.name).order_by(t.c.id)).fetchall()
+        assert result == [(1, "Ada"), (2, "Grace")]

duckdb_sqlalchemy/tests/test_helpers.py

@@ -0,0 +1,71 @@
+from urllib.parse import parse_qs
+
+import sqlalchemy
+from packaging.version import Version
+from pytest import raises
+from sqlalchemy import select
+
+from duckdb_sqlalchemy import URL, Dialect, make_url, read_parquet
+from duckdb_sqlalchemy.config import get_core_config
+
+
+def test_url_helper_round_trip() -> None:
+    url = URL(database=":memory:", read_only=True)
+    assert url.drivername == "duckdb"
+    assert url.database == ":memory:"
+    assert str(url.query["read_only"]).lower() == "true"
+
+    base = make_url(database=":memory:")
+    rendered = base.render_as_string(hide_password=False)
+    assert rendered.startswith("duckdb:///")
+
+
+def test_read_parquet_helper() -> None:
+    if Version(sqlalchemy.__version__) < Version("1.4.0"):
+        with raises(NotImplementedError):
+            read_parquet("data/events.parquet", columns=["event_id"])
+        return
+
+    parquet = read_parquet("data/events.parquet", columns=["event_id"])
+    stmt = select(parquet.c.event_id).select_from(parquet)
+    sql = str(stmt.compile(dialect=Dialect(), compile_kwargs={"literal_binds": True}))
+    assert "read_parquet" in sql
+
+
+def test_motherduck_config_env_and_ttl(monkeypatch) -> None:
+    get_core_config()
+
+    captured = {}
+
+    class DummyConn:
+        def execute(self, *args, **kwargs):
+            return self
+
+        def register_filesystem(self, filesystem):
+            return None
+
+        def close(self):
+            return None
+
+    def fake_connect(*cargs, **cparams):
+        captured.update(cparams)
+        return DummyConn()
+
+    import duckdb_sqlalchemy
+
+    monkeypatch.setenv("motherduck_token", "token123")
+    monkeypatch.delenv("MOTHERDUCK_TOKEN", raising=False)
+    monkeypatch.setattr(duckdb_sqlalchemy.duckdb, "connect", fake_connect)
+
+    dialect = Dialect()
+    dialect.connect(
+        database="md:my_db",
+        url_config={"dbinstance_inactivity_ttl": "1h"},
+        config={},
+    )
+
+    assert captured["config"]["motherduck_token"] == "token123"
+    database, query = captured["database"].split("?", 1)
+    assert database == "md:my_db"
+    assert parse_qs(query)["dbinstance_inactivity_ttl"] == ["1h"]
+    assert "motherduck_dbinstance_inactivity_ttl" not in captured["config"]

duckdb_sqlalchemy/tests/test_integration.py

@@ -0,0 +1,46 @@
+import duckdb
+import pandas as pd
+from pytest import mark, raises
+from sqlalchemy import __version__, text
+from sqlalchemy.engine import create_engine
+from sqlalchemy.engine.base import Connection
+from sqlalchemy.exc import ProgrammingError
+
+df = pd.DataFrame([{"a": 1}])
+
+
+@mark.skipif(not hasattr(Connection, "exec_driver_sql"), reason="Needs exec_driver_sql")
+def test_register_driver(conn: Connection) -> None:
+    conn.exec_driver_sql("register", ("test_df_driver", df))  # type: ignore[arg-type]
+    conn.execute(text("select * from test_df_driver"))
+
+
+def test_plain_register(conn: Connection) -> None:
+    if __version__.startswith("1.3"):
+        conn.execute(text("register"), {"name": "test_df", "df": df})
+    else:
+        conn.execute(text("register(:name, :df)"), {"name": "test_df", "df": df})
+    conn.execute(text("select * from test_df"))
+
+
+duckdb_version = duckdb.__version__
+
+
+@mark.remote_data
+@mark.skipif(
+    "dev" in duckdb_version, reason="md extension not available for dev builds"
+)
+@mark.skipif(
+    duckdb_version != "1.1.1", reason="md extension not available for this version"
+)
+def test_motherduck() -> None:
+    engine = create_engine(
+        "duckdb:///md:motherdb",
+        connect_args={"config": {"motherduck_token": "motherduckdb_token"}},
+    )
+
+    with raises(
+        ProgrammingError,
+        match="Jwt is not in the form of Header.Payload.Signature with two dots and 3 sections",
+    ):
+        engine.connect()

duckdb_sqlalchemy/tests/test_pandas.py

@@ -0,0 +1,146 @@
+#! /usr/bin/env python3
+# vim:fenc=utf-8
+
+"""
+Test pandas functionality.
+"""
+
+import pathlib
+import random
+from collections import OrderedDict
+from datetime import datetime, timezone
+from itertools import product
+from typing import Dict, List, Optional, Tuple, Union, cast
+
+import pandas as pd
+from pandas.testing import assert_frame_equal
+from pytest import importorskip, mark
+from sqlalchemy import create_engine
+
+from .util import sqlalchemy_1_only
+
+pytestmark = sqlalchemy_1_only
+
+_possible_args = OrderedDict(
+    {
+        "chunksize": [None, 1, 10, 100],
+        "if_exists": ["fail", "replace", "append"],
+        "method": [
+            None,
+            "multi",
+        ],  # TODO Implement a callable insert method?
+    }
+)
+
+args = {
+    "to_sql": ["chunksize", "if_exists", "method"],
+    "read_sql": ["chunksize"],
+}
+params = {
+    k: list(
+        product(*(_v for _k, _v in _possible_args.items() if _k in args[k]))  # type: ignore
+    )
+    for k in args
+}
+params_strings = {k: (",".join([str(_k) for _k in args[k]])) for k in args}
+
+# Generate a DataFrame of 100 rows.
+sample_data: Dict[str, List[Union[datetime, str, int, float]]] = {
+    "datetime": [],
+    "int": [],
+    "str": [],
+    "float": [],
+}
+sample_rowcount = max(
+    cs
+    for cs in cast(List[Optional[int]], _possible_args["chunksize"])
+    if cs is not None
+)
+for i in range(sample_rowcount):
+    sample_data["datetime"].append(datetime.now(timezone.utc))
+    sample_data["int"].append(random.randint(0, 100))
+    sample_data["float"].append(round(random.random(), 5))
+    sample_data["str"].append("foo")
+
+sample_df: pd.DataFrame = pd.DataFrame(sample_data)
+
+
+@mark.parametrize(params_strings["to_sql"], params["to_sql"])
+def test_to_sql(
+    chunksize: Optional[int],
+    if_exists: str,
+    method: Optional[str],
+    index: bool = False,
+) -> None:
+    eng = create_engine("duckdb:///:memory:")
+    try:
+        sample_df.to_sql(
+            name="foo",
+            con=eng,
+            if_exists=if_exists,
+            chunksize=chunksize,
+            index=index,
+        )
+    except ValueError as e:
+        if if_exists != "fail":
+            raise e
+
+
+table_name = "test_read"
+
+
+# Perform the test twice:
+# Once for reading the table name (testing reflection),
+@mark.parametrize(params_strings["read_sql"], params["read_sql"])
+def test_read_sql_reflection(
+    chunksize: Tuple[Optional[int]],
+) -> None:
+    run_query(table_name, chunksize[0])
+
+
+# and once for directly executing a SQL query.
+@mark.parametrize(params_strings["read_sql"], params["read_sql"])
+def test_read_sql(
+    chunksize: Tuple[Optional[int]],
+) -> None:
+    run_query(f"SELECT * FROM {table_name}", chunksize[0])
+
+
+def run_query(query: str, chunksize: Optional[int]) -> None:
+    eng = create_engine("duckdb:///:memory:")
+
+    sample_df.to_sql(name=table_name, con=eng, if_exists="replace")
+
+    result = pd.read_sql(query, eng, chunksize=chunksize)
+    chunks = [result] if chunksize is None else list(result)
+    if chunksize is None:
+        df = cast(pd.DataFrame, chunks[0])
+        assert len(df) == sample_rowcount
+    else:
+        importorskip("duckdb", "0.5.0")
+
+        # Assert that the chunks are the size specified.
+        df = cast(pd.DataFrame, chunks[0])
+        assert len(df) == chunksize
+        # Assert that the expected number of chunks was returned.
+        assert (sample_rowcount / chunksize) == len(chunks)
+
+
+def test_read_sql_duckdb_table(tmp_path: pathlib.Path) -> None:
+    import duckdb
+
+    db = str(tmp_path / "test_db.duckdb")
+    con = duckdb.connect(database=db, read_only=False)
+    df = pd.DataFrame(
+        {"a": [1, 2, 3, 4], "b": [0.1, 0.2, 0.3, 0.4], "c": ["a", "b", "c", "d"]}
+    )
+    con.register("df_view", df)
+    con.execute("CREATE TABLE test_data AS SELECT * FROM df_view;")
+    con.close()
+
+    engine = create_engine(f"duckdb:///{db}")
+
+    result = pd.read_sql("SELECT * FROM test_data", engine)
+    assert_frame_equal(result, df)
+    result = pd.read_sql("test_data", engine)
+    assert_frame_equal(result, df)

duckdb_sqlalchemy/tests/test_pyarrow.py

@@ -0,0 +1,61 @@
+from pyarrow import RecordBatch, RecordBatchReader
+from pyarrow import Table as ArrowTable
+from sqlalchemy import MetaData, Table, create_engine, text
+
+
+def test_fetch_arrow() -> None:
+    engine = create_engine("duckdb:///:memory:")
+    with engine.begin() as con:
+        con.execute(text("CREATE TABLE tbl (label VARCHAR, value DOUBLE)"))
+        con.execute(
+            text(
+                "INSERT INTO tbl VALUES ('xx',-1.0), ('ww',-4.5), ('zz',6.0), ('yy',2.5)"
+            )
+        )
+
+    md = MetaData()
+    t = Table("tbl", md, autoload_with=engine)
+    stmt = t.select().where(t.c.value > -4.0).order_by(t.c.label)
+
+    # rows
+    with engine.begin() as con:
+        res = con.execute(stmt).cursor.fetchall()
+        assert res == [("xx", -1.0), ("yy", 2.5), ("zz", 6.0)]
+
+    # arrow table
+    with engine.begin() as con:
+        res = con.execute(stmt).cursor.fetch_arrow_table()
+        assert isinstance(res, ArrowTable)
+        assert res == ArrowTable.from_pydict(
+            {"label": ["xx", "yy", "zz"], "value": [-1.0, 2.5, 6.0]}
+        )
+
+    # arrow batches
+    with engine.begin() as con:
+        res = con.execute(stmt).cursor.fetch_record_batch()
+        assert isinstance(res, RecordBatchReader)
+        assert res.read_all() == ArrowTable.from_pydict(
+            {"label": ["xx", "yy", "zz"], "value": [-1.0, 2.5, 6.0]}
+        )
+        res = con.execute(stmt).cursor.fetch_record_batch(rows_per_batch=2)
+        assert res.read_next_batch() == RecordBatch.from_pydict(
+            {"label": ["xx", "yy"], "value": [-1.0, 2.5]}
+        )
+        assert res.read_next_batch() == RecordBatch.from_pydict(
+            {"label": ["zz"], "value": [6.0]}
+        )
+
+
+def test_arrow_execution_option() -> None:
+    engine = create_engine("duckdb:///:memory:")
+    with engine.begin() as con:
+        con.execute(text("CREATE TABLE tbl (label VARCHAR, value DOUBLE)"))
+        con.execute(text("INSERT INTO tbl VALUES ('aa', 1.0), ('bb', 2.0)"))
+
+    with engine.connect().execution_options(duckdb_arrow=True) as con:
+        result = con.execute(text("SELECT * FROM tbl ORDER BY label"))
+        table = result.arrow
+        assert isinstance(table, ArrowTable)
+        assert table == ArrowTable.from_pydict(
+            {"label": ["aa", "bb"], "value": [1.0, 2.0]}
+        )

duckdb_sqlalchemy/tests/util.py

@@ -0,0 +1,9 @@
+import sqlalchemy
+from packaging.version import Version
+from pytest import mark
+
+is_sqlalchemy_1 = Version(sqlalchemy.__version__).major == 1
+sqlalchemy_1_only = mark.skipif(
+    not is_sqlalchemy_1,
+    reason="Pandas doesn't yet support sqlalchemy 2+",
+)

duckdb_sqlalchemy/url.py

@@ -0,0 +1,59 @@
+from typing import Any, Mapping, Optional, Sequence
+
+from sqlalchemy.engine import URL as SAURL
+
+__all__ = ["URL", "make_url"]
+
+
+def URL(
+    *,
+    database: str = ":memory:",
+    query: Optional[Mapping[str, Any]] = None,
+    **kwargs: Any,
+) -> SAURL:
+    """
+    Build a SQLAlchemy URL for duckdb-sqlalchemy.
+
+    All keyword arguments are treated as DuckDB config options (URL query params).
+    """
+
+    def _stringify(value: Any) -> str:
+        if isinstance(value, bool):
+            return "true" if value else "false"
+        return str(value)
+
+    def _coerce(value: Any) -> Any:
+        if value is None:
+            return None
+        if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
+            return tuple(_stringify(v) for v in value)
+        return _stringify(value)
+
+    query_dict: dict[str, str | tuple[str, ...]] = {}
+    if query:
+        query_dict.update(
+            {
+                k: v
+                for k, v in ((k, _coerce(v)) for k, v in query.items())
+                if v is not None
+            }
+        )
+    if kwargs:
+        query_dict.update(
+            {
+                k: v
+                for k, v in ((k, _coerce(v)) for k, v in kwargs.items())
+                if v is not None
+            }
+        )
+
+    return SAURL.create("duckdb", database=database, query=query_dict)
+
+
+def make_url(
+    *,
+    database: str = ":memory:",
+    query: Optional[Mapping[str, Any]] = None,
+    **kwargs: Any,
+) -> SAURL:
+    return URL(database=database, query=query, **kwargs)