Merge pull request #341 from Mause/bugfix/use-singletonthreadpool

feat: use SingletonThreadPool for :memory: connections

duckdb_engine/__init__.py

@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type
 
 import duckdb
+from sqlalchemy import pool
 from sqlalchemy import types as sqltypes
 from sqlalchemy import util
 from sqlalchemy.dialects.postgresql import dialect as postgres_dialect
@@ -184,6 +185,13 @@ def do_execute(
     ) -> None:
         cursor.execute(statement, parameters, context)
 
+    @classmethod
+    def get_pool_class(cls, url: URL) -> Type[pool.Pool]:
+        if url.database == ":memory:":
+            return pool.SingletonThreadPool
+        else:
+            return pool.QueuePool
+
     def do_executemany(
         self,
         cursor: ConnectionWrapper,

duckdb_engine/tests/test_basic.py

@@ -4,7 +4,7 @@
 
 from hypothesis import assume, given, settings
 from hypothesis.strategies import text
-from pytest import fixture, mark
+from pytest import fixture, importorskip, mark
 from sqlalchemy import (
     Column,
     ForeignKey,
@@ -188,7 +188,7 @@ def test_reflect(session: Session, engine: Engine) -> None:
 def test_commit(session: Session, engine: Engine) -> None:
     session.execute("commit;")
 
-    from IPython.core.interactiveshell import InteractiveShell
+    InteractiveShell = importorskip("IPython.core.interactiveshell").InteractiveShell
 
     shell = InteractiveShell()
     assert not shell.run_line_magic("load_ext", "sql")
@@ -262,3 +262,15 @@ def test_sessions(session: Session) -> None:
     c.field = timedelta(days=5)
     session.flush()
     session.commit()
+
+
+def test_inmemory() -> None:
+    InteractiveShell = importorskip("IPython.core.interactiveshell").InteractiveShell
+
+    shell = InteractiveShell()
+    shell.run_cell("""import sqlalchemy as sa""")
+    shell.run_cell("""eng = sa.create_engine("duckdb:///:memory:")""")
+    shell.run_cell("""eng.execute("CREATE TABLE t (x int)")""")
+    res = shell.run_cell("""eng.execute("SHOW TABLES").fetchall()""")
+
+    assert res.result == [("t",)]