From 13a92e1fa7d6bdb5777b46c234cb00a150978e9c Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 00:50:58 +0800 Subject: [PATCH 01/22] feat: allow preloading of extensions --- duckdb_engine/__init__.py | 36 ++++++++++++++++++++++++++----- duckdb_engine/config.py | 14 ++++++++++++ duckdb_engine/tests/test_basic.py | 17 ++++++++++++++- 3 files changed, 61 insertions(+), 6 deletions(-) create mode 100644 duckdb_engine/config.py diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index f1bd0299..345e237a 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type import duckdb -from sqlalchemy import pool +from sqlalchemy import String, pool from sqlalchemy import types as sqltypes from sqlalchemy import util from sqlalchemy.dialects.postgresql.base import PGInspector, PGTypeCompiler @@ -10,6 +10,7 @@ from sqlalchemy.ext.compiler import compiles from . import datatypes +from .config import get_core_config __version__ = "0.5.0" @@ -33,6 +34,9 @@ class DBAPI: # this is being fixed upstream to add a proper exception hierarchy Error = getattr(duckdb, "Error", RuntimeError) + IOException = getattr(duckdb, "IOException", RuntimeError) + CatalogException = getattr(duckdb, "CatalogException", RuntimeError) + @staticmethod def Binary(x: Any) -> Any: return x @@ -144,12 +148,34 @@ class Dialect(PGDialect_psycopg2): }, ) - def __init__(self, *args: Any, **kwargs: Any) -> None: + def __init__(self, **kwargs: Any) -> None: kwargs["use_native_hstore"] = False - super().__init__(*args, **kwargs) + super().__init__(**kwargs) def connect(self, *cargs: Any, **cparams: Any) -> "Connection": - return ConnectionWrapper(duckdb.connect(*cargs, **cparams)) + + core_keys = get_core_config() + preload_extensions = cparams.pop("preload_extensions", []) + config = cparams.get("config", {}) + + ext = {k: config.pop(k) for k in list(config) if k not in core_keys} + + conn = duckdb.connect(*cargs, **cparams) + + for extension in preload_extensions: + try: + conn.execute(f"LOAD {extension}") + except self.dbapi().IOException: + pass + + for k, v in ext.items(): + v = String().literal_processor(dialect=self)(v) + try: + conn.execute(f"SET {k} = {v}") + except self.dbapi().CatalogException: + pass + + return ConnectionWrapper(conn) def on_connect(self) -> None: pass @@ -189,7 +215,7 @@ def get_view_names( connection: Any, schema: Optional[Any] = ..., include: Any = ..., - **kw: Any + **kw: Any, ) -> Any: s = "SELECT name FROM sqlite_master WHERE type='view' ORDER BY name" rs = connection.exec_driver_sql(s) diff --git a/duckdb_engine/config.py b/duckdb_engine/config.py new file mode 100644 index 00000000..c0100e51 --- /dev/null +++ b/duckdb_engine/config.py @@ -0,0 +1,14 @@ +from functools import lru_cache +from typing import Set + +import duckdb + + +@lru_cache() +def get_core_config() -> Set[str]: + rows = ( + duckdb.connect(":memory:") + .execute("SELECT name FROM duckdb_settings()") + .fetchall() + ) + return {name for name, in rows} diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 5c1b4e41..0b3ebd2a 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -7,7 +7,7 @@ import duckdb from hypothesis import assume, given, settings from hypothesis.strategies import text as text_strat -from pytest import LogCaptureFixture, fixture, importorskip, mark, raises +from pytest import LogCaptureFixture, fixture, importorskip, mark, raises, skip from sqlalchemy import ( Column, ForeignKey, @@ -151,6 +151,21 @@ def test_get_views(engine: Engine) -> None: assert views == ["test"] +def test_preload_extension() -> None: + try: + duckdb.default_connection.execute("INSTALL https") + except Exception as e: + skip(str(e)) + engine = create_engine( + "duckdb:///", + connect_args={ + "preload_extensions": ["httpfs"], + "config": {"s3_region": "ap-southeast-2"}, + }, + ) + engine.connect() + + @fixture def inspector(engine: Engine, session: Session) -> Inspector: session.execute(text("create table test (id int);")) From bd116172862fe2cc5c982ef2be4f0495543eef5e Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 00:59:57 +0800 Subject: [PATCH 02/22] chore: verbose pytest --- tox.ini | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tox.ini b/tox.ini index a5ece71b..c84e2b4f 100644 --- a/tox.ini +++ b/tox.ini @@ -21,10 +21,10 @@ whitelist_externals = poetry commands = poetry install -v poetry run pip install duckdb --pre -U - poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml + poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml -v [testenv] whitelist_externals = poetry commands = poetry install -v - poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml + poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml -v From b8cf370a75105d8450fcd1578c7ccafacb653f4a Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 01:09:33 +0800 Subject: [PATCH 03/22] chore: set unit test check name --- .github/workflows/pythonapp.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/pythonapp.yaml b/.github/workflows/pythonapp.yaml index eef24da0..87872dfb 100644 --- a/.github/workflows/pythonapp.yaml +++ b/.github/workflows/pythonapp.yaml @@ -52,6 +52,7 @@ jobs: if: always() with: junit_files: results.xml + check_name: "Test Results - ${{ join(matrix.*, ' - ') }}" - uses: codecov/codecov-action@v3 with: files: ./coverage.xml From 3bb817c3f8f17f55a0b54831d78d370061650394 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 01:17:21 +0800 Subject: [PATCH 04/22] chore: print skip reasons --- tox.ini | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tox.ini b/tox.ini index c84e2b4f..cf75560f 100644 --- a/tox.ini +++ b/tox.ini @@ -21,10 +21,10 @@ whitelist_externals = poetry commands = poetry install -v poetry run pip install duckdb --pre -U - poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml -v + poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml --verbose -rs [testenv] whitelist_externals = poetry commands = poetry install -v - poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml -v + poetry run pytest --junitxml=results.xml --cov --cov-report xml:coverage.xml --verbose -rs From 694bf79e4b8b35c01b033eaaa0e14f79126593ca Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 01:20:59 +0800 Subject: [PATCH 05/22] chore: fix typo --- duckdb_engine/tests/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 0b3ebd2a..059b0ac3 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -153,7 +153,7 @@ def test_get_views(engine: Engine) -> None: def test_preload_extension() -> None: try: - duckdb.default_connection.execute("INSTALL https") + duckdb.default_connection.execute("INSTALL httpfs") except Exception as e: skip(str(e)) engine = create_engine( From 62f7286d9e61c317aec89402f6f9b0441ac98531 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 01:34:13 +0800 Subject: [PATCH 06/22] chore: tweak skip --- duckdb_engine/tests/test_basic.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 059b0ac3..a16e59ae 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -1,4 +1,5 @@ import logging +import os import zlib from datetime import timedelta from pathlib import Path @@ -7,7 +8,7 @@ import duckdb from hypothesis import assume, given, settings from hypothesis.strategies import text as text_strat -from pytest import LogCaptureFixture, fixture, importorskip, mark, raises, skip +from pytest import LogCaptureFixture, fixture, importorskip, mark, raises from sqlalchemy import ( Column, ForeignKey, @@ -151,11 +152,9 @@ def test_get_views(engine: Engine) -> None: assert views == ["test"] +@mark.skipif(os.uname().machine == "aarch64", reason="not supported on aarch64") def test_preload_extension() -> None: - try: - duckdb.default_connection.execute("INSTALL httpfs") - except Exception as e: - skip(str(e)) + duckdb.default_connection.execute("INSTALL httpfs") engine = create_engine( "duckdb:///", connect_args={ From 2f028f2c2544ba5f573c99d57dc799ff0e0fbdbf Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 01:36:59 +0800 Subject: [PATCH 07/22] chore: try to assert extension was loaded --- duckdb_engine/tests/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index a16e59ae..233c42b5 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -162,7 +162,7 @@ def test_preload_extension() -> None: "config": {"s3_region": "ap-southeast-2"}, }, ) - engine.connect() + assert engine.connect().execute("select * From duckdb_extensions() where loaded") @fixture From 73de97f940c66de44def2ee300b072d6111c2849 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 01:45:18 +0800 Subject: [PATCH 08/22] chore: extend assertions --- duckdb_engine/tests/test_basic.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 233c42b5..ccb456db 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -162,7 +162,13 @@ def test_preload_extension() -> None: "config": {"s3_region": "ap-southeast-2"}, }, ) - assert engine.connect().execute("select * From duckdb_extensions() where loaded") + with engine.connect() as conn: + assert conn.execute( + "select * from duckdb_extensions() where loaded and name='httpfs'" + ).fetchone() + assert not conn.execute( + "select * from duckdb_extensions() where loaded and name='fts'" + ).fetchone() @fixture From bb600ffac8e3afe383fcc88581cca8374de022e1 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 19:31:41 +0800 Subject: [PATCH 09/22] chore: tweak extension load check --- duckdb_engine/tests/test_basic.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index ccb456db..7402b409 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -162,13 +162,12 @@ def test_preload_extension() -> None: "config": {"s3_region": "ap-southeast-2"}, }, ) - with engine.connect() as conn: - assert conn.execute( - "select * from duckdb_extensions() where loaded and name='httpfs'" - ).fetchone() - assert not conn.execute( - "select * from duckdb_extensions() where loaded and name='fts'" - ).fetchone() + + # check that we get an error indicating that the extension was loaded + with engine.connect() as conn, raises(DBAPI.Error, match="unreachable"): + conn.execute( + "SELECT * FROM read_parquet('https:///path/to/file.parquet');'" + ) @fixture From 43c5cd84116a9cfdcae58b85fb92b6d909f53d48 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 19:35:34 +0800 Subject: [PATCH 10/22] chore: fix URL --- duckdb_engine/tests/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 7402b409..6651d049 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -166,7 +166,7 @@ def test_preload_extension() -> None: # check that we get an error indicating that the extension was loaded with engine.connect() as conn, raises(DBAPI.Error, match="unreachable"): conn.execute( - "SELECT * FROM read_parquet('https:///path/to/file.parquet');'" + "SELECT * FROM read_parquet('https://domain/path/to/file.parquet');'" ) From 1b0fcd301d512f102868b0c0c17cd82d2fa60d94 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 19:38:51 +0800 Subject: [PATCH 11/22] chore: remove trailing quote --- duckdb_engine/tests/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 6651d049..414d15f2 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -166,7 +166,7 @@ def test_preload_extension() -> None: # check that we get an error indicating that the extension was loaded with engine.connect() as conn, raises(DBAPI.Error, match="unreachable"): conn.execute( - "SELECT * FROM read_parquet('https://domain/path/to/file.parquet');'" + "SELECT * FROM read_parquet('https://domain/path/to/file.parquet');" ) From 624a60ff7a90da074dba166500c9e430d1aa0988 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 19:41:57 +0800 Subject: [PATCH 12/22] chore: update expected error message --- duckdb_engine/tests/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 414d15f2..9d732e1c 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -164,7 +164,7 @@ def test_preload_extension() -> None: ) # check that we get an error indicating that the extension was loaded - with engine.connect() as conn, raises(DBAPI.Error, match="unreachable"): + with engine.connect() as conn, raises(DBAPI.Error, match="HTTP HEAD error"): conn.execute( "SELECT * FROM read_parquet('https://domain/path/to/file.parquet');" ) From a2955f823054041041f108e507cd308e13fbbfae Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 19:45:00 +0800 Subject: [PATCH 13/22] chore: remove unused try catches --- duckdb_engine/__init__.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index 345e237a..24428f60 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -163,17 +163,11 @@ def connect(self, *cargs: Any, **cparams: Any) -> "Connection": conn = duckdb.connect(*cargs, **cparams) for extension in preload_extensions: - try: - conn.execute(f"LOAD {extension}") - except self.dbapi().IOException: - pass + conn.execute(f"LOAD {extension}") for k, v in ext.items(): v = String().literal_processor(dialect=self)(v) - try: - conn.execute(f"SET {k} = {v}") - except self.dbapi().CatalogException: - pass + conn.execute(f"SET {k} = {v}") return ConnectionWrapper(conn) From eea31684647783debf19d4f943765c8b2577e192 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 20:06:09 +0800 Subject: [PATCH 14/22] chore: correct exception --- duckdb_engine/tests/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 9d732e1c..cdcadbd4 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -164,7 +164,7 @@ def test_preload_extension() -> None: ) # check that we get an error indicating that the extension was loaded - with engine.connect() as conn, raises(DBAPI.Error, match="HTTP HEAD error"): + with engine.connect() as conn, raises(DBAPIError, match="HTTP HEAD error"): conn.execute( "SELECT * FROM read_parquet('https://domain/path/to/file.parquet');" ) From 0ac0e44a0d1c5a9b174c57d10fc049a291902e91 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 20:10:53 +0800 Subject: [PATCH 15/22] chore: correct exception type again --- duckdb_engine/tests/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index cdcadbd4..3da94fe5 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -164,7 +164,7 @@ def test_preload_extension() -> None: ) # check that we get an error indicating that the extension was loaded - with engine.connect() as conn, raises(DBAPIError, match="HTTP HEAD error"): + with engine.connect() as conn, raises(Exception, match="HTTP HEAD error"): conn.execute( "SELECT * FROM read_parquet('https://domain/path/to/file.parquet');" ) From 05e852d218ef0e5775e211c21f07066d5c291251 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 22:45:35 +0800 Subject: [PATCH 16/22] chore: undo dud changes --- duckdb_engine/__init__.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index 24428f60..235abfe4 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -34,9 +34,6 @@ class DBAPI: # this is being fixed upstream to add a proper exception hierarchy Error = getattr(duckdb, "Error", RuntimeError) - IOException = getattr(duckdb, "IOException", RuntimeError) - CatalogException = getattr(duckdb, "CatalogException", RuntimeError) - @staticmethod def Binary(x: Any) -> Any: return x @@ -148,9 +145,9 @@ class Dialect(PGDialect_psycopg2): }, ) - def __init__(self, **kwargs: Any) -> None: + def __init__(self, *args, **kwargs: Any) -> None: kwargs["use_native_hstore"] = False - super().__init__(**kwargs) + super().__init__(*args, **kwargs) def connect(self, *cargs: Any, **cparams: Any) -> "Connection": @@ -209,7 +206,7 @@ def get_view_names( connection: Any, schema: Optional[Any] = ..., include: Any = ..., - **kw: Any, + **kw: Any ) -> Any: s = "SELECT name FROM sqlite_master WHERE type='view' ORDER BY name" rs = connection.exec_driver_sql(s) From 3d0f2e7caf824c2960efb41b162621c9cc8c0091 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 20 Aug 2022 14:45:54 +0000 Subject: [PATCH 17/22] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- duckdb_engine/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index 235abfe4..dbc0c57e 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -206,7 +206,7 @@ def get_view_names( connection: Any, schema: Optional[Any] = ..., include: Any = ..., - **kw: Any + **kw: Any, ) -> Any: s = "SELECT name FROM sqlite_master WHERE type='view' ORDER BY name" rs = connection.exec_driver_sql(s) From 0f600801d05c213a09d6999b43f316b434be1fc1 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sat, 20 Aug 2022 22:46:28 +0800 Subject: [PATCH 18/22] chore: undo another change --- duckdb_engine/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index dbc0c57e..d02ec963 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -145,7 +145,7 @@ class Dialect(PGDialect_psycopg2): }, ) - def __init__(self, *args, **kwargs: Any) -> None: + def __init__(self, *args: Any, **kwargs: Any) -> None: kwargs["use_native_hstore"] = False super().__init__(*args, **kwargs) @@ -206,7 +206,7 @@ def get_view_names( connection: Any, schema: Optional[Any] = ..., include: Any = ..., - **kw: Any, + **kw: Any ) -> Any: s = "SELECT name FROM sqlite_master WHERE type='view' ORDER BY name" rs = connection.exec_driver_sql(s) From 1e7f9d0c74fd2913c1c9f61769766ed5a26826c0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 20 Aug 2022 14:46:47 +0000 Subject: [PATCH 19/22] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- duckdb_engine/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index d02ec963..54c92527 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -206,7 +206,7 @@ def get_view_names( connection: Any, schema: Optional[Any] = ..., include: Any = ..., - **kw: Any + **kw: Any, ) -> Any: s = "SELECT name FROM sqlite_master WHERE type='view' ORDER BY name" rs = connection.exec_driver_sql(s) From 7ac783334d8c395961fae6de69c2edaa7f196b39 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sun, 21 Aug 2022 15:55:45 +0800 Subject: [PATCH 20/22] chore: move apply_config to config.py --- duckdb_engine/__init__.py | 9 +++------ duckdb_engine/config.py | 12 +++++++++++- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index 54c92527..8413eb5f 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type import duckdb -from sqlalchemy import String, pool +from sqlalchemy import pool from sqlalchemy import types as sqltypes from sqlalchemy import util from sqlalchemy.dialects.postgresql.base import PGInspector, PGTypeCompiler @@ -10,7 +10,7 @@ from sqlalchemy.ext.compiler import compiles from . import datatypes -from .config import get_core_config +from .config import apply_config, get_core_config __version__ = "0.5.0" @@ -132,7 +132,6 @@ class Dialect(PGDialect_psycopg2): supports_statement_cache = False supports_comments = False supports_sane_rowcount = False - supports_comments = False inspector = DuckDBInspector # colspecs TODO: remap types to duckdb types colspecs = util.update_copy( @@ -162,9 +161,7 @@ def connect(self, *cargs: Any, **cparams: Any) -> "Connection": for extension in preload_extensions: conn.execute(f"LOAD {extension}") - for k, v in ext.items(): - v = String().literal_processor(dialect=self)(v) - conn.execute(f"SET {k} = {v}") + apply_config(self, conn, ext) return ConnectionWrapper(conn) diff --git a/duckdb_engine/config.py b/duckdb_engine/config.py index c0100e51..9783a016 100644 --- a/duckdb_engine/config.py +++ b/duckdb_engine/config.py @@ -1,7 +1,9 @@ from functools import lru_cache -from typing import Set +from typing import Dict, Set import duckdb +from sqlalchemy import String +from sqlalchemy.engine import Dialect @lru_cache() @@ -12,3 +14,11 @@ def get_core_config() -> Set[str]: .fetchall() ) return {name for name, in rows} + + +def apply_config( + dialect: Dialect, conn: duckdb.DuckDBPyConnection, ext: Dict[str, str] +) -> None: + process = String().literal_processor(dialect=dialect) + for k, v in ext.items(): + conn.execute(f"SET {process(k)} = {process(v)}") From c0f2a993ee826fa91cab2add7321c52e437bb5a8 Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sun, 21 Aug 2022 16:03:13 +0800 Subject: [PATCH 21/22] docs: document preload_extensions config parameter --- README.md | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/README.md b/README.md index cc7ea815..63417b74 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,20 @@ Basic SQLAlchemy driver for [DuckDB](https://duckdb.org/) + +* [duckdb_engine](#duckdb_engine) + * [Installation](#installation) + * [Usage](#usage) + * [Configuration](#configuration) + * [How to register a pandas DataFrame](#how-to-register-a-pandas-dataframe) + * [Things to keep in mind](#things-to-keep-in-mind) + * [Auto-incrementing ID columns](#auto-incrementing-id-columns) + * [Pandas read_sql() chunksize](#pandas-read_sql-chunksize) + * [Unsigned integer support](#unsigned-integer-support) + * [Preloading extensions (experimental)](#preloading-extensions-experimental) + * [The name](#the-name) + + ## Installation ```sh $ pip install duckdb-engine @@ -110,6 +124,24 @@ The `pandas.read_sql()` method can read tables from `duckdb_engine` into DataFra Unsigned integers are supported by DuckDB, and are available in [`duckdb_engine.datatypes`](duckdb_engine/datatypes.py). +## Preloading extensions (experimental) + +Until the DuckDB python client allows you to natively preload extensions, I've added experimental support via a `connect_args` parameter + +```python +from sqlalchemy import create_engine + +create_engine( + 'duckdb:///:memory:', + connect_args={ + 'preload_extensions': ['https'], + 'config': { + 's3_region': 'ap-southeast-1' + } + } +) +``` + ## The name Yes, I'm aware this package should be named `duckdb-driver` or something, I wasn't thinking when I named it and it's too hard to change the name now From ba3fbe864ec9b225e995863f759f8377a5f8546b Mon Sep 17 00:00:00 2001 From: Elliana May Date: Sun, 21 Aug 2022 16:08:03 +0800 Subject: [PATCH 22/22] chore: remove overzealous escaping --- duckdb_engine/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/duckdb_engine/config.py b/duckdb_engine/config.py index 9783a016..fe197139 100644 --- a/duckdb_engine/config.py +++ b/duckdb_engine/config.py @@ -21,4 +21,4 @@ def apply_config( ) -> None: process = String().literal_processor(dialect=dialect) for k, v in ext.items(): - conn.execute(f"SET {process(k)} = {process(v)}") + conn.execute(f"SET {k} = {process(v)}")