feat: Add Firebolt Core connection support

setup.cfg

@@ -27,7 +27,7 @@ project_urls =
 [options]
 packages = find:
 install_requires =
-    firebolt-sdk>=1.5.0
+    firebolt-sdk@git+https://github.com/firebolt-db/firebolt-python-sdk.git@core-support
     sqlalchemy>=1.0.0
 python_requires = >=3.8
 package_dir =

src/firebolt_db/firebolt_dialect.py

@@ -5,7 +5,12 @@
 
 import firebolt.db as dbapi
 import sqlalchemy.types as sqltypes
-from firebolt.client.auth import Auth, ClientCredentials, UsernamePassword
+from firebolt.client.auth import (
+    Auth,
+    ClientCredentials,
+    FireboltCore,
+    UsernamePassword,
+)
 from firebolt.db import Cursor
 from sqlalchemy.engine import Connection as AlchemyConnection
 from sqlalchemy.engine import ExecutionContext, default
@@ -145,18 +150,33 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]:
         """
         Build firebolt-sdk compatible connection arguments.
         URL format : firebolt://id:secret@host:port/db_name
+        For Core: firebolt://db_name?url=http://localhost:8080
+        (full URL including scheme, host, port in url parameter)
         """
         parameters = dict(url.query)
+
+        is_core_connection = "url" in parameters
+        core_url = parameters.pop("url", None) if is_core_connection else None
+
         # 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")))
-        auth = _determine_auth(url.username, url.password, token_cache_flag)
+
+        if is_core_connection:
+            auth = _determine_auth("", "", token_cache_flag, is_core=True)
+        else:
+            auth = _determine_auth(url.username, url.password, token_cache_flag)
+
         kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]] = {
             "database": url.host or None,
             "auth": auth,
             "engine_name": url.database,
             "additional_parameters": {},
         }
+
+        if core_url:
+            kwargs["url"] = core_url
+
         additional_parameters = {}
         if "account_name" in parameters:
             kwargs["account_name"] = parameters.pop("account_name")
@@ -366,8 +386,12 @@ def get_is_nullable(column_is_nullable: int) -> bool:
     return column_is_nullable == 1
 
 
-def _determine_auth(key: str, secret: str, token_cache_flag: bool = True) -> Auth:
-    if "@" in key:
+def _determine_auth(
+    key: str, secret: str, token_cache_flag: bool = True, is_core: bool = False
+) -> Auth:
+    if is_core:
+        return FireboltCore()
+    elif "@" in key:
         return UsernamePassword(key, secret, token_cache_flag)
     else:
         return ClientCredentials(key, secret, token_cache_flag)

tests/unit/test_firebolt_dialect.py

@@ -4,6 +4,7 @@
 import sqlalchemy
 import sqlalchemy.types as sqltypes
 from conftest import MockCursor, MockDBApi
+from firebolt.client.auth import FireboltCore
 from pytest import mark, raises
 from sqlalchemy.engine import url
 from sqlalchemy.exc import ArgumentError
@@ -302,6 +303,43 @@ def test_unicode_description(
     ):
         assert dialect._check_unicode_description(connection)
 
+    def test_create_connect_args_core_connection(self, dialect: FireboltDialect):
+        connection_url = (
+            "test_engine://test_db_name/test_engine_name?" "url=http://localhost:8080"
+        )
+        u = url.make_url(connection_url)
+        result_list, result_dict = dialect.create_connect_args(u)
+
+        assert result_dict["engine_name"] == "test_engine_name"
+        assert result_dict["database"] == "test_db_name"
+        assert result_dict["url"] == "http://localhost:8080"
+        assert isinstance(result_dict["auth"], FireboltCore)
+        assert "account_name" not in result_dict
+        assert result_list == []
+
+    def test_create_connect_args_core_connection_with_database(
+        self, dialect: FireboltDialect
+    ):
+        connection_url = "test_engine://test_db_name?" "url=http://localhost:8080"
+        u = url.make_url(connection_url)
+        result_list, result_dict = dialect.create_connect_args(u)
+
+        assert result_dict["database"] == "test_db_name"
+        assert result_dict["url"] == "http://localhost:8080"
+        assert isinstance(result_dict["auth"], FireboltCore)
+        assert result_dict["engine_name"] is None
+        assert result_list == []
+
+    def test_create_connect_args_core_no_credentials_required(
+        self, dialect: FireboltDialect
+    ):
+        connection_url = "test_engine://test_db_name?url=http://localhost:8080"
+        u = url.make_url(connection_url)
+
+        result_list, result_dict = dialect.create_connect_args(u)
+        assert isinstance(result_dict["auth"], FireboltCore)
+        assert "account_name" not in result_dict
+
 
 def test_get_is_nullable():
     assert firebolt_db.firebolt_dialect.get_is_nullable(1)