feat: Turn off token cache param (#54)

Author: ptiurin

src/firebolt_db/firebolt_dialect.py

@@ -1,4 +1,5 @@
 import os
+from distutils.util import strtobool
 from types import ModuleType
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -112,6 +113,12 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]:
         parameters = dict(url.query)
         if "account_name" in parameters:
             kwargs["account_name"] = parameters.pop("account_name")
+        if "use_token_cache" in parameters:
+            # parameters are all passed as a string, we need to convert it
+            # to boolean for SDK compatibility
+            kwargs["use_token_cache"] = bool(
+                strtobool(parameters.pop("use_token_cache"))
+            )
         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:

tests/unit/test_firebolt_dialect.py

@@ -3,6 +3,7 @@
 
 import sqlalchemy
 from conftest import MockCursor, MockDBApi
+from pytest import mark
 from sqlalchemy.engine import url
 from sqlalchemy.sql import text
 
@@ -46,6 +47,7 @@ def test_create_connect_args(self, dialect: FireboltDialect):
         with mock.patch.dict(os.environ, {}, clear=True):
             result_list, result_dict = dialect.create_connect_args(u)
             assert "api_endpoint" not in result_dict
+            assert "use_token_cache" not in result_dict
             assert result_list == []
 
     def test_create_connect_args_set_params(self, dialect: FireboltDialect):
@@ -60,6 +62,24 @@ 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"}
 
+    @mark.parametrize(
+        "token,expected", [("false", False), ("0", False), ("true", True)]
+    )
+    def test_create_connect_args_token_cache(
+        self, token, expected, dialect: FireboltDialect
+    ):
+        connection_url = (
+            "test_engine://test_user@email:test_password@test_db_name/test_engine_name"
+            f"?use_token_cache={token}&param1=1&param2=2"
+        )
+        u = url.make_url(connection_url)
+        result_list, result_dict = dialect.create_connect_args(u)
+        assert (
+            "use_token_cache" in result_dict
+        ), "use_token_cache was not parsed correctly from connection string"
+        assert result_dict["use_token_cache"] == expected
+        assert dialect._set_parameters == {"param1": "1", "param2": "2"}
+
     def test_do_execute(
         self, dialect: FireboltDialect, cursor: mock.Mock(spec=MockCursor)
     ):