test: Compatibility with alchemy 2.0.1 (#70)

Author: ptiurin

setup.cfg

@@ -42,6 +42,7 @@ sqlalchemy.dialects =
 [options.extras_require]
 dev =
     devtools==0.7.0
+    greenlet==2.0.2
     mock==4.0.3
     mypy==0.910
     pre-commit==2.15.0

src/firebolt_db/firebolt_dialect.py

@@ -148,6 +148,7 @@ def has_table(
         connection: AlchemyConnection,
         table_name: str,
         schema: Optional[str] = None,
+        **kw: Any
     ) -> bool:
         query = """
             select count(*) > 0 as exists_

tests/integration/conftest.py

@@ -3,7 +3,7 @@
 from os import environ
 
 from pytest import fixture
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, text
 from sqlalchemy.engine.base import Connection, Engine
 from sqlalchemy.ext.asyncio import create_async_engine
 
@@ -157,28 +157,32 @@ def setup_test_tables(
     dimension_table_name: str,
 ):
     connection.execute(
-        f"""
+        text(
+            f"""
         CREATE FACT TABLE IF NOT EXISTS {fact_table_name}
         (
             idx INT,
             dummy TEXT
         ) PRIMARY INDEX idx;
         """
+        )
     )
     connection.execute(
-        f"""
+        text(
+            f"""
         CREATE DIMENSION TABLE IF NOT EXISTS {dimension_table_name}
         (
             idx INT,
             dummy TEXT
         );
         """
+        )
     )
-    assert engine.dialect.has_table(engine, fact_table_name)
-    assert engine.dialect.has_table(engine, dimension_table_name)
+    assert engine.dialect.has_table(connection, fact_table_name)
+    assert engine.dialect.has_table(connection, dimension_table_name)
     yield
     # Teardown
-    connection.execute(f"DROP TABLE IF EXISTS {fact_table_name} CASCADE;")
-    connection.execute(f"DROP TABLE IF EXISTS {dimension_table_name} CASCADE;")
-    assert not engine.dialect.has_table(engine, fact_table_name)
-    assert not engine.dialect.has_table(engine, dimension_table_name)
+    connection.execute(text(f"DROP TABLE IF EXISTS {fact_table_name} CASCADE;"))
+    connection.execute(text(f"DROP TABLE IF EXISTS {dimension_table_name} CASCADE;"))
+    assert not engine.dialect.has_table(connection, fact_table_name)
+    assert not engine.dialect.has_table(connection, dimension_table_name)

tests/integration/test_sqlalchemy_integration.py

@@ -2,7 +2,7 @@
 from decimal import Decimal
 
 import pytest
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, text
 from sqlalchemy.engine.base import Connection, Engine
 from sqlalchemy.exc import OperationalError
 
@@ -15,11 +15,11 @@ def test_create_ex_table(
         ex_table_query: str,
         ex_table_name: str,
     ):
-        connection.execute(ex_table_query)
-        assert engine.dialect.has_table(engine, ex_table_name)
+        connection.execute(text(ex_table_query))
+        assert engine.dialect.has_table(connection, ex_table_name)
         # Cleanup
-        connection.execute(f"DROP TABLE {ex_table_name}")
-        assert not engine.dialect.has_table(engine, ex_table_name)
+        connection.execute(text(f"DROP TABLE {ex_table_name}"))
+        assert not engine.dialect.has_table(connection, ex_table_name)
 
     def test_set_params(
         self, username: str, password: str, database_name: str, engine_name: str
@@ -28,81 +28,95 @@ def test_set_params(
             f"firebolt://{username}:{password}@{database_name}/{engine_name}"
         )
         with engine.connect() as connection:
-            connection.execute("SET advanced_mode=1")
-            connection.execute("SET use_standard_sql=0")
-            result = connection.execute("SELECT sleepEachRow(1) from numbers(1)")
+            connection.execute(text("SET advanced_mode=1"))
+            connection.execute(text("SET use_standard_sql=0"))
+            result = connection.execute(text("SELECT sleepEachRow(1) from numbers(1)"))
             assert len(result.fetchall()) == 1
         engine.dispose()
 
     def test_data_write(self, connection: Connection, fact_table_name: str):
         connection.execute(
-            f"INSERT INTO {fact_table_name}(idx, dummy) VALUES (1, 'some_text')"
+            text(f"INSERT INTO {fact_table_name}(idx, dummy) VALUES (1, 'some_text')")
+        )
+        result = connection.execute(
+            text(f"SELECT * FROM {fact_table_name} WHERE idx=1")
         )
-        result = connection.execute(f"SELECT * FROM {fact_table_name} WHERE idx=?", 1)
         assert result.fetchall() == [(1, "some_text")]
-        result = connection.execute(f"SELECT * FROM {fact_table_name}")
+        result = connection.execute(text(f"SELECT * FROM {fact_table_name}"))
         assert len(result.fetchall()) == 1
         # Update not supported
         with pytest.raises(OperationalError):
             connection.execute(
-                f"UPDATE {fact_table_name} SET dummy='some_other_text' WHERE idx=1"
+                text(
+                    f"UPDATE {fact_table_name} SET dummy='some_other_text' WHERE idx=1"
+                )
             )
         # Delete works but is not officially supported yet
         # with pytest.raises(OperationalError):
         #     connection.execute(f"DELETE FROM {fact_table_name} WHERE idx=1")
 
     def test_firebolt_types(self, connection: Connection):
-        result = connection.execute("SELECT '1896-01-01' :: DATE_EXT")
+        result = connection.execute(text("SELECT '1896-01-01' :: DATE_EXT"))
         assert result.fetchall() == [(date(1896, 1, 1),)]
-        result = connection.execute("SELECT '1896-01-01 00:01:00' :: TIMESTAMP_EXT")
+        result = connection.execute(
+            text("SELECT '1896-01-01 00:01:00' :: TIMESTAMP_EXT")
+        )
         assert result.fetchall() == [(datetime(1896, 1, 1, 0, 1, 0, 0),)]
-        result = connection.execute("SELECT 100.76 :: DECIMAL(5, 2)")
+        result = connection.execute(text("SELECT 100.76 :: DECIMAL(5, 2)"))
         assert result.fetchall() == [(Decimal("100.76"),)]
 
     def test_agg_index(self, connection: Connection, fact_table_name: str):
         # Test if sql parsing allows it
         agg_index = "idx_agg_max"
         connection.execute(
-            f"""
+            text(
+                f"""
             CREATE AGGREGATING INDEX {agg_index} ON {fact_table_name} (
                 dummy,
                 max(idx)
             );
             """
+            )
         )
-        connection.execute(f"DROP AGGREGATING INDEX {agg_index}")
+        connection.execute(text(f"DROP AGGREGATING INDEX {agg_index}"))
 
     def test_join_index(self, connection: Connection, dimension_table_name: str):
         # Test if sql parsing allows it
         join_index = "idx_join"
         connection.execute(
-            f"""
+            text(
+                f"""
             CREATE JOIN INDEX {join_index} ON {dimension_table_name} (
                 idx,
                 dummy
             );
             """
+            )
         )
-        connection.execute(f"DROP JOIN INDEX {join_index}")
+        connection.execute(text(f"DROP JOIN INDEX {join_index}"))
 
     def test_get_schema_names(self, engine: Engine, database_name: str):
         results = engine.dialect.get_schema_names(engine)
         assert "public" in results
 
-    def test_has_table(self, engine: Engine, fact_table_name: str):
-        results = engine.dialect.has_table(engine, fact_table_name)
+    def test_has_table(
+        self, engine: Engine, connection: Connection, fact_table_name: str
+    ):
+        results = engine.dialect.has_table(connection, fact_table_name)
         assert results == 1
 
-    def test_get_table_names(self, engine: Engine):
-        results = engine.dialect.get_table_names(engine)
+    def test_get_table_names(self, engine: Engine, connection: Connection):
+        results = engine.dialect.get_table_names(connection)
         assert len(results) > 0
-        results = engine.dialect.get_table_names(engine, "public")
+        results = engine.dialect.get_table_names(connection, "public")
         assert len(results) > 0
-        results = engine.dialect.get_table_names(engine, "non_existing_schema")
+        results = engine.dialect.get_table_names(connection, "non_existing_schema")
         assert len(results) == 0
 
-    def test_get_columns(self, engine: Engine, fact_table_name: str):
-        results = engine.dialect.get_columns(engine, fact_table_name)
+    def test_get_columns(
+        self, engine: Engine, connection: Connection, fact_table_name: str
+    ):
+        results = engine.dialect.get_columns(connection, fact_table_name)
         assert len(results) > 0
         row = results[0]
         assert isinstance(row, dict)
@@ -113,5 +127,5 @@ def test_get_columns(self, engine: Engine, fact_table_name: str):
         assert row_keys[3] == "default"
 
     def test_service_account_connect(self, connection_service_account: Connection):
-        result = connection_service_account.execute("SELECT 1")
+        result = connection_service_account.execute(text("SELECT 1"))
         assert result.fetchall() == [(1,)]