Merge pull request #502 from Mause/refactor/type-cleanliness

refactor: type cleanliness for SQLA2 support

Author: kodiakhq[bot]

duckdb_engine/__init__.py

@@ -2,7 +2,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
 
 import duckdb
-from sqlalchemy import pool
+from sqlalchemy import pool, text
 from sqlalchemy import types as sqltypes
 from sqlalchemy import util
 from sqlalchemy.dialects.postgresql.base import PGInspector
@@ -61,7 +61,7 @@ def cursor(self) -> "Connection":
         return self
 
     def fetchmany(self, size: Optional[int] = None) -> List:
-        if hasattr(self.c, "fetchmany"):
+        if hasattr(self.__c, "fetchmany"):
             # fetchmany was only added in 0.5.0
             if size is None:
                 return self.__c.fetchmany()
@@ -223,8 +223,10 @@ def get_view_names(
         include: Optional[Any] = None,
         **kw: Any,
     ) -> Any:
-        s = "SELECT table_name FROM information_schema.tables WHERE table_type='VIEW' and table_schema=?"
-        rs = connection.execute(s, schema if schema is not None else "main")
+        s = "SELECT table_name FROM information_schema.tables WHERE table_type='VIEW' and table_schema=:schema_name"
+        rs = connection.execute(
+            text(s), {"schema_name": schema if schema is not None else "main"}
+        )
 
         return [row[0] for row in rs]
 

duckdb_engine/tests/test_basic.py

@@ -156,12 +156,11 @@ def test_get_views(engine: Engine) -> None:
     views = engine.dialect.get_view_names(con)
     assert views == []
 
-    engine.execute(text("create view test as select 1"))
-    engine.execute(
+    con.execute(text("create view test as select 1"))
+    con.execute(
         text("create schema scheme; create view scheme.schema_test as select 1")
     )
 
-    con = engine.connect()
     views = engine.dialect.get_view_names(con)
     assert views == ["test"]
 
@@ -183,7 +182,7 @@ def test_preload_extension() -> None:
     # check that we get an error indicating that the extension was loaded
     with engine.connect() as conn, raises(Exception, match="HTTP HEAD"):
         conn.execute(
-            "SELECT * FROM read_parquet('https://domain/path/to/file.parquet');"
+            text("SELECT * FROM read_parquet('https://domain/path/to/file.parquet');")
         )
 
 
@@ -307,8 +306,9 @@ def test_inmemory() -> None:
     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()""")
+    shell.run_cell("""conn = eng.connect()""")
+    shell.run_cell("""conn.execute(sa.text("CREATE TABLE t (x int)"))""")
+    res = shell.run_cell("""conn.execute(sa.text("SHOW TABLES")).fetchall()""")
 
     assert res.result == [("t",)]
 
@@ -329,7 +329,8 @@ def test_config(tmp_path: Path) -> None:
         DBAPIError,
         match='Cannot execute statement of type "CREATE" (on database "test" which is attached )?in read-only mode!',
     ):
-        eng.execute("create table hello2 (i int)")
+        with eng.connect() as conn:
+            conn.execute(text("create table hello2 (i int)"))
 
 
 def test_do_ping(tmp_path: Path, caplog: LogCaptureFixture) -> None:

duckdb_engine/tests/test_integration.py

@@ -1,8 +1,14 @@
 import pandas as pd
+from sqlalchemy import text
 from sqlalchemy.engine import Engine
 
 
 def test_integration(engine: Engine) -> None:
-    engine.execute("register", ("test_df", pd.DataFrame([{"a": 1}])))
+    with engine.connect() as conn:
+        execute = (
+            conn.exec_driver_sql if hasattr(conn, "exec_driver_sql") else conn.execute
+        )
+        params = ("test_df", pd.DataFrame([{"a": 1}]))
+        execute("register", params)  # type: ignore[operator]
 
-    engine.execute("select * from test_df")
+        conn.execute(text("select * from test_df"))