Skip to content

Commit

Permalink
fix: support boolean and integer config values
Browse files Browse the repository at this point in the history
  • Loading branch information
Mause committed Aug 23, 2022
1 parent ac19e34 commit 4a2c639
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 5 deletions.
16 changes: 12 additions & 4 deletions duckdb_engine/config.py
Original file line number Diff line number Diff line change
@@ -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()
Expand All @@ -17,8 +20,13 @@ 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)
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)}")
2 changes: 1 addition & 1 deletion duckdb_engine/tests/test_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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},
},
)

Expand Down

0 comments on commit 4a2c639

Please sign in to comment.