Add Firebolt Core connection support

- Add FireboltCore import from firebolt.client.auth
- Modify create_connect_args to detect 'url' parameter for Core connections
- Update _determine_auth to return FireboltCore auth when is_core=True
- Skip client_id/secret and account_name validation for Core connections
- Pass Core URL to SDK via kwargs['url']
- Add comprehensive unit tests for Core connection behavior
- Maintain backward compatibility with existing connection patterns

Co-Authored-By: [email protected] <[email protected]>

Author: devin-ai-integration[bot]

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 Auth, ClientCredentials, UsernamePassword
+from firebolt.client.auth import Auth, ClientCredentials, UsernamePassword, FireboltCore
 from firebolt.db import Cursor
 from sqlalchemy.engine import Connection as AlchemyConnection
 from sqlalchemy.engine import ExecutionContext, default
@@ -145,18 +145,32 @@ 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://host:port/db_name?url=core_url
         """
         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 +380,10 @@ 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

@@ -18,6 +18,7 @@
 )
 from firebolt_db.firebolt_dialect import dialect as dialect_definition
 from firebolt_db.firebolt_dialect import resolve_type
+from firebolt.client.auth import FireboltCore
 
 
 class TestFireboltDialect:
@@ -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)