feat: Add Firebolt Core connection support (#106)

devin-ai-integration[bot] · ptiurin · web-flow · commit f708f518c9b5 · 2025-07-10T09:03:19.000+01:00
Co-authored-by: Devin AI &lt;158243242+devin-ai-integration[bot]@users.noreply.github.com&gt;
Co-authored-by: petro.tiurin@firebolt.io &lt;petro.tiurin@firebolt.io&gt;
diff --git a/.github/resources/core/config.json b/.github/resources/core/config.json
@@ -0,0 +1,7 @@
+{
+  "nodes": [
+    {
+      "host": "firebolt-core"
+    }
+  ]
+}
diff --git a/.github/resources/core/docker-compose.yml b/.github/resources/core/docker-compose.yml
@@ -0,0 +1,16 @@
+name: firebolt-core
+
+services:
+  firebolt-core:
+    image: ghcr.io/firebolt-db/firebolt-core:${IMAGE_TAG}
+    container_name: firebolt-core
+    command: --node 0
+    privileged: true
+    restart: no
+    ulimits:
+      memlock: 8589934592
+    ports:
+      - 3473:3473
+    volumes:
+      - ${BASE_DIR}/.github/resources/core/config.json:/firebolt-core/config.json:ro
+      - ${BASE_DIR}/firebolt-core:/firebolt-core/data
diff --git a/.github/workflows/integration-tests-core.yml b/.github/workflows/integration-tests-core.yml
@@ -0,0 +1,111 @@
+name: Core integration tests
+
+on:
+  workflow_dispatch:
+    inputs:
+      tag_version:
+        description: 'The docker image tag for the firebolt core'
+        required: false
+        type: string
+        default: 'preview-rc'
+      python_version:
+        description: 'Python version'
+        required: false
+        type: string
+        default: '3.8'
+  workflow_call:
+    inputs:
+      tag_version:
+        description: 'The docker image tag for the firebolt core'
+        required: false
+        type: string
+        default: 'preview-rc'
+      python_version:
+        description: 'Python version'
+        required: false
+        type: string
+        default: '3.8'
+
+jobs:
+  run-core-integration-tests:
+    runs-on: ubuntu-latest
+    env:
+      DOCKER_COMPOSE_FILE: ${{ github.workspace }}/.github/resources/core/docker-compose.yml
+      SERVICE_PORT: 3473
+      SERVICE_URL: http://localhost:3473
+      MAX_RETRIES: 30
+      RETRY_INTERVAL: 2
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ inputs.python_version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ".[dev]"
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Prepare docker-compose.yml
+        run: |
+          sed -i "s|\${IMAGE_TAG}|${{ inputs.tag_version }}|g" "$DOCKER_COMPOSE_FILE"
+          sed -i "s|\${BASE_DIR}|${{ github.workspace }}|g" "$DOCKER_COMPOSE_FILE"
+
+      - name: Start service container
+        run: |
+          docker compose -f "$DOCKER_COMPOSE_FILE" up -d
+          docker compose -f "$DOCKER_COMPOSE_FILE" ps
+
+      - name: Wait for service to be ready
+        run: |
+          for i in $(seq 1 $MAX_RETRIES); do
+            if curl --silent --fail "$SERVICE_URL" --data-binary "SELECT 1" | grep -q "1"; then
+              echo "Service is up and responding!"
+              exit 0
+            fi
+            echo "Waiting for service... ($i/$MAX_RETRIES)"
+            sleep $RETRY_INTERVAL
+          done
+          echo "Error: Service failed to start within timeout"
+          docker compose -f "$DOCKER_COMPOSE_FILE" logs
+          exit 1
+
+      - name: Run Core integration tests
+        env:
+          CORE_URL: "http://localhost:3473"
+        run: |
+          pytest -o log_cli=true -o log_cli_level=INFO tests/integration -k "core" --alluredir=allure-results
+
+      - name: Stop container
+        if: always()
+        run: |
+          docker compose -f "$DOCKER_COMPOSE_FILE" down
+
+      # Need to pull the pages branch in order to fetch the previous runs
+      - name: Get Allure history
+        uses: actions/checkout@v4
+        if: always()
+        continue-on-error: true
+        with:
+            ref: gh-pages
+            path: gh-pages
+
+      - name: Allure Report
+        uses: firebolt-db/action-allure-report@v1
+        if: always()
+        with:
+            github-key: ${{ secrets.GITHUB_TOKEN }}
+            test-type: integration
+            allure-dir: allure-results
+            pages-branch: gh-pages
+            repository-name: firebolt-sqlalchemy
diff --git a/.github/workflows/integration-tests-v1.yml b/.github/workflows/integration-tests-v1.yml
@@ -16,10 +16,10 @@ jobs:
       - name: Check out code
         uses: actions/checkout@v2
 
-      - name: Set up Python 3.8
+      - name: Set up Python 3.9
         uses: actions/setup-python@v5
         with:
-          python-version: 3.8
+          python-version: 3.9
 
       - name: Install dependencies
         run: |
@@ -56,7 +56,7 @@ jobs:
           FIREBOLT_BASE_URL: "api.staging.firebolt.io"
           ACCOUNT_NAME: "firebolt"
         run: |
-          pytest --last-failed -o log_cli=true -o log_cli_level=INFO tests/integration
+          pytest --last-failed -o log_cli=true -o log_cli_level=INFO tests/integration -k "not core"
 
       - name: Save failed tests
         id: cache-tests-save
diff --git a/.github/workflows/integration-tests-v2.yml b/.github/workflows/integration-tests-v2.yml
@@ -28,10 +28,10 @@ jobs:
         with:
           repository: firebolt-db/firebolt-sqlalchemy
 
-      - name: Set up Python 3.8
+      - name: Set up Python 3.9
         uses: actions/setup-python@v5
         with:
-          python-version: 3.8
+          python-version: 3.9
 
       - name: Install dependencies
         run: |
@@ -65,7 +65,7 @@ jobs:
           FIREBOLT_BASE_URL: "api.staging.firebolt.io"
           ACCOUNT_NAME: ${{ vars.FIREBOLT_ACCOUNT }}
         run: |
-          pytest --last-failed -o log_cli=true -o log_cli_level=INFO tests/integration --alluredir=allure-results
+          pytest --last-failed -o log_cli=true -o log_cli_level=INFO tests/integration -k "not core" --alluredir=allure-results
 
       - name: Save failed tests
         id: cache-tests-save
diff --git a/.github/workflows/python-integration-tests.yml b/.github/workflows/python-integration-tests.yml
@@ -19,3 +19,5 @@ jobs:
   integration-test-v2:
     uses: ./.github/workflows/integration-tests-v2.yml
     secrets: inherit
+  integration-test-core:
+    uses: ./.github/workflows/integration-tests-core.yml
diff --git a/pytest.ini b/pytest.ini
@@ -0,0 +1,4 @@
+[pytest]
+trio_mode = true
+markers =
+    core: mark test to run only on Firebolt Core.
diff --git a/setup.cfg b/setup.cfg
@@ -27,7 +27,7 @@ project_urls =
 [options]
 packages = find:
 install_requires =
-    firebolt-sdk>=1.5.0
+    firebolt-sdk>=1.13.0
     sqlalchemy>=1.0.0
 python_requires = >=3.8
 package_dir =
diff --git a/src/firebolt_db/firebolt_dialect.py b/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,36 +150,95 @@ 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)
-        # 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)
+        is_core_connection = "url" in parameters
+
+        if is_core_connection:
+            self._validate_core_connection(url, parameters)
+
+        token_cache_flag = self._parse_token_cache_flag(parameters)
+        auth = _determine_auth(url, token_cache_flag)
+        kwargs = self._build_connection_kwargs(
+            url, parameters, auth, is_core_connection
+        )
+
+        return ([], kwargs)
+
+    def _validate_core_connection(self, url: URL, parameters: Dict[str, str]) -> None:
+        """Validate that Core connection parameters are correct.
+
+        Only validates credentials since FireboltCore auth handles other parameters.
+        """
+        if url.username or url.password:
+            raise ArgumentError(
+                "Core connections do not support username/password authentication"
+            )
+
+    def _parse_token_cache_flag(self, parameters: Dict[str, str]) -> bool:
+        """Parse and remove token cache flag from parameters."""
+        return bool(strtobool(parameters.pop("use_token_cache", "True")))
+
+    def _build_connection_kwargs(
+        self, url: URL, parameters: Dict[str, str], auth: Auth, is_core_connection: bool
+    ) -> Dict[str, Union[str, Auth, Dict[str, Any], None]]:
+        """Build connection kwargs for the SDK.
+
+        SQLAlchemy URL mapping:
+        - url.host -> database (Firebolt database name)
+        - url.database -> engine_name (Firebolt engine name)
+        """
         kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]] = {
             "database": url.host or None,
             "auth": auth,
             "engine_name": url.database,
             "additional_parameters": {},
         }
-        additional_parameters = {}
+
+        if is_core_connection:
+            kwargs["url"] = parameters.pop("url")
+
+        self._handle_account_name(parameters, auth, kwargs)
+        self._handle_environment_config(kwargs)
+        kwargs["additional_parameters"] = self._build_additional_parameters(parameters)
+        self._set_parameters = parameters
+
+        return kwargs
+
+    def _handle_account_name(
+        self,
+        parameters: Dict[str, str],
+        auth: Auth,
+        kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]],
+    ) -> None:
+        """Handle account_name parameter and validation."""
         if "account_name" in parameters:
             kwargs["account_name"] = parameters.pop("account_name")
         elif isinstance(auth, ClientCredentials):
-            # account_name is required for client credentials authentication
             raise ArgumentError(
                 "account_name parameter must be provided to authenticate"
             )
-        self._set_parameters = parameters
-        # If URL override is not provided leave it to the sdk to determine the endpoint
+
+    def _handle_environment_config(
+        self, kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]]
+    ) -> None:
+        """Handle environment-based configuration."""
         if "FIREBOLT_BASE_URL" in os.environ:
             kwargs["api_endpoint"] = os.environ["FIREBOLT_BASE_URL"]
-        # Tracking information
+
+    def _build_additional_parameters(
+        self, parameters: Dict[str, str]
+    ) -> Dict[str, Any]:
+        """Build additional parameters including tracking information."""
+        additional_parameters: Dict[str, Any] = {}
+
         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)
+
+        return additional_parameters
 
     def get_schema_names(
         self, connection: AlchemyConnection, **kwargs: Any
@@ -366,8 +430,13 @@ 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:
-        return UsernamePassword(key, secret, token_cache_flag)
+def _determine_auth(url: URL, token_cache_flag: bool = True) -> Auth:
+    parameters = dict(url.query)
+    is_core_connection = "url" in parameters
+
+    if is_core_connection:
+        return FireboltCore()
+    elif "@" in (url.username or ""):
+        return UsernamePassword(url.username, url.password, token_cache_flag)
     else:
-        return ClientCredentials(key, secret, token_cache_flag)
+        return ClientCredentials(url.username, url.password, token_cache_flag)
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -190,7 +190,7 @@ def dimension_table_name() -> str:
     return "test_alchemy_dimension"
 
 
-@fixture(scope="class", autouse=True)
+@fixture(scope="class")
 def setup_test_tables(
     connection: Connection,
     engine: Engine,
@@ -233,3 +233,19 @@ def setup_test_tables(
     assert not engine.dialect.has_table(connection, fact_table_name)
     assert not engine.dialect.has_table(connection, dimension_table_name)
     assert not engine.dialect.has_table(connection, type_table_name)
+
+
+@fixture(scope="session")
+def core_url() -> str:
+    return environ.get("CORE_URL", "http://localhost:3473")
+
+
+@fixture(scope="session")
+def core_engine(core_url: str) -> Engine:
+    return create_engine(f"firebolt://firebolt?url={core_url}")
+
+
+@fixture(scope="session")
+def core_connection(core_engine: Engine) -> Connection:
+    with core_engine.connect() as c:
+        yield c
diff --git a/tests/integration/test_core_integration.py b/tests/integration/test_core_integration.py
diff --git a/tests/integration/test_sqlalchemy_async_integration.py b/tests/integration/test_sqlalchemy_async_integration.py
diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py
diff --git a/tests/unit/test_firebolt_dialect.py b/tests/unit/test_firebolt_dialect.py

-Original file line number
+Diff line change
@@ @@ -0,0 +1,7 @@ @@
 +{
 +  "nodes": [
 +    {
 +      "host": "firebolt-core"
 +    }
 +  ]
 +}