feat: Adding set parameter support (#46)

Author: ptiurin

src/firebolt_db/firebolt_async_dialect.py

@@ -8,8 +8,9 @@
 from firebolt.async_db import Connection
 
 # Ignoring type since sqlalchemy-stubs doesn't cover AdaptedConnection
+# and util.concurrency
 from sqlalchemy.engine import AdaptedConnection  # type: ignore[attr-defined]
-from sqlalchemy.util.concurrency import await_only
+from sqlalchemy.util.concurrency import await_only  # type: ignore[import]
 
 from firebolt_db.firebolt_dialect import FireboltDialect
 

src/firebolt_db/firebolt_dialect.py

@@ -4,6 +4,7 @@
 
 import firebolt.db as dbapi
 import sqlalchemy.types as sqltypes
+from firebolt.db import Cursor
 from sqlalchemy.engine import Connection as AlchemyConnection
 from sqlalchemy.engine import ExecutionContext, default
 from sqlalchemy.engine.url import URL
@@ -87,6 +88,7 @@ class FireboltDialect(default.DefaultDialect):
     returns_unicode_strings = True
     description_encoding = None
     supports_native_boolean = True
+    _set_parameters: Optional[Dict[str, Any]] = None
 
     def __init__(
         self, context: Optional[ExecutionContext] = None, *args: Any, **kwargs: Any
@@ -107,6 +109,10 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]:
             "password": url.password or None,
             "engine_name": url.database,
         }
+        parameters = dict(url.query)
+        if "account_name" in parameters:
+            kwargs["account_name"] = parameters.pop("account_name")
+        self._set_parameters = parameters
         # If URL override is not provided leave it to the sdk to determine the endpoint
         if "FIREBOLT_BASE_URL" in os.environ:
             kwargs["api_endpoint"] = os.environ["FIREBOLT_BASE_URL"]
@@ -257,6 +263,15 @@ def get_view_definition(
     ) -> str:
         pass
 
+    def do_execute(
+        self,
+        cursor: Cursor,
+        statement: str,
+        parameters: Tuple[str, Any],
+        context: Optional[ExecutionContext] = None,
+    ) -> None:
+        cursor.execute(statement, set_parameters=self._set_parameters)
+
     def do_rollback(self, dbapi_connection: AlchemyConnection) -> None:
         pass
 

tests/integration/test_sqlalchemy_integration.py

@@ -1,4 +1,5 @@
 import pytest
+from sqlalchemy import create_engine
 from sqlalchemy.engine.base import Connection, Engine
 from sqlalchemy.exc import OperationalError
 
@@ -17,6 +18,18 @@ def test_create_ex_table(
         connection.execute(f"DROP TABLE {ex_table_name}")
         assert not engine.dialect.has_table(engine, ex_table_name)
 
+    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"
+        )
+        with engine.connect() as connection:
+            result = connection.execute("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')"

tests/unit/conftest.py

@@ -37,6 +37,20 @@ def connect():
         pass
 
 
+class MockCursor:
+    def execute():
+        pass
+
+    def executemany():
+        pass
+
+    def fetchall():
+        pass
+
+    def close():
+        pass
+
+
 class MockAsyncDBApi:
     class DatabaseError:
         pass
@@ -100,6 +114,11 @@ def connection() -> mock.Mock(spec=MockDBApi):
     return mock.Mock(spec=MockDBApi)
 
 
+@fixture
+def cursor() -> mock.Mock(spec=MockCursor):
+    return mock.Mock(spec=MockCursor)
+
+
 @fixture
 def async_api() -> AsyncMock(spec=MockAsyncDBApi):
     return AsyncMock(spec=MockAsyncDBApi)

tests/unit/test_firebolt_dialect.py

@@ -2,7 +2,7 @@
 from unittest import mock
 
 import sqlalchemy
-from conftest import MockDBApi
+from conftest import MockCursor, MockDBApi
 from sqlalchemy.engine import url
 from sqlalchemy.sql import text
 
@@ -48,6 +48,25 @@ def test_create_connect_args(self, dialect: FireboltDialect):
             assert "api_endpoint" not in result_dict
             assert result_list == []
 
+    def test_create_connect_args_set_params(self, dialect: FireboltDialect):
+        connection_url = (
+            "test_engine://test_user@email:test_password@test_db_name/test_engine_name"
+            "?account_name=FB&param1=1&param2=2"
+        )
+        u = url.make_url(connection_url)
+        result_list, result_dict = dialect.create_connect_args(u)
+        assert (
+            "account_name" in result_dict
+        ), "account_name was not parsed correctly from connection string"
+        assert dialect._set_parameters == {"param1": "1", "param2": "2"}
+
+    def test_do_execute(
+        self, dialect: FireboltDialect, cursor: mock.Mock(spec=MockCursor)
+    ):
+        dialect._set_parameters = {"a": "b"}
+        dialect.do_execute(cursor, "SELECT *", None, None)
+        cursor.execute.assert_called_once_with("SELECT *", set_parameters={"a": "b"})
+
     def test_schema_names(
         self, dialect: FireboltDialect, connection: mock.Mock(spec=MockDBApi)
     ):