feat: Allow set parameters as queries (#61)

Author: ptiurin

src/firebolt_db/firebolt_async_dialect.py

@@ -102,6 +102,14 @@ def fetchall(self) -> List[List]:
         self._rows[:] = []
         return retval
 
+    @property
+    def _set_parameters(self) -> Dict[str, Any]:
+        return self._cursor._set_parameters
+
+    @_set_parameters.setter
+    def _set_parameters(self, value: Dict[str, Any]) -> None:
+        self._cursor._set_parameters = value
+
 
 class AsyncConnectionWrapper(AdaptedConnection):
     await_ = staticmethod(await_only)

src/firebolt_db/firebolt_dialect.py

@@ -90,7 +90,7 @@ class FireboltDialect(default.DefaultDialect):
     returns_unicode_strings = True
     description_encoding = None
     supports_native_boolean = True
-    _set_parameters: Optional[Dict[str, Any]] = None
+    _set_parameters: Dict[str, Any] = dict()
 
     def __init__(
         self, context: Optional[ExecutionContext] = None, *args: Any, **kwargs: Any
@@ -283,9 +283,10 @@ def do_execute(
         parameters: Tuple[str, Any],
         context: Optional[ExecutionContext] = None,
     ) -> None:
-        cursor.execute(
-            statement, parameters=parameters, set_parameters=self._set_parameters
-        )
+        cursor._set_parameters = self._set_parameters
+        cursor.execute(statement, parameters=parameters)
+        # Persist set parameters across calls
+        self._set_parameters = cursor._set_parameters
 
     def do_rollback(self, dbapi_connection: AlchemyConnection) -> None:
         pass

tests/integration/test_sqlalchemy_async_integration.py

@@ -50,6 +50,17 @@ async def test_data_write(self, async_connection: Connection, fact_table_name: s
                 text(f"DELETE FROM {fact_table_name} WHERE idx=1")
             )
 
+    @pytest.mark.asyncio
+    async def test_set_params(self, async_connection: Engine):
+        await async_connection.execute(text(f"SET advanced_mode=1"))
+        await async_connection.execute(text(f"SET use_standard_sql=0"))
+        result = await async_connection.execute(
+            text(f"SELECT sleepEachRow(1) from numbers(1)")
+        )
+        assert len(result.fetchall()) == 1
+        await async_connection.execute(text(f"SET use_standard_sql=1"))
+        await async_connection.execute(text(f"SET advanced_mode=0"))
+
     @pytest.mark.asyncio
     async def test_get_table_names(self, async_connection: Connection):
         def get_table_names(conn: Connection) -> bool:

tests/integration/test_sqlalchemy_integration.py

@@ -22,10 +22,11 @@ def test_set_params(
         self, username: str, password: str, database_name: str, engine_name: str
     ):
         engine = create_engine(
-            f"firebolt://{username}:{password}@{database_name}/{engine_name}?"
-            "advanced_mode=1&use_standard_sql=0"
+            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)")
             assert len(result.fetchall()) == 1
         engine.dispose()

tests/unit/test_firebolt_dialect.py

@@ -104,14 +104,13 @@ def test_do_execute(
     ):
         dialect._set_parameters = {"a": "b"}
         dialect.do_execute(cursor, "SELECT *", None)
-        cursor.execute.assert_called_once_with(
-            "SELECT *", parameters=None, set_parameters={"a": "b"}
-        )
+        cursor.execute.assert_called_once_with("SELECT *", parameters=None)
+        assert cursor._set_parameters == {"a": "b"}, "Set parameters were not set"
+        cursor._set_parameters = {}
         cursor.execute.reset_mock()
         dialect.do_execute(cursor, "SELECT *", (1, 22), None)
-        cursor.execute.assert_called_once_with(
-            "SELECT *", parameters=(1, 22), set_parameters={"a": "b"}
-        )
+        cursor.execute.assert_called_once_with("SELECT *", parameters=(1, 22))
+        assert cursor._set_parameters == {"a": "b"}, "Set parameters were not set"
 
     def test_schema_names(
         self, dialect: FireboltDialect, connection: mock.Mock(spec=MockDBApi)