duckdb-sqlalchemy 0.19.0__py3-none-any.whl

duckdb_sqlalchemy/_supports.py

@@ -0,0 +1,23 @@
+from functools import lru_cache
+
+import duckdb
+
+from .capabilities import get_capabilities
+
+_capabilities = get_capabilities(duckdb.__version__)
+duckdb_version = _capabilities.version
+has_uhugeint_support = _capabilities.supports_uhugeint
+
+
+@lru_cache()
+def has_comment_support() -> bool:
+    """
+    See https://github.com/duckdb/duckdb/pull/10372
+    """
+    try:
+        with duckdb.connect(":memory:") as con:
+            con.execute("CREATE TABLE t (i INTEGER);")
+            con.execute("COMMENT ON TABLE t IS 'test';")
+    except duckdb.ParserException:
+        return False
+    return True

duckdb_sqlalchemy/bulk.py

@@ -0,0 +1,192 @@
+import csv
+import tempfile
+from pathlib import Path
+from typing import Any, Iterable, Mapping, Optional, Sequence, Tuple, Union
+
+TableLike = Union[str, Any]
+
+
+def _quote_literal(value: Any) -> str:
+    if isinstance(value, Path):
+        value = str(value)
+    if isinstance(value, bool):
+        return "TRUE" if value else "FALSE"
+    if isinstance(value, (int, float)):
+        return str(value)
+    if value is None:
+        return "NULL"
+    return "'" + str(value).replace("'", "''") + "'"
+
+
+def _format_copy_options(options: Mapping[str, Any]) -> str:
+    if not options:
+        return ""
+    parts = []
+    for key, value in options.items():
+        if value is None:
+            continue
+        opt_key = str(key).upper()
+        if isinstance(value, (list, tuple)):
+            inner = ", ".join(_quote_literal(v) for v in value)
+            parts.append(f"{opt_key} ({inner})")
+        else:
+            parts.append(f"{opt_key} {_quote_literal(value)}")
+    if not parts:
+        return ""
+    return " (" + ", ".join(parts) + ")"
+
+
+def _format_table(connection: Any, table: TableLike) -> str:
+    if hasattr(table, "name"):
+        preparer = getattr(
+            getattr(connection, "dialect", None), "identifier_preparer", None
+        )
+        if preparer is not None:
+            return preparer.format_table(table)
+        schema = getattr(table, "schema", None)
+        name = getattr(table, "name", None)
+        if schema:
+            return f"{schema}.{name}"
+        return str(name)
+    return str(table)
+
+
+def _format_columns(connection: Any, columns: Optional[Sequence[str]]) -> str:
+    if not columns:
+        return ""
+    preparer = getattr(
+        getattr(connection, "dialect", None), "identifier_preparer", None
+    )
+    if preparer is None:
+        cols = ", ".join(columns)
+    else:
+        cols = ", ".join(preparer.quote_identifier(col) for col in columns)
+    return f" ({cols})"
+
+
+def _execute_sql(connection: Any, statement: str) -> Any:
+    if hasattr(connection, "exec_driver_sql"):
+        return connection.exec_driver_sql(statement)
+    return connection.execute(statement)
+
+
+def copy_from_parquet(
+    connection: Any,
+    table: TableLike,
+    path: Union[str, Path],
+    *,
+    columns: Optional[Sequence[str]] = None,
+    **options: Any,
+) -> Any:
+    return _copy_from_file(
+        connection,
+        table,
+        path,
+        format_name="parquet",
+        columns=columns,
+        **options,
+    )
+
+
+def copy_from_csv(
+    connection: Any,
+    table: TableLike,
+    path: Union[str, Path],
+    *,
+    columns: Optional[Sequence[str]] = None,
+    **options: Any,
+) -> Any:
+    return _copy_from_file(
+        connection,
+        table,
+        path,
+        format_name="csv",
+        columns=columns,
+        **options,
+    )
+
+
+def _copy_from_file(
+    connection: Any,
+    table: TableLike,
+    path: Union[str, Path],
+    *,
+    format_name: str,
+    columns: Optional[Sequence[str]] = None,
+    **options: Any,
+) -> Any:
+    table_name = _format_table(connection, table)
+    column_clause = _format_columns(connection, columns)
+    path_literal = _quote_literal(path)
+    copy_options = {"format": format_name, **options}
+    options_clause = _format_copy_options(copy_options)
+    statement = f"COPY {table_name}{column_clause} FROM {path_literal}{options_clause}"
+    return _execute_sql(connection, statement)
+
+
+def copy_from_rows(
+    connection: Any,
+    table: TableLike,
+    rows: Iterable[Union[Mapping[str, Any], Sequence[Any]]],
+    *,
+    columns: Optional[Sequence[str]] = None,
+    chunk_size: int = 100000,
+    include_header: bool = False,
+    **copy_options: Any,
+) -> Any:
+    iterator = iter(rows)
+    first = next(iterator, None)
+    if first is None:
+        return None
+
+    if isinstance(first, Mapping):
+        if columns is None:
+            columns = list(first.keys())
+
+        def row_to_seq(row: Mapping[str, Any]) -> Sequence[Any]:
+            return [row.get(col) for col in columns or []]
+    else:
+
+        def row_to_seq(row: Sequence[Any]) -> Sequence[Any]:
+            return row
+
+    def open_writer() -> Tuple[Any, Any, int]:
+        tmp = tempfile.NamedTemporaryFile("w", newline="", suffix=".csv", delete=False)
+        writer = csv.writer(tmp)
+        if include_header and columns:
+            writer.writerow(columns)
+        return tmp, writer, 0
+
+    def flush_chunk(tmp: Any) -> None:
+        tmp.flush()
+        path = tmp.name
+        tmp.close()
+        try:
+            copy_from_csv(
+                connection,
+                table,
+                path,
+                columns=columns if columns else None,
+                **copy_options,
+            )
+        finally:
+            try:
+                Path(path).unlink()
+            except FileNotFoundError:
+                pass
+
+    copy_options = {"header": include_header, **copy_options}
+
+    tmp, writer, count = open_writer()
+    writer.writerow(row_to_seq(first))
+    count += 1
+
+    for row in iterator:
+        writer.writerow(row_to_seq(row))
+        count += 1
+        if chunk_size and count >= chunk_size:
+            flush_chunk(tmp)
+            tmp, writer, count = open_writer()
+
+    if count:
+        flush_chunk(tmp)

duckdb_sqlalchemy/capabilities.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from typing import Union
+
+from packaging.version import Version
+
+
+@dataclass(frozen=True)
+class DuckDBCapabilities:
+    version: Version
+    supports_attach: bool
+    supports_user_agent: bool
+    supports_uhugeint: bool
+    supports_varint: bool
+
+
+def get_capabilities(version: Union[str, Version]) -> DuckDBCapabilities:
+    resolved = version if isinstance(version, Version) else Version(version)
+    return DuckDBCapabilities(
+        version=resolved,
+        supports_attach=resolved >= Version("0.7.0"),
+        supports_user_agent=resolved >= Version("0.9.2"),
+        supports_uhugeint=resolved >= Version("0.10.0"),
+        supports_varint=resolved > Version("1.0.0"),
+    )

duckdb_sqlalchemy/config.py

@@ -0,0 +1,70 @@
+import os
+from decimal import Decimal
+from functools import lru_cache
+from typing import Dict, Set, Type, Union
+
+import duckdb
+from sqlalchemy import Boolean, Float, Integer, String
+from sqlalchemy.engine import Dialect
+from sqlalchemy.sql.type_api import TypeEngine
+
+TYPES: Dict[Type, TypeEngine] = {
+    bool: Boolean(),
+    int: Integer(),
+    float: Float(),
+    str: String(),
+}
+
+
+@lru_cache()
+def get_core_config() -> Set[str]:
+    # List of connection string parameters that are supported by MotherDuck
+    # See: https://motherduck.com/docs/key-tasks/authenticating-and-connecting-to-motherduck/authenticating-to-motherduck/
+    motherduck_config_keys = {
+        "motherduck_token",
+        "attach_mode",
+        "saas_mode",
+        "session_hint",
+        "access_mode",
+        "dbinstance_inactivity_ttl",
+        "motherduck_dbinstance_inactivity_ttl",
+    }
+
+    with duckdb.connect(":memory:") as conn:
+        rows = conn.execute("SELECT name FROM duckdb_settings()").fetchall()
+    return {name for (name,) in rows} | motherduck_config_keys
+
+
+def apply_config(
+    dialect: Dialect,
+    conn: duckdb.DuckDBPyConnection,
+    ext: Dict[str, Union[str, int, bool, float, None]],
+) -> None:
+    # TODO: does sqlalchemy have something that could do this for us?
+    processors = [
+        (typ, type_engine.literal_processor(dialect=dialect))
+        for typ, type_engine in TYPES.items()
+    ]
+    string_processor = String().literal_processor(dialect=dialect)
+
+    for k, v in ext.items():
+        if v is None:
+            conn.execute(f"SET {k} = NULL")
+            continue
+        if isinstance(v, os.PathLike):
+            v = os.fspath(v)
+            process = string_processor
+        elif isinstance(v, Decimal):
+            v = str(v)
+            process = string_processor
+        else:
+            process = None
+            for typ, processor in processors:
+                if isinstance(v, typ):
+                    process = processor
+                    break
+            if process is None:
+                v = str(v)
+                process = string_processor
+        assert process, f"Not able to configure {k} with {v}"
+        conn.execute(f"SET {k} = {process(v)}")

duckdb_sqlalchemy/datatypes.py

@@ -0,0 +1,332 @@
+"""
+See https://duckdb.org/docs/sql/data_types/numeric for more information
+
+Also
+```sql
+select * from duckdb_types where type_category = 'NUMERIC';
+```
+"""
+
+import typing
+from typing import Any, Dict, Optional, Tuple, Type
+
+import duckdb
+from packaging.version import Version
+from sqlalchemy import exc, util
+from sqlalchemy.dialects.postgresql.base import PGIdentifierPreparer, PGTypeCompiler
+from sqlalchemy.engine import Dialect
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql import sqltypes, type_api
+from sqlalchemy.sql.type_api import TypeEngine
+from sqlalchemy.types import BigInteger, Integer, SmallInteger, String
+
+# INTEGER	INT4, INT, SIGNED	-2147483648	2147483647
+# SMALLINT	INT2, SHORT	-32768	32767
+# BIGINT	INT8, LONG	-9223372036854775808	9223372036854775807
+(BigInteger, SmallInteger)  # pure reexport
+
+duckdb_version = duckdb.__version__
+
+IS_GT_1 = Version(duckdb_version) > Version("1.0.0")
+
+
+class DuckDBInteger(Integer):
+    cache_ok = True
+
+
+class UInt64(DuckDBInteger):
+    pass
+
+
+class UInt32(DuckDBInteger):
+    pass
+
+
+class UInt16(DuckDBInteger):
+    "AKA USMALLINT"
+
+
+class UInt8(DuckDBInteger):
+    pass
+
+
+class UTinyInteger(DuckDBInteger):
+    "AKA UInt1"
+
+    name = "UTinyInt"
+    # UTINYINT	-	0	255
+
+
+class TinyInteger(DuckDBInteger):
+    "AKA Int1"
+
+    name = "TinyInt"
+    # TINYINT	INT1	-128	127
+
+
+class USmallInteger(DuckDBInteger):
+    name = "usmallint"
+    "AKA UInt2"
+    # USMALLINT	-	0	65535
+    min = 0
+    max = 2**15
+
+
+class UBigInteger(DuckDBInteger):
+    name = "UBigInt"
+    min = 0
+    max = 18446744073709551615
+
+
+class HugeInteger(DuckDBInteger):
+    name = "HugeInt"
+    # HUGEINT	 	-170141183460469231731687303715884105727*	170141183460469231731687303715884105727
+
+
+class UHugeInteger(DuckDBInteger):
+    name = "UHugeInt"
+
+
+class UInteger(DuckDBInteger):
+    # UINTEGER	-	0	4294967295
+    pass
+
+
+if IS_GT_1:
+
+    class VarInt(DuckDBInteger):
+        pass
+
+
+def compile_uint(element: Integer, compiler: PGTypeCompiler, **kw: Any) -> str:
+    return getattr(element, "name", type(element).__name__)
+
+
+types = [
+    subclass
+    for subclass in DuckDBInteger.__subclasses__()
+    if subclass.__module__ == UInt64.__module__
+]
+assert types
+
+
+TV = typing.Union[Type[TypeEngine], TypeEngine]
+
+
+def _normalize_fields(
+    fields: Optional[Dict[str, TV]],
+) -> Optional[Tuple[Tuple[str, TypeEngine], ...]]:
+    if fields is None:
+        return None
+    if isinstance(fields, dict):
+        items = fields.items()
+    else:
+        items = fields
+    return tuple(
+        (str(key), type_api.to_instance(value))
+        for key, value in items  # type: ignore[arg-type]
+    )
+
+
+def _normalize_fields_cache_key(
+    fields: Optional[Tuple[Tuple[str, TypeEngine], ...]],
+) -> Optional[Tuple[Tuple[str, Any], ...]]:
+    if fields is None:
+        return None
+    return tuple(
+        (key, value._static_cache_key if isinstance(value, TypeEngine) else value)
+        for key, value in fields
+    )
+
+
+class Struct(TypeEngine):
+    """
+    Represents a STRUCT type in DuckDB
+
+    ```python
+    from duckdb_sqlalchemy.datatypes import Struct
+    from sqlalchemy import Table, Column, String
+
+    Table(
+        'hello',
+        Column('name', Struct({'first': String, 'last': String})
+    )
+    ```
+
+    :param fields: only optional due to limitations with how much type information DuckDB returns to us in the description field
+    """
+
+    __visit_name__ = "struct"
+    cache_ok = True
+
+    def __init__(self, fields: Optional[Dict[str, TV]] = None):
+        self.fields = fields
+        self._fields = _normalize_fields(fields)
+        self._fields_cache_key = _normalize_fields_cache_key(self._fields)
+
+    @util.memoized_property
+    def _static_cache_key(self):  # type: ignore[override]
+        if self._fields_cache_key is None:
+            return (self.__class__,)
+        return (self.__class__, ("fields", self._fields_cache_key))
+
+
+class Map(TypeEngine):
+    """
+    Represents a MAP type in DuckDB
+
+    ```python
+    from duckdb_sqlalchemy.datatypes import Map
+    from sqlalchemy import Table, Column, String
+
+    Table(
+        'hello',
+        Column('name', Map(String, String)
+    )
+    ```
+    """
+
+    __visit_name__ = "map"
+    cache_ok = True
+    key_type: TV
+    value_type: TV
+
+    def __init__(self, key_type: TV, value_type: TV):
+        self.key_type = type_api.to_instance(key_type)
+        self.value_type = type_api.to_instance(value_type)
+
+    def bind_processor(self, dialect: Dialect) -> Any:
+        return lambda value: (
+            {"key": list(value), "value": list(value.values())} if value else None
+        )
+
+    def result_processor(self, dialect: Dialect, coltype: object) -> Any:
+        if IS_GT_1:
+            return lambda value: value
+        else:
+            return (
+                lambda value: dict(zip(value["key"], value["value"])) if value else {}
+            )
+
+
+class Union(TypeEngine):
+    """
+    Represents a UNION type in DuckDB
+
+    ```python
+    from duckdb_sqlalchemy.datatypes import Union
+    from sqlalchemy import Table, Column, String
+
+    Table(
+        'hello',
+        Column('name', Union({"name": String, "age": String})
+    )
+    ```
+    """
+
+    __visit_name__ = "union"
+    cache_ok = True
+    fields: Dict[str, TV]
+
+    def __init__(self, fields: Dict[str, TV]):
+        self.fields = fields
+        self._fields = _normalize_fields(fields)
+        self._fields_cache_key = _normalize_fields_cache_key(self._fields)
+
+    @util.memoized_property
+    def _static_cache_key(self):  # type: ignore[override]
+        if self._fields_cache_key is None:
+            return (self.__class__,)
+        return (self.__class__, ("fields", self._fields_cache_key))
+
+
+ISCHEMA_NAMES = {
+    "bignum": sqltypes.Numeric,
+    "hugeint": HugeInteger,
+    "uhugeint": UHugeInteger,
+    "tinyint": TinyInteger,
+    "utinyint": UTinyInteger,
+    "int8": BigInteger,
+    "int4": Integer,
+    "int2": SmallInteger,
+    "timetz": sqltypes.TIME,
+    "timestamptz": sqltypes.TIMESTAMP,
+    "float4": sqltypes.FLOAT,
+    "float8": sqltypes.FLOAT,
+    "usmallint": USmallInteger,
+    "uinteger": UInteger,
+    "ubigint": UBigInteger,
+    "timestamp_s": sqltypes.TIMESTAMP,
+    "timestamp_ms": sqltypes.TIMESTAMP,
+    "timestamp_ns": sqltypes.TIMESTAMP,
+    "enum": sqltypes.Enum,
+    "bool": sqltypes.BOOLEAN,
+    "varchar": String,
+}
+if IS_GT_1:
+    ISCHEMA_NAMES["varint"] = VarInt
+
+
+def register_extension_types() -> None:
+    for subclass in types:
+        compiles(subclass, "duckdb")(compile_uint)
+
+
+@compiles(Struct, "duckdb")  # type: ignore[misc]
+def visit_struct(
+    instance: Struct,
+    compiler: PGTypeCompiler,
+    identifier_preparer: PGIdentifierPreparer,
+    **kw: Any,
+) -> str:
+    return "STRUCT" + struct_or_union(instance, compiler, identifier_preparer, **kw)
+
+
+@compiles(Union, "duckdb")  # type: ignore[misc]
+def visit_union(
+    instance: Union,
+    compiler: PGTypeCompiler,
+    identifier_preparer: PGIdentifierPreparer,
+    **kw: Any,
+) -> str:
+    return "UNION" + struct_or_union(instance, compiler, identifier_preparer, **kw)
+
+
+def struct_or_union(
+    instance: typing.Union[Union, Struct],
+    compiler: PGTypeCompiler,
+    identifier_preparer: PGIdentifierPreparer,
+    **kw: Any,
+) -> str:
+    fields = getattr(instance, "_fields", None)
+    if fields is None:
+        fields = instance.fields
+    if fields is None:
+        raise exc.CompileError(f"DuckDB {repr(instance)} type requires fields")
+    return "({})".format(
+        ", ".join(
+            "{} {}".format(
+                identifier_preparer.quote_identifier(key),
+                process_type(
+                    value, compiler, identifier_preparer=identifier_preparer, **kw
+                ),
+            )
+            for key, value in (fields.items() if isinstance(fields, dict) else fields)
+        )
+    )
+
+
+def process_type(
+    value: typing.Union[TypeEngine, Type[TypeEngine]],
+    compiler: PGTypeCompiler,
+    **kw: Any,
+) -> str:
+    return compiler.process(type_api.to_instance(value), **kw)
+
+
+@compiles(Map, "duckdb")  # type: ignore[misc]
+def visit_map(instance: Map, compiler: PGTypeCompiler, **kw: Any) -> str:
+    return "MAP({}, {})".format(
+        process_type(instance.key_type, compiler, **kw),
+        process_type(instance.value_type, compiler, **kw),
+    )