feat: Client tracking args

src/firebolt_db/firebolt_dialect.py

@@ -5,7 +5,7 @@
 
 import firebolt.db as dbapi
 import sqlalchemy.types as sqltypes
-from firebolt.client.auth import UsernamePassword
+from firebolt.client.auth import Auth, UsernamePassword
 from firebolt.db import Cursor
 from sqlalchemy.engine import Connection as AlchemyConnection
 from sqlalchemy.engine import ExecutionContext, default
@@ -111,17 +111,24 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]:
         # parameters are all passed as a string, we need to convert
         # bool flag to boolean for SDK compatibility
         token_cache_flag = bool(strtobool(parameters.pop("use_token_cache", "True")))
-        kwargs = {
+        kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]] = {
             "database": url.host or None,
             "auth": UsernamePassword(url.username, url.password, token_cache_flag),
             "engine_name": url.database,
+            "additional_parameters": {},
         }
+        additional_parameters = {}
         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"]
+        # Tracking information
+        if "user_clients" in parameters or "user_drivers" in parameters:
+            additional_parameters["user_drivers"] = parameters.pop("user_drivers", [])
+            additional_parameters["user_clients"] = parameters.pop("user_clients", [])
+        kwargs["additional_parameters"] = additional_parameters
         return ([], kwargs)
 
     def get_schema_names(

tests/integration/test_sqlalchemy_async_integration.py

@@ -51,23 +51,21 @@ async def test_data_write(self, async_connection: Connection, fact_table_name: s
             )
 
     @pytest.mark.asyncio
-    async def test_get_table_names(
-        self, async_connection: Connection, database_name: str
-    ):
+    async def test_get_table_names(self, async_connection: Connection):
         def get_table_names(conn: Connection) -> bool:
             inspector = inspect(conn)
-            return inspector.get_table_names(database_name)
+            return inspector.get_table_names()
 
         results = await async_connection.run_sync(get_table_names)
         assert len(results) > 0
 
     @pytest.mark.asyncio
     async def test_get_columns(
-        self, async_connection: Connection, database_name: str, fact_table_name: str
+        self, async_connection: Connection, fact_table_name: str
     ):
         def get_columns(conn: Connection) -> List[Dict]:
             inspector = inspect(conn)
-            return inspector.get_columns(fact_table_name, database_name)
+            return inspector.get_columns(fact_table_name)
 
         results = await async_connection.run_sync(get_columns)
         assert len(results) > 0

tests/integration/test_sqlalchemy_integration.py

@@ -51,18 +51,16 @@ def test_get_schema_names(self, engine: Engine, database_name: str):
         results = engine.dialect.get_schema_names(engine)
         assert database_name in results
 
-    def test_has_table(self, engine: Engine, database_name: str, fact_table_name: str):
-        results = engine.dialect.has_table(engine, fact_table_name, database_name)
+    def test_has_table(self, engine: Engine, fact_table_name: str):
+        results = engine.dialect.has_table(engine, fact_table_name)
         assert results == 1
 
-    def test_get_table_names(self, engine: Engine, database_name: str):
-        results = engine.dialect.get_table_names(engine, database_name)
+    def test_get_table_names(self, engine: Engine):
+        results = engine.dialect.get_table_names(engine)
         assert len(results) > 0
 
-    def test_get_columns(
-        self, engine: Engine, database_name: str, fact_table_name: str
-    ):
-        results = engine.dialect.get_columns(engine, fact_table_name, database_name)
+    def test_get_columns(self, engine: Engine, fact_table_name: str):
+        results = engine.dialect.get_columns(engine, fact_table_name)
         assert len(results) > 0
         row = results[0]
         assert isinstance(row, dict)

tests/unit/test_firebolt_dialect.py

@@ -65,6 +65,23 @@ def test_create_connect_args_set_params(self, dialect: FireboltDialect):
         ), "account_name was not parsed correctly from connection string"
         assert dialect._set_parameters == {"param1": "1", "param2": "2"}
 
+    def test_create_connect_args_driver_override(self, dialect: FireboltDialect):
+        connection_url = (
+            "test_engine://test_user@email:test_password@test_db_name/test_engine_name"
+            "?user_drivers=DriverA:1.0.2&user_clients=ClientB:2.0.9"
+        )
+        u = url.make_url(connection_url)
+        result_list, result_dict = dialect.create_connect_args(u)
+        assert (
+            "additional_parameters" in result_dict
+        ), "additional_parameters were not parsed correctly from connection string"
+        assert (
+            result_dict["additional_parameters"].get("user_drivers") == "DriverA:1.0.2"
+        )
+        assert (
+            result_dict["additional_parameters"].get("user_clients") == "ClientB:2.0.9"
+        )
+
     @mark.parametrize(
         "token,expected", [("false", False), ("0", False), ("true", True)]
     )