From 15161efc131bd912bbce05b5f53aeee4c18dcdbd Mon Sep 17 00:00:00 2001 From: Hayssam Saleh Date: Mon, 10 Nov 2025 15:00:11 +0100 Subject: [PATCH 1/3] Add the ability to run actions right after connecting to the database --- README.md | 22 +++++++++++++++++++++- duckdb_engine/__init__.py | 4 ++++ duckdb_engine/tests/test_basic.py | 15 +++++++++++++++ 3 files changed, 40 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 1149927c6..e5990f402 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ Basic SQLAlchemy driver for [DuckDB](https://duckdb.org/) - [Alembic Integration](#alembic-integration) - [Preloading extensions (experimental)](#preloading-extensions-experimental) - [Registering Filesystems](#registering-filesystems) + - [Running actions right after connecting](#running-actions-right-after-connecting) - [The name](#the-name) @@ -176,7 +177,8 @@ create_engine( 'preload_extensions': ['https'], 'config': { 's3_region': 'ap-southeast-1' - } + }, + 'pre_actions': ["ATTACH 'file.db' AS file_db';"] } ) ``` @@ -199,6 +201,24 @@ create_engine( ) ``` +## Running actions right after connecting +You can run arbitrary SQL commands right after connecting by passing a list of SQL commands to the `pre_actions` parameter in `connect_args` + +```python +from sqlalchemy import create_engine +create_engine( + 'duckdb:///:memory:', + connect_args={ + 'pre_actions': [ + "ATTACH 'file.db' AS file_db';", + "SET some_config_option='some_value';" + ] + } +) +``` + ## 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 + + diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index e6b1680e2..f779018bd 100644 --- a/duckdb_engine/__init__.py +++ b/duckdb_engine/__init__.py @@ -285,6 +285,7 @@ def type_descriptor(self, typeobj: Type[sqltypes.TypeEngine]) -> Any: # type: i def connect(self, *cargs: Any, **cparams: Any) -> "Connection": core_keys = get_core_config() preload_extensions = cparams.pop("preload_extensions", []) + pre_actions = cparams.pop("pre_actions", []) config = dict(cparams.get("config", {})) cparams["config"] = config config.update(cparams.pop("url_config", {})) @@ -306,6 +307,9 @@ def connect(self, *cargs: Any, **cparams: Any) -> "Connection": for filesystem in filesystems: conn.register_filesystem(filesystem) + for action in pre_actions: + conn.execute(action) + apply_config(self, conn, ext) return ConnectionWrapper(conn) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 15fda1fea..3c8abff89 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -278,6 +278,21 @@ def test_preload_extension() -> None: text("SELECT * FROM read_parquet('https://domain/path/to/file.parquet');") ) +def test_pre_actions() -> None: + engine = create_engine( + "duckdb:///", + connect_args={ + "pre_actions": ["INSTALL SPATIAL", "LOAD SPATIAL"], + "config": {"s3_region": "ap-southeast-2", "s3_use_ssl": True}, + }, + ) + + # check that we get an error indicating that the extension was loaded + with engine.connect() as conn: + conn.execute( + text("SELECT ST_Affine(ST_Point(1, 1),1, 0, 0, 1, 2, 3);") + ) + @fixture def inspector(engine: Engine, session: Session) -> Inspector: From 3efe3ca7d172774cdf7448ce06f7918898bc4ae7 Mon Sep 17 00:00:00 2001 From: Hayssam Saleh Date: Mon, 10 Nov 2025 15:05:48 +0100 Subject: [PATCH 2/3] fix ruff formatt --- duckdb_engine/tests/test_basic.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index 3c8abff89..066bb7a69 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -278,6 +278,7 @@ def test_preload_extension() -> None: text("SELECT * FROM read_parquet('https://domain/path/to/file.parquet');") ) + def test_pre_actions() -> None: engine = create_engine( "duckdb:///", @@ -289,9 +290,7 @@ def test_pre_actions() -> None: # check that we get an error indicating that the extension was loaded with engine.connect() as conn: - conn.execute( - text("SELECT ST_Affine(ST_Point(1, 1),1, 0, 0, 1, 2, 3);") - ) + conn.execute(text("SELECT ST_Affine(ST_Point(1, 1),1, 0, 0, 1, 2, 3);")) @fixture From 1296df751ea1ddcd309f3e1f4da716c95fef542d Mon Sep 17 00:00:00 2001 From: Hayssam Saleh Date: Mon, 10 Nov 2025 20:07:47 +0100 Subject: [PATCH 3/3] fix commment in test --- 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 066bb7a69..d70af2021 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/duckdb_engine/tests/test_basic.py @@ -288,7 +288,7 @@ def test_pre_actions() -> None: }, ) - # check that we get an error indicating that the extension was loaded + # check that we can use spatial functions with engine.connect() as conn: conn.execute(text("SELECT ST_Affine(ST_Point(1, 1),1, 0, 0, 1, 2, 3);"))