diff --git a/duckdb_engine/config.py b/duckdb_engine/config.py index fe197139..a32f72ef 100644 --- a/duckdb_engine/config.py +++ b/duckdb_engine/config.py @@ -1,9 +1,12 @@ from functools import lru_cache -from typing import Dict, Set +from typing import Dict, Set, Type, Union import duckdb -from sqlalchemy import String +from sqlalchemy import Boolean, Integer, String from sqlalchemy.engine import Dialect +from sqlalchemy.sql.type_api import TypeEngine + +TYPES: Dict[Type, TypeEngine] = {int: Integer(), str: String(), bool: Boolean()} @lru_cache() @@ -17,8 +20,14 @@ def get_core_config() -> Set[str]: def apply_config( - dialect: Dialect, conn: duckdb.DuckDBPyConnection, ext: Dict[str, str] + dialect: Dialect, + conn: duckdb.DuckDBPyConnection, + ext: Dict[str, Union[str, int, bool]], ) -> None: - process = String().literal_processor(dialect=dialect) + # TODO: does sqlalchemy have something that could do this for us? + processors = {k: v.literal_processor(dialect=dialect) for k, v in TYPES.items()} + for k, v in ext.items(): + process = processors[type(v)] + assert process, f"Not able to configure {k} with {v}" conn.execute(f"SET {k} = {process(v)}") diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 3da94fe5..6a405d99 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -159,7 +159,7 @@ def test_preload_extension() -> None: "duckdb:///", connect_args={ "preload_extensions": ["httpfs"], - "config": {"s3_region": "ap-southeast-2"}, + "config": {"s3_region": "ap-southeast-2", "s3_use_ssl": True}, }, )