Merge pull request #691 from Mause/motherduck

fix: allow motherduck to pass their special token parameter at extension load time

Author: Mause

duckdb_engine/config.py

@@ -16,7 +16,8 @@ def get_core_config() -> Set[str]:
         .execute("SELECT name FROM duckdb_settings()")
         .fetchall()
     )
-    return {name for name, in rows}
+    # special case for motherduck here - they accept this config at extension load time
+    return {name for name, in rows} | {"motherduck_token"}
 
 
 def apply_config(

duckdb_engine/tests/test_integration.py

@@ -1,6 +1,16 @@
+import sys
+from subprocess import PIPE, Popen
+from traceback import print_exc
+
+import duckdb
 import pandas as pd
+from pytest import importorskip, mark, raises
 from sqlalchemy import text
-from sqlalchemy.engine import Engine
+from sqlalchemy.engine import Engine, create_engine
+from sqlalchemy.exc import ProgrammingError
+
+SEGFAULT = -6
+SUCCESS = 0
 
 
 def test_integration(engine: Engine) -> None:
@@ -12,3 +22,44 @@ def test_integration(engine: Engine) -> None:
         execute("register", params)  # type: ignore[operator]
 
         conn.execute(text("select * from test_df"))
+
+
+@mark.remote_data
+@mark.skipif(
+    "dev" in duckdb.__version__, reason="md extension not available for dev builds"  # type: ignore[attr-defined]
+)
+def test_motherduck() -> None:
+    importorskip("duckdb", "0.7.1")
+
+    # we're running this test in a sub process due to occasional segfaults in the motherduck extension
+
+    proc = Popen([sys.executable, __file__], stderr=PIPE, stdout=PIPE, text=True)
+    wait = proc.wait(5000)
+    assert wait == SUCCESS or wait == SEGFAULT, (proc.stdout, proc.stderr)
+
+
+def hook() -> None:
+    try:
+        _test_motherduck()
+    except Exception:
+        print_exc()
+        sys.exit(-1)
+    else:
+        sys.exit(0)
+
+
+def _test_motherduck() -> None:
+    engine = create_engine(
+        "duckdb:///md:motherdb",
+        connect_args={"config": {"motherduck_token": "motherduckdb_token"}},
+    )
+
+    with raises(
+        ProgrammingError,
+        match="Jwt is not in the form of Header.Payload.Signature with two dots and 3 sections",
+    ):
+        engine.connect()
+
+
+if __name__ == "__main__":
+    _test_motherduck()