diff --git a/.github/resources/core/README.md b/.github/resources/core/README.md new file mode 100644 index 00000000000..c5e53379333 --- /dev/null +++ b/.github/resources/core/README.md @@ -0,0 +1,26 @@ +# Running tests for firebolt core over https + +In order to run the integration tests using https protocol, we need to bring up an nginx reverse proxy. On this proxy we will install the server certificate. In our case it will be a self signed cert (localhost.pem) and its private key (localhost-key.pem). + +The cert and its key will be stored in GitHub since we cannot have secrets stored in code repository, and it would be easier to rotate them (no code changes required): +- FIREBOLT_CORE_DEV_CERT is the certificate +- FIREBOLT_CORE_DEV_CERT_PRIVATE_KEY is the private key + +This is the certificate information that we currently have. Note that it wille expire in 2 years (when it will expire we will generate a new certificate and key and set them in the Git repository secrets) + +Certificate Information: +- Common Name: +- Subject Alternative Names: localhost +- Organization: mkcert development certificate +- Organization Unit: george@Mac.localdomain (George's Blue Mac) +- Locality: +- State: +- Country: +- Valid From: May 22, 2025 +- Valid To: August 22, 2027 +- Issuer: mkcert george@Mac.localdomain (George's Blue Mac), mkcert development CA +- Key Size: 2048 bit +- Serial Number: 5b19e46ed1a5a912da2cef0129924c41 + + +The client will connect over https to the nginx server which will do the TLS handshake and the TLS termination. It will then call the firebolt-core over http on port 3473. diff --git a/.github/resources/core/config.json b/.github/resources/core/config.json new file mode 100644 index 00000000000..d7fe906fd12 --- /dev/null +++ b/.github/resources/core/config.json @@ -0,0 +1,7 @@ +{ + "nodes": [ + { + "host": "firebolt-core" + } + ] +} \ No newline at end of file diff --git a/.github/resources/core/default.conf b/.github/resources/core/default.conf new file mode 100644 index 00000000000..81591b38732 --- /dev/null +++ b/.github/resources/core/default.conf @@ -0,0 +1,17 @@ +# this is the nginx configuration +# for http we are already exposing port 3473 on the docker so it will connect to the firebolt directly without going through nginx + +# HTTPS server +server { + listen 443 ssl; + server_name localhost; + + # On github these the localhost + ssl_certificate /etc/nginx/certs/server.pem; + ssl_certificate_key /etc/nginx/certs/server.key; + + location / { + proxy_pass http://firebolt-core:3473; + proxy_set_header Host $host; + } +} \ No newline at end of file diff --git a/.github/resources/core/docker-compose.yml b/.github/resources/core/docker-compose.yml new file mode 100644 index 00000000000..85092e6add4 --- /dev/null +++ b/.github/resources/core/docker-compose.yml @@ -0,0 +1,28 @@ +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 # equivalent to --rm (no persistence) + ulimits: + memlock: 8589934592 + ports: + - 3473:3473 + volumes: + # Mount the config file into the container. + - ${BASE_DIR}/.github/resources/core/config.json:/firebolt-core/config.json:ro + # Create an anonymous volume for Firebolt's internal database files. Not using a volume would have a performance impact. + - ${BASE_DIR}/firebolt-core:/firebolt-core/data + + nginx: + image: nginx:alpine + ports: + - "443:443" + volumes: + - ${BASE_DIR}/.github/resources/core/certs:/etc/nginx/certs:ro + - ${BASE_DIR}/.github/resources/core/default.conf:/etc/nginx/conf.d/default.conf:ro + depends_on: + - firebolt-core \ No newline at end of file diff --git a/.github/workflows/integration-tests-core.yml b/.github/workflows/integration-tests-core.yml new file mode 100644 index 00000000000..894ef659d18 --- /dev/null +++ b/.github/workflows/integration-tests-core.yml @@ -0,0 +1,199 @@ +name: Core integration tests + +on: + workflow_dispatch: + inputs: + tag_version: + description: 'The docker image tag for the firebolt core' + required: false + type: string + python_version: + description: 'Python version' + required: false + type: string + default: '3.10' + os_name: + description: 'The operating system' + required: false + type: string + default: 'ubuntu-latest' + workflow_call: + inputs: + tag_version: + description: 'The docker image tag for the firebolt core' + required: false + type: string + python_version: + description: 'Python version' + required: false + type: string + default: '3.10' + os_name: + description: 'Operating system' + required: false + type: string + default: 'ubuntu-latest' + sendSlackNotifications: + description: 'Send Slack notifications on failure' + required: false + type: boolean + default: false + secrets: + SLACK_BOT_TOKEN: + required: false +env: + DEFAULT_IMAGE_TAG: ${{ vars.DEFAULT_CORE_IMAGE_TAG }} +jobs: + run-core-integration-tests: + runs-on: ${{ inputs.os_name }} + 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 + with: + repository: 'firebolt-db/firebolt-python-sdk' + + - 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: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Cache Docker layers + uses: actions/cache@v3 + with: + path: /tmp/.buildx-cache + key: ${{ runner.os }}-buildx-${{ github.sha }} + restore-keys: | + ${{ runner.os }}-buildx- + + - name: Write certificate and certificate key to file + run: | + mkdir "${{ github.workspace }}/.github/resources/core/certs" + pip install trustme + # Generate a self-signed certificate for localhost + python3 -m trustme -d "${{ github.workspace }}/.github/resources/core/certs/" + + - name: Install certs to keystore + run: | + sudo cp ${GITHUB_WORKSPACE}/.github/resources/core/certs/client.pem /usr/local/share/ca-certificates/client.crt + sudo update-ca-certificates + + # if no image tag was passed in, then use the image tag from the defaults + - name: Set image tag + id: set-tag + run: | + IMAGE_TAG="${{ inputs.tag_version }}" + if [ -z "$IMAGE_TAG" ]; then + IMAGE_TAG="$DEFAULT_IMAGE_TAG" + fi + echo "tag=$IMAGE_TAG" >> $GITHUB_OUTPUT + + - name: Prepare docker-compose.yml + run: | + if [ ! -f "$DOCKER_COMPOSE_FILE" ]; then + echo "Error: Docker compose file not found at $DOCKER_COMPOSE_FILE" + exit 1 + fi + sed -i "s|\${IMAGE_TAG}|${{ steps.set-tag.outputs.tag }}|g" "$DOCKER_COMPOSE_FILE" + sed -i "s|\${BASE_DIR}|${{ github.workspace }}|g" "$DOCKER_COMPOSE_FILE" + echo "Docker compose file prepared:" + cat "$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 integration tests HTTP + env: + SERVICE_ID: ${{ secrets.FIREBOLT_CLIENT_ID_STG_NEW_IDN }} + SERVICE_SECRET: ${{ secrets.FIREBOLT_CLIENT_SECRET_STG_NEW_IDN }} + DATABASE_NAME: "firebolt" + ENGINE_NAME: "" + STOPPED_ENGINE_NAME: "" + API_ENDPOINT: "" + ACCOUNT_NAME: "" + CORE_URL: "http://localhost:3473" + run: | + pytest -o log_cli=true -o log_cli_level=WARNING tests/integration -k "core" --alluredir=allure-results/ + + - name: Run integration tests HTTPS + env: + SERVICE_ID: ${{ secrets.FIREBOLT_CLIENT_ID_STG_NEW_IDN }} + SERVICE_SECRET: ${{ secrets.FIREBOLT_CLIENT_SECRET_STG_NEW_IDN }} + DATABASE_NAME: "firebolt" + ENGINE_NAME: "" + STOPPED_ENGINE_NAME: "" + API_ENDPOINT: "" + ACCOUNT_NAME: "" + CORE_URL: "https://localhost:443" + run: | + pytest -o log_cli=true -o log_cli_level=WARNING tests/integration -k "core" --alluredir=allure-results-https/ + + - 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: python-sdk + + - name: Allure Report HTTPS + uses: firebolt-db/action-allure-report@v1 + if: always() + with: + github-key: ${{ secrets.GITHUB_TOKEN }} + test-type: integration_https + allure-dir: allure-results-https + pages-branch: gh-pages + repository-name: python-sdk \ No newline at end of file diff --git a/.github/workflows/integration-tests-v1.yml b/.github/workflows/integration-tests-v1.yml index 908f52385b1..0b8f131c483 100644 --- a/.github/workflows/integration-tests-v1.yml +++ b/.github/workflows/integration-tests-v1.yml @@ -59,7 +59,7 @@ jobs: API_ENDPOINT: "api.staging.firebolt.io" ACCOUNT_NAME: "firebolt" run: | - pytest --last-failed -n 6 --dist loadgroup --timeout_method "signal" -o log_cli=true -o log_cli_level=INFO tests/integration -k "not V2" --runslow + pytest --last-failed -n 6 --dist loadgroup --timeout_method "signal" -o log_cli=true -o log_cli_level=INFO tests/integration -k "not V2 and not core" --runslow - name: Save failed tests id: cache-tests-save diff --git a/.github/workflows/integration-tests-v2.yml b/.github/workflows/integration-tests-v2.yml index 0e62e582844..df881134a55 100644 --- a/.github/workflows/integration-tests-v2.yml +++ b/.github/workflows/integration-tests-v2.yml @@ -57,7 +57,7 @@ jobs: API_ENDPOINT: "api.staging.firebolt.io" ACCOUNT_NAME: ${{ vars.FIREBOLT_ACCOUNT }} run: | - pytest -n 6 --dist loadgroup --timeout_method "signal" -o log_cli=true -o log_cli_level=WARNING tests/integration -k "not V1" --runslow --alluredir=allure-results + pytest -n 6 --dist loadgroup --timeout_method "signal" -o log_cli=true -o log_cli_level=WARNING tests/integration -k "not V1 and not core" --runslow --alluredir=allure-results # Need to pull the pages branch in order to fetch the previous runs - name: Get Allure history diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index cb0a63dbab9..839d112deec 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -23,8 +23,10 @@ jobs: integration-test-v2: uses: ./.github/workflows/integration-tests-v2.yml secrets: inherit + integration-test-core: + uses: ./.github/workflows/integration-tests-core.yml report-test-results: - needs: [integration-test-v1, integration-test-v2] + needs: [integration-test-v1, integration-test-v2, integration-test-core] if: always() runs-on: ubuntu-latest steps: diff --git a/docsrc/Connecting_and_queries.rst b/docsrc/Connecting_and_queries.rst index b4e55c2fa21..5fb815ed1e6 100644 --- a/docsrc/Connecting_and_queries.rst +++ b/docsrc/Connecting_and_queries.rst @@ -1,4 +1,3 @@ - ############################### Connecting and running queries ############################### @@ -7,8 +6,8 @@ This topic provides a walkthrough and examples for how to use the Firebolt Pytho connect to Firebolt resources to run commands and query data. -Setting up a connection -========================= +Setting up a connection (Cloud) +=============================== To connect to a Firebolt database to run queries or command, you must provide your account credentials through a connection request. @@ -149,6 +148,94 @@ To get started, follow the steps below: The ``cursor`` object can be used to send queries and commands to your Firebolt database and engine. See below for examples of functions using the ``cursor`` object. +Setting up a connection (Core) +=============================== + +Firebolt Core is a Docker-based version of Firebolt that can be run locally or remotely. To connect to Firebolt Core, you need to use the ``FireboltCore`` authentication class, which doesn't require actual credentials. + +To get started, follow the steps below: + +**1. Import modules** + + The Firebolt Python SDK requires you to import the following modules before making + any command or query requests to your Firebolt Core instance. + +.. _required_core_connection_imports: + + :: + + from firebolt.db import connect + from firebolt.client.auth import FireboltCore + +**2. Connect to your database** + + To connect to Firebolt Core, you need to create a ``FireboltCore`` auth object and use it + to establish a connection. + + A connection requires the following parameters: + + +------------------------------------+---------------------------------------------------------------------------------------------------------------+ + | ``auth`` | Auth object of type FireboltCore. This is a special authentication type that doesn't require credentials. | + +------------------------------------+---------------------------------------------------------------------------------------------------------------+ + | ``database`` | Optional. The name of the database you would like to connect to. Defaults to "firebolt" if not specified. | + +------------------------------------+---------------------------------------------------------------------------------------------------------------+ + | ``url`` | Optional. The URL of the Firebolt Core instance of form ://:. | + | | Defaults to "http://localhost:3473" if not specified. | + +------------------------------------+---------------------------------------------------------------------------------------------------------------+ + + Here's how to create a connection to Firebolt Core: + + :: + + # Connect to Firebolt Core + # The database parameter defaults to 'firebolt' if not specified + with connect( + auth=FireboltCore(), + url="http://localhost:3473", + database="firebolt" + ) as connection: + # Create a cursor + cursor = connection.cursor() + + # Execute a simple test query + cursor.execute("SELECT 1") + +.. note:: + + Firebolt Core is assumed to be running locally on the default port (3473). For instructions + on how to run Firebolt Core locally using Docker, refer to the + `official docs `_. + + +**2.1. Connecting to an HTTPS Server** + + If you are connecting to an HTTPS server running Firebolt Core, ensure the server's certificate is properly configured. Follow these steps: + + - **Obtain the Certificate**: Ensure you have the certificate for the HTTPS server you are connecting to. + + - **Install the Certificate in the System Certificate Store**: For Ubuntu users, copy the certificate to `/usr/local/share/ca-certificates/` and run the following command: + + ```bash + sudo update-ca-certificates + ``` + + .. note:: + Installing the certificate is also possible on other systems, but you need to use the correct path for your operating system's certificate store. + + - **Provide the Certificate via Environment Variable**: Alternatively, set the `SSL_CERT_FILE` environment variable to the path of your certificate file. For example: + + ```bash + export SSL_CERT_FILE=/path/to/your/certificate.pem + ``` + + - **Python Version Considerations**: The system certificate store is only available for users running Python 3.10 and above. If you are using an older version of Python, you must explicitly set the `SSL_CERT_FILE` environment variable to use the certificate. + + +**3. Execute commands using the cursor** + + The ``cursor`` object can be used to send queries and commands to your Firebolt + database and engine. See below for examples of functions using the ``cursor`` object. + Synchronous command and query examples ================================================== @@ -769,4 +856,3 @@ function will raise a ``QueryTimeoutError`` exception. ) **Warning**: If running multiple queries, and one of queries times out, all the previous queries will not be rolled back and their result will persist. All the remaining queries will be cancelled. - diff --git a/examples/firebolt_core.ipynb b/examples/firebolt_core.ipynb new file mode 100644 index 00000000000..bb0eaa7e02b --- /dev/null +++ b/examples/firebolt_core.ipynb @@ -0,0 +1,333 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "11fe8ceb", + "metadata": {}, + "source": [ + "# Firebolt Core Examples\n", + "\n", + "This notebook demonstrates using the Firebolt Python SDK with Firebolt Core, a Docker-based version of Firebolt for local or remote use." + ] + }, + { + "cell_type": "markdown", + "id": "a0ff488e", + "metadata": {}, + "source": [ + "## Prerequisites\n", + "\n", + "Firebolt Core can be run locally using Docker. Refer to the [official docs](https://docs.firebolt.io/firebolt-core/firebolt-core-get-started) on how to run it.\n", + "\n", + "For this notebook, we'll assume Firebolt Core is running locally on the default port (3473)." + ] + }, + { + "cell_type": "markdown", + "id": "8caa8e5a", + "metadata": {}, + "source": [ + "## Connection Types\n", + "\n", + "There are two ways to connect to Firebolt Core:\n", + "\n", + "1. **Synchronous connection** - using the standard DB-API 2.0 interface\n", + "2. **Asynchronous connection** - using the async equivalent API\n", + "\n", + "Both methods use the `FireboltCore` authentication class, which doesn't require actual credentials." + ] + }, + { + "cell_type": "markdown", + "id": "14314b52", + "metadata": {}, + "source": [ + "### Required imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "159d4236", + "metadata": {}, + "outputs": [], + "source": [ + "from firebolt.db import connect\n", + "from firebolt.client.auth import FireboltCore" + ] + }, + { + "cell_type": "markdown", + "id": "57eaa46b", + "metadata": {}, + "source": [ + "### Connection settings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc50522c", + "metadata": {}, + "outputs": [], + "source": [ + "# For Firebolt Core, we don't need credentials\n", + "# but we can specify the database and host if needed\n", + "database = \"firebolt\" # Default database name\n", + "url = \"http://localhost:3473\" # Default URL" + ] + }, + { + "cell_type": "markdown", + "id": "a2159041", + "metadata": {}, + "source": [ + "### Connecting to a database and creating cursor" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7debd869", + "metadata": {}, + "outputs": [], + "source": [ + "# Create a FireboltCore auth object - no credentials needed\n", + "auth = FireboltCore()\n", + "\n", + "# Connect to Firebolt Core\n", + "connection = connect(auth=auth, database=database, url=url)\n", + "\n", + "# Create a cursor\n", + "cursor = connection.cursor()" + ] + }, + { + "cell_type": "markdown", + "id": "5a498df4", + "metadata": {}, + "source": [ + "### Executing a query" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "769ff379", + "metadata": {}, + "outputs": [], + "source": [ + "# Create a test table\n", + "cursor.execute(\n", + " \"CREATE TABLE IF NOT EXISTS example_table (id INT, name TEXT, value FLOAT, created_at TIMESTAMP)\"\n", + ")\n", + "\n", + "# Insert some test data\n", + "cursor.execute(\n", + " \"INSERT INTO example_table VALUES \"\n", + " \"(1, 'Item 1', 10.5, '2023-01-01 12:00:00'), \"\n", + " \"(2, 'Item 2', 20.75, '2023-01-02 14:30:00'), \"\n", + " \"(3, 'Item 3', 15.25, '2023-01-03 09:45:00')\"\n", + ")\n", + "\n", + "# Execute a simple test query\n", + "cursor.execute(\"SELECT 1\")\n", + "\n", + "# Fetch and display results\n", + "result = cursor.fetchall()\n", + "print(f\"Query result: {result}\")\n", + "\n", + "# Get column names\n", + "print(f\"Column names: {cursor.description[0][0]}\")\n", + "\n", + "# Show connection parameters (filtered for Firebolt Core)\n", + "print(f\"Connection URL: {connection.engine_url}\")" + ] + }, + { + "cell_type": "markdown", + "id": "5439c170", + "metadata": {}, + "source": [ + "### Parameterized query" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98377601", + "metadata": {}, + "outputs": [], + "source": [ + "# Execute with a single parameter set\n", + "cursor.execute(\n", + " \"INSERT INTO example_table VALUES (?, ?, ?, ?)\",\n", + " (4, \"Parameter Example\", 30.5, \"2023-01-04 10:00:00\"),\n", + ")\n", + "\n", + "# Execute with multiple parameter sets\n", + "cursor.executemany(\n", + " \"INSERT INTO example_table VALUES (?, ?, ?, ?)\",\n", + " [\n", + " (5, \"Multi Param 1\", 25.5, \"2023-01-05 11:00:00\"),\n", + " (6, \"Multi Param 2\", 35.75, \"2023-01-06 12:00:00\"),\n", + " ],\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "7ef390a9", + "metadata": {}, + "source": [ + "### Getting query description, rowcount" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ae2366ff", + "metadata": {}, + "outputs": [], + "source": [ + "cursor.execute(\"SELECT * FROM example_table\")\n", + "print(\"Description: \", cursor.description)\n", + "print(\"Rowcount: \", cursor.rowcount)" + ] + }, + { + "cell_type": "markdown", + "id": "355028a9", + "metadata": {}, + "source": [ + "### Fetch query results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a992d4aa", + "metadata": {}, + "outputs": [], + "source": [ + "# Re-run the query to reset the cursor position\n", + "cursor.execute(\"SELECT * FROM example_table\")\n", + "\n", + "# Different fetch methods\n", + "print(\"fetchone():\", cursor.fetchone())\n", + "print(\"fetchmany(2):\", cursor.fetchmany(2))\n", + "print(\"fetchall():\", cursor.fetchall())" + ] + }, + { + "cell_type": "markdown", + "id": "f73c88b3", + "metadata": {}, + "source": [ + "### Closing the connection" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7871d2fa", + "metadata": {}, + "outputs": [], + "source": [ + "connection.close()" + ] + }, + { + "cell_type": "markdown", + "id": "e1a0c339", + "metadata": {}, + "source": [ + "### Async interface\n", + "**NOTE**: In order to make async examples work in jupyter, you would need to install [trio-jupyter](https://github.com/mehaase/trio-jupyter) library and select **Python 3 Trio** kernel" + ] + }, + { + "cell_type": "markdown", + "id": "46abd8b7", + "metadata": {}, + "source": [ + "#### Required imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dc86010e", + "metadata": {}, + "outputs": [], + "source": [ + "from firebolt.async_db import connect as async_connect\n", + "from firebolt.client.auth import FireboltCore" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2173ff49", + "metadata": {}, + "outputs": [], + "source": [ + "auth = FireboltCore()\n", + "# Connect to Firebolt Core asynchronously\n", + "connection = await async_connect(auth=auth, database=\"firebolt\")\n", + "\n", + "# Create a cursor\n", + "cursor = connection.cursor()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4e77f75f", + "metadata": {}, + "outputs": [], + "source": [ + "await cursor.execute(\"SELECT 2\")\n", + "\n", + "# Fetch and display results\n", + "result = await cursor.fetchall()\n", + "print(f\"Async query result: {result}\")\n", + "\n", + "# Get column names\n", + "print(f\"Column names: {cursor.description[0][0]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2788ad42", + "metadata": {}, + "outputs": [], + "source": [ + "# Close connection when done\n", + "await connection.aclose()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/setup.cfg b/setup.cfg index 04140ea41fd..d2c5a975853 100755 --- a/setup.cfg +++ b/setup.cfg @@ -39,6 +39,7 @@ install_requires = readerwriterlock>=1.0.9 sqlparse>=0.4.2 trio>=0.22.0 + truststore>=0.10;python_version>="3.10" python_requires = >=3.8 include_package_data = True package_dir = diff --git a/src/firebolt/async_db/connection.py b/src/firebolt/async_db/connection.py index 62f6f215330..202aa7530ea 100644 --- a/src/firebolt/async_db/connection.py +++ b/src/firebolt/async_db/connection.py @@ -1,7 +1,8 @@ from __future__ import annotations +from ssl import SSLContext from types import TracebackType -from typing import Any, Dict, List, Optional, Type +from typing import Any, Dict, List, Optional, Type, Union from httpx import Timeout @@ -9,6 +10,7 @@ from firebolt.async_db.util import _get_system_engine_url_and_params from firebolt.client import DEFAULT_API_URL from firebolt.client.auth import Auth +from firebolt.client.auth.base import FireboltAuthVersion from firebolt.client.client import AsyncClient, AsyncClientV1, AsyncClientV2 from firebolt.common.base_connection import ( ASYNC_QUERY_CANCEL, @@ -26,6 +28,11 @@ ConnectionClosedError, FireboltError, ) +from firebolt.utils.firebolt_core import ( + get_core_certificate_context, + parse_firebolt_core_url, + validate_firebolt_core_parameters, +) from firebolt.utils.usage_tracker import get_user_agent_header from firebolt.utils.util import fix_url_schema, validate_engine_name_and_url_v1 @@ -209,6 +216,7 @@ async def connect( engine_url: Optional[str] = None, api_endpoint: str = DEFAULT_API_URL, disable_cache: bool = False, + url: Optional[str] = None, additional_parameters: Dict[str, Any] = {}, ) -> Connection: # auth parameter is optional in function signature @@ -224,10 +232,20 @@ async def connect( user_agent_header = get_user_agent_header(user_drivers, user_clients) if disable_cache: _firebolt_system_engine_cache.disable() - # Use v2 if auth is ClientCredentials - # Use v1 if auth is ServiceAccount or UsernamePassword + # Use CORE if auth is FireboltCore + # Use V2 if auth is ClientCredentials + # Use V1 if auth is ServiceAccount or UsernamePassword auth_version = auth.get_firebolt_version() - if auth_version == 2: + if auth_version == FireboltAuthVersion.CORE: + # Verify that Core-incompatible parameters are not provided + validate_firebolt_core_parameters(account_name, engine_name, engine_url) + return connect_core( + auth=auth, + user_agent_header=user_agent_header, + database=database, + connection_url=url, + ) + elif auth_version == FireboltAuthVersion.V2: assert account_name is not None return await connect_v2( auth=auth, @@ -237,7 +255,7 @@ async def connect( engine_name=engine_name, api_endpoint=api_endpoint, ) - elif auth_version == 1: + elif auth_version == FireboltAuthVersion.V1: return await connect_v1( auth=auth, user_agent_header=user_agent_header, @@ -379,3 +397,48 @@ async def connect_v1( headers={"User-Agent": user_agent_header}, ) return Connection(engine_url, database, client, CursorV1, api_endpoint) + + +def connect_core( + auth: Auth, + user_agent_header: str, + database: Optional[str] = None, + connection_url: Optional[str] = None, +) -> Connection: + """Connect to Firebolt Core. + + Args: + auth (Auth): Authentication object (must be FireboltCore) + user_agent_header (str): User agent header string + database (Optional[str]): Name of the database to connect to + (defaults to 'firebolt') + connection_url (Optional[str]): URL in format protocol://host:port + Protocol defaults to http, host defaults to localhost, port + defaults to 3473. + + Returns: + Connection: A connection to Firebolt Core + """ + connection_params = parse_firebolt_core_url(connection_url) + + ctx: Union[SSLContext, bool] = True # Default context + if connection_params.scheme == "https": + ctx = get_core_certificate_context() + + verified_url = connection_params.geturl() + client = AsyncClientV2( + auth=auth, + account_name="", # FireboltCore does not require an account name + base_url=verified_url, + timeout=Timeout(DEFAULT_TIMEOUT_SECONDS, read=None), + headers={"User-Agent": user_agent_header}, + verify=ctx, + ) + + return Connection( + engine_url=verified_url, + database=database, + client=client, + cursor_type=CursorV2, + api_endpoint=verified_url, + ) diff --git a/src/firebolt/client/auth/__init__.py b/src/firebolt/client/auth/__init__.py index 0caa8b71e37..512efacb2e7 100644 --- a/src/firebolt/client/auth/__init__.py +++ b/src/firebolt/client/auth/__init__.py @@ -1,5 +1,6 @@ from firebolt.client.auth.base import Auth from firebolt.client.auth.client_credentials import ClientCredentials +from firebolt.client.auth.firebolt_core import FireboltCore from firebolt.client.auth.service_account import ServiceAccount from firebolt.client.auth.token import Token from firebolt.client.auth.username_password import UsernamePassword diff --git a/src/firebolt/client/auth/base.py b/src/firebolt/client/auth/base.py index a2adfb1d063..74df42202db 100644 --- a/src/firebolt/client/auth/base.py +++ b/src/firebolt/client/auth/base.py @@ -1,4 +1,5 @@ from abc import abstractmethod +from enum import IntEnum from time import time from typing import AsyncGenerator, Generator, Optional @@ -10,6 +11,14 @@ from firebolt.utils.util import Timer, cached_property, get_internal_error_code +class FireboltAuthVersion(IntEnum): + """Enum for Firebolt authentication versions.""" + + V1 = 1 # Service Account, Username Password + V2 = 2 # Client Credentials + CORE = 3 # Firebolt Core + + class AuthRequest(Request): """Class to distinguish auth requests from regular""" @@ -56,13 +65,13 @@ def token(self) -> Optional[str]: """ return self._token - def get_firebolt_version(self) -> int: + @abstractmethod + def get_firebolt_version(self) -> FireboltAuthVersion: """Get Firebolt version from auth. Returns: - int: Firebolt version + FireboltAuthVersion: The authentication version enum """ - return -1 @property def expired(self) -> bool: diff --git a/src/firebolt/client/auth/client_credentials.py b/src/firebolt/client/auth/client_credentials.py index e8f86bfa1dc..92181fb9188 100644 --- a/src/firebolt/client/auth/client_credentials.py +++ b/src/firebolt/client/auth/client_credentials.py @@ -1,6 +1,6 @@ from typing import Optional -from firebolt.client.auth.base import AuthRequest +from firebolt.client.auth.base import AuthRequest, FireboltAuthVersion from firebolt.client.auth.request_auth_base import _RequestBasedAuth from firebolt.utils.token_storage import TokenSecureStorage from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL @@ -53,13 +53,13 @@ def copy(self) -> "ClientCredentials": self.client_id, self.client_secret, self._use_token_cache ) - def get_firebolt_version(self) -> int: + def get_firebolt_version(self) -> FireboltAuthVersion: """Get Firebolt version from auth. Returns: - int: Firebolt version + FireboltAuthVersion: V2 for Client Credentials authentication """ - return 2 + return FireboltAuthVersion.V2 @cached_property def _token_storage(self) -> Optional[TokenSecureStorage]: diff --git a/src/firebolt/client/auth/firebolt_core.py b/src/firebolt/client/auth/firebolt_core.py new file mode 100644 index 00000000000..8cffe64aa4f --- /dev/null +++ b/src/firebolt/client/auth/firebolt_core.py @@ -0,0 +1,96 @@ +from typing import AsyncGenerator, Generator + +from httpx import Request, Response + +from firebolt.client.auth.base import Auth, FireboltAuthVersion + + +class FireboltCore(Auth): + """Authentication class for Firebolt Core. + + Represents authentication for local/remote Docker-based deployments of Firebolt, + which do not require authentication. + """ + + __slots__ = ( + "_token", + "_expires", + "_use_token_cache", + ) + + def __init__(self) -> None: + # Initialize with no token caching + super().__init__(use_token_cache=False) + + # FireboltCore doesn't need a token, but we provide an empty one + # to satisfy the Auth interface requirements + self._token = "" + self._expires = None + + def copy(self) -> "FireboltCore": + """Make another auth object with same URL. + + Returns: + FireboltCore: Auth object + """ + return FireboltCore() + + def get_firebolt_version(self) -> FireboltAuthVersion: + """Get Firebolt version from auth. + + Returns: + FireboltAuthVersion: CORE for Firebolt Core authentication + """ + return FireboltAuthVersion.CORE + + def get_new_token_generator(self) -> Generator: + """FireboltCore doesn't need token authentication. + + Yields: + No requests are yielded for FireboltCore + + Raises: + No exceptions are raised + """ + yield from [] # No requests to yield, as no authentication is needed + + @property + def token(self) -> str: + """Get token for Firebolt Core. + + For FireboltCore, this returns an empty string since no auth is needed. + + Returns: + str: Empty string (no token needed) + """ + return "" + + def auth_flow(self, request: Request) -> Generator[Request, Response, None]: + """Override auth flow for Firebolt Core to avoid sending auth headers. + + This implementation ensures no Authorization headers are sent. + + Args: + request: The request to authenticate + + Yields: + The request without authentication headers + """ + # Yield the request without authentication + yield request + + async def async_auth_flow( + self, request: Request + ) -> AsyncGenerator[Request, Response]: + """Override async auth flow for Firebolt Core to avoid sending auth headers. + + This implementation ensures no Authorization headers are sent. + + Args: + request: The request to authenticate + + Yields: + The request without authentication headers + """ + # Yield the request without authentication + yield request diff --git a/src/firebolt/client/auth/service_account.py b/src/firebolt/client/auth/service_account.py index c0a85569774..6d670c00d0a 100644 --- a/src/firebolt/client/auth/service_account.py +++ b/src/firebolt/client/auth/service_account.py @@ -1,6 +1,6 @@ from typing import Optional -from firebolt.client.auth.base import AuthRequest +from firebolt.client.auth.base import AuthRequest, FireboltAuthVersion from firebolt.client.auth.request_auth_base import _RequestBasedAuth from firebolt.utils.token_storage import TokenSecureStorage from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL @@ -43,13 +43,13 @@ def __init__( self.client_secret = client_secret super().__init__(use_token_cache) - def get_firebolt_version(self) -> int: + def get_firebolt_version(self) -> FireboltAuthVersion: """Get Firebolt version from auth. Returns: - int: Firebolt version + FireboltAuthVersion: V1 for Service Account authentication """ - return 1 + return FireboltAuthVersion.V1 def copy(self) -> "ServiceAccount": """Make another auth object with same credentials. diff --git a/src/firebolt/client/auth/token.py b/src/firebolt/client/auth/token.py index 73761fdeec5..4ad7bd13ce3 100644 --- a/src/firebolt/client/auth/token.py +++ b/src/firebolt/client/auth/token.py @@ -3,6 +3,7 @@ from httpx import Request, Response from firebolt.client.auth import Auth +from firebolt.client.auth.base import FireboltAuthVersion from firebolt.utils.exception import AuthorizationError @@ -23,13 +24,13 @@ def __init__(self, token: str): super().__init__(use_token_cache=False) self._token = token - def get_firebolt_version(self) -> int: + def get_firebolt_version(self) -> FireboltAuthVersion: """Get Firebolt version from auth. Returns: - int: Firebolt version + FireboltAuthVersion: V1 for Token authentication """ - return 1 + return FireboltAuthVersion.V1 def copy(self) -> "Token": """Make another auth object with same credentials. diff --git a/src/firebolt/client/auth/username_password.py b/src/firebolt/client/auth/username_password.py index 3ec2c644e5d..79f1548294a 100644 --- a/src/firebolt/client/auth/username_password.py +++ b/src/firebolt/client/auth/username_password.py @@ -1,6 +1,6 @@ from typing import Optional -from firebolt.client.auth.base import AuthRequest +from firebolt.client.auth.base import AuthRequest, FireboltAuthVersion from firebolt.client.auth.request_auth_base import _RequestBasedAuth from firebolt.utils.token_storage import TokenSecureStorage from firebolt.utils.urls import AUTH_URL @@ -43,13 +43,13 @@ def __init__( self.password = password super().__init__(use_token_cache) - def get_firebolt_version(self) -> int: + def get_firebolt_version(self) -> FireboltAuthVersion: """Get Firebolt version from auth. Returns: - int: Firebolt version + FireboltAuthVersion: V1 for Username Password authentication """ - return 1 + return FireboltAuthVersion.V1 def copy(self) -> "UsernamePassword": """Make another auth object with same credentials. diff --git a/src/firebolt/client/client.py b/src/firebolt/client/client.py index dd6cb620943..b8d71280245 100644 --- a/src/firebolt/client/client.py +++ b/src/firebolt/client/client.py @@ -106,7 +106,11 @@ def clone(self) -> "Client": class Client(FireboltClientMixin, HttpxClient, metaclass=ABCMeta): def __init__(self, *args: Any, **kwargs: Any): - super().__init__(*args, **kwargs, transport=KeepaliveTransport()) + super().__init__( + *args, + **kwargs, + transport=KeepaliveTransport(verify=kwargs.get("verify", True)), + ) @property @abstractmethod @@ -269,7 +273,11 @@ def _resolve_engine_url(self, engine_name: str) -> str: class AsyncClient(FireboltClientMixin, HttpxAsyncClient, metaclass=ABCMeta): def __init__(self, *args: Any, **kwargs: Any): - super().__init__(*args, **kwargs, transport=AsyncKeepaliveTransport()) + super().__init__( + *args, + **kwargs, + transport=AsyncKeepaliveTransport(verify=kwargs.get("verify", True)), + ) @property @abstractmethod diff --git a/src/firebolt/db/connection.py b/src/firebolt/db/connection.py index 07b10f4ddc1..ef7b5ff62db 100644 --- a/src/firebolt/db/connection.py +++ b/src/firebolt/db/connection.py @@ -1,14 +1,16 @@ from __future__ import annotations import logging +from ssl import SSLContext from types import TracebackType -from typing import Any, Dict, List, Optional, Type +from typing import Any, Dict, List, Optional, Type, Union from warnings import warn from httpx import Timeout from firebolt.client import DEFAULT_API_URL, Client, ClientV1, ClientV2 from firebolt.client.auth import Auth +from firebolt.client.auth.base import FireboltAuthVersion from firebolt.common.base_connection import ( ASYNC_QUERY_CANCEL, ASYNC_QUERY_STATUS_REQUEST, @@ -27,6 +29,11 @@ ConnectionClosedError, FireboltError, ) +from firebolt.utils.firebolt_core import ( + get_core_certificate_context, + parse_firebolt_core_url, + validate_firebolt_core_parameters, +) from firebolt.utils.usage_tracker import get_user_agent_header from firebolt.utils.util import fix_url_schema, validate_engine_name_and_url_v1 @@ -41,6 +48,7 @@ def connect( engine_url: Optional[str] = None, api_endpoint: str = DEFAULT_API_URL, disable_cache: bool = False, + url: Optional[str] = None, additional_parameters: Dict[str, Any] = {}, ) -> Connection: # auth parameter is optional in function signature @@ -57,9 +65,20 @@ def connect( auth_version = auth.get_firebolt_version() if disable_cache: _firebolt_system_engine_cache.disable() - # Use v2 if auth is ClientCredentials - # Use v1 if auth is ServiceAccount or UsernamePassword - if auth_version == 2: + # Use CORE if auth is FireboltCore + # Use V2 if auth is ClientCredentials + # Use V1 if auth is ServiceAccount or UsernamePassword + if auth_version == FireboltAuthVersion.CORE: + # Verify that Core-incompatible parameters are not provided + validate_firebolt_core_parameters(account_name, engine_name, engine_url) + + return connect_core( + auth=auth, + user_agent_header=user_agent_header, + database=database, + connection_url=url, + ) + elif auth_version == FireboltAuthVersion.V2: assert account_name is not None return connect_v2( auth=auth, @@ -69,7 +88,7 @@ def connect( engine_name=engine_name, api_endpoint=api_endpoint, ) - elif auth_version == 1: + elif auth_version == FireboltAuthVersion.V1: return connect_v1( auth=auth, user_agent_header=user_agent_header, @@ -384,3 +403,49 @@ def connect_v1( headers={"User-Agent": user_agent_header}, ) return Connection(engine_url, database, client, CursorV1, api_endpoint) + + +def connect_core( + auth: Auth, + user_agent_header: str, + database: Optional[str] = None, + connection_url: Optional[str] = None, +) -> Connection: + """Connect to Firebolt Core. + + Args: + auth (Auth): Authentication object (must be FireboltCore) + user_agent_header (str): User agent header string + database (Optional[str]): Name of the database to connect to + (defaults to 'firebolt') + connection_url (Optional[str]): URL in format protocol://host:port + Protocol defaults to http, host defaults to localhost, port + defaults to 3473. + + Returns: + Connection: A connection to Firebolt Core + """ + connection_params = parse_firebolt_core_url(connection_url) + + ctx: Union[SSLContext, bool] = True # Default context + if connection_params.scheme == "https": + ctx = get_core_certificate_context() + + verified_url = connection_params.geturl() + + client = ClientV2( + auth=auth, + account_name="", # FireboltCore does not require an account name + base_url=verified_url, + timeout=Timeout(DEFAULT_TIMEOUT_SECONDS, read=None), + headers={"User-Agent": user_agent_header}, + verify=ctx, + ) + + return Connection( + engine_url=verified_url, + database=database, + client=client, + cursor_type=CursorV2, + api_endpoint=verified_url, + ) diff --git a/src/firebolt/service/manager.py b/src/firebolt/service/manager.py index 435b072ef50..432ece5dd3b 100644 --- a/src/firebolt/service/manager.py +++ b/src/firebolt/service/manager.py @@ -5,13 +5,14 @@ from firebolt.client import ( DEFAULT_API_URL, - Auth, ClientV1, ClientV2, log_request, log_response, raise_on_4xx_5xx, ) +from firebolt.client.auth import Auth +from firebolt.client.auth.base import FireboltAuthVersion from firebolt.common import Settings from firebolt.db import connect from firebolt.service.V1.binding import BindingService @@ -94,10 +95,10 @@ def __init__( assert auth is not None version = auth.get_firebolt_version() - if version == 2: + if version == FireboltAuthVersion.V2: client_class = ClientV2 assert account_name is not None - elif version == 1: + elif version == FireboltAuthVersion.V1: client_class = ClientV1 else: raise ValueError(f"Unsupported Firebolt version: {version}") diff --git a/src/firebolt/utils/firebolt_core.py b/src/firebolt/utils/firebolt_core.py new file mode 100644 index 00000000000..fcf7c092c99 --- /dev/null +++ b/src/firebolt/utils/firebolt_core.py @@ -0,0 +1,117 @@ +import os +import sys +from ssl import ( + PROTOCOL_TLS_CLIENT, + Purpose, + SSLContext, + TLSVersion, + create_default_context, +) +from typing import Optional, Union +from urllib.parse import ParseResult, urlparse + +from firebolt.utils.exception import ConfigurationError + + +def get_core_certificate_context() -> Union[SSLContext, bool]: + """Get the SSL context for Firebolt Core connections.""" + ctx: Union[SSLContext, bool] = True # Default context for SSL verification + if os.getenv("SSL_CERT_FILE"): + ctx = create_default_context( + Purpose.SERVER_AUTH, cafile=os.getenv("SSL_CERT_FILE") + ) + ctx.minimum_version = TLSVersion.TLSv1_2 + elif sys.version_info >= (3, 10): + # Can import truststore only if python is 3.10 or higher + import truststore + + ctx = truststore.SSLContext(PROTOCOL_TLS_CLIENT) + else: + raise ConfigurationError( + "Not able to use system certificate store for Firebolt Core." + " on Python < 3.10, you may need to set" + " SSL_CERT_FILE environment variable pointing to your .pem file." + ) + return ctx + + +def validate_firebolt_core_parameters( + account_name: Optional[str] = None, + engine_name: Optional[str] = None, + engine_url: Optional[str] = None, +) -> None: + """Validate that no incompatible parameters are provided with + FireboltCore authentication. + + Args: + account_name (Optional[str]): Account name parameter to validate + engine_name (Optional[str]): Engine name parameter to validate + engine_url (Optional[str]): Engine URL parameter to validate + + Raises: + ConfigurationError: If any incompatible parameters are provided + """ + forbidden_params = {} + if account_name: + forbidden_params["account_name"] = account_name + if engine_name: + forbidden_params["engine_name"] = engine_name + if engine_url: + forbidden_params["engine_url"] = engine_url + + if forbidden_params: + param_list = ", ".join(f"'{p}'" for p in forbidden_params.keys()) + raise ConfigurationError( + f"Parameters {param_list} are not compatible with Firebolt Core " + "authentication. These parameters should not be provided when " + "using Firebolt Core." + ) + + +def parse_firebolt_core_url(url: Optional[str] = None) -> ParseResult: + """Parse a Firebolt Core URL into its components. + + Args: + url (str, optional): URL in format protocol://host:port + Protocol defaults to http, host defaults to localhost, port + defaults to 3473. + + Returns: + Dict[str, str]: Dictionary with protocol, host, and port keys + + Raises: + ConfigurationError: If the URL is invalid + """ + if not url: + # Default values + return ParseResult( + scheme="http", + netloc="localhost:3473", + path="", + params="", + query="", + fragment="", + ) + + # Parse URL + try: + parsed = urlparse(url) + # Protocol + if parsed.scheme: + if parsed.scheme not in ["http", "https"]: + raise ConfigurationError( + f"Invalid protocol: {parsed.scheme}. Must be 'http' or 'https'." + ) + + if parsed.port and not (1 <= parsed.port <= 65535): + raise ConfigurationError( + f"Invalid port: {parsed.port}. Must be between 1 and 65535." + ) + except Exception as e: + if isinstance(e, ConfigurationError): + raise + raise ConfigurationError( + f"Invalid URL format: {url}. Expected format: protocol://host:port" + ) from e + + return parsed diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 18545f70768..3894fce433a 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -1,11 +1,12 @@ from logging import getLogger -from os import environ +from os import environ, getenv from time import time from typing import Optional from pytest import fixture, mark from firebolt.client.auth import ClientCredentials +from firebolt.client.auth.firebolt_core import FireboltCore from firebolt.client.auth.username_password import UsernamePassword LOGGER = getLogger(__name__) @@ -21,6 +22,7 @@ PASSWORD_ENV = "PASSWORD" ENGINE_URL_ENV = "ENGINE_URL" STOPPED_ENGINE_URL_ENV = "STOPPED_ENGINE_URL" +CORE_URL_ENV = "CORE_URL" # https://docs.pytest.org/en/latest/example/simple.html#control-skipping-of-tests-according-to-command-line-option # Adding slow marker to tests @@ -120,6 +122,11 @@ def auth(service_id: str, service_secret: Secret) -> ClientCredentials: return ClientCredentials(service_id, service_secret.value) +@fixture(scope="session") +def core_auth() -> FireboltCore: + return FireboltCore() + + @fixture(scope="session") def username() -> str: return must_env(USER_NAME_ENV) @@ -145,6 +152,11 @@ def stopped_engine_url() -> str: return must_env(STOPPED_ENGINE_URL_ENV) +@fixture(scope="session") +def core_url() -> str: + return getenv(CORE_URL_ENV, "") + + @fixture(scope="function") def minimal_time(): limit: Optional[float] = None diff --git a/tests/integration/dbapi/async/V2/conftest.py b/tests/integration/dbapi/async/V2/conftest.py index eb14724c758..7265dc1596f 100644 --- a/tests/integration/dbapi/async/V2/conftest.py +++ b/tests/integration/dbapi/async/V2/conftest.py @@ -10,20 +10,36 @@ from tests.integration.conftest import Secret -@fixture +@fixture(params=["remote", "core"]) async def connection( engine_name: str, database_name: str, auth: Auth, + core_auth: Auth, account_name: str, api_endpoint: str, + core_url: str, + request: Any, ) -> Connection: + if request.param == "core": + kwargs = { + "engine_name": None, + "database": "firebolt", + "auth": core_auth, + "url": core_url, + "account_name": None, + "api_endpoint": None, + } + else: + kwargs = { + "engine_name": engine_name, + "database": database_name, + "auth": auth, + "account_name": account_name, + "api_endpoint": api_endpoint, + } async with await connect( - engine_name=engine_name, - database=database_name, - auth=auth, - account_name=account_name, - api_endpoint=api_endpoint, + **kwargs, ) as connection: yield connection diff --git a/tests/integration/dbapi/async/V2/test_queries_async.py b/tests/integration/dbapi/async/V2/test_queries_async.py index d782d54e7e6..c1f339a5592 100644 --- a/tests/integration/dbapi/async/V2/test_queries_async.py +++ b/tests/integration/dbapi/async/V2/test_queries_async.py @@ -11,6 +11,7 @@ from firebolt.client.auth.base import Auth from firebolt.common._types import ColType from firebolt.common.row_set.types import Column +from firebolt.utils.exception import FireboltStructuredError from tests.integration.dbapi.utils import assert_deep_eq VALS_TO_INSERT_2 = ",".join( @@ -50,7 +51,7 @@ async def test_select( async with connection.cursor() as c: # For timestamptz test assert ( - await c.execute(f"SET time_zone={timezone_name}") == -1 + await c.execute(f"SET timezone={timezone_name}") == -1 ), "Invalid set statment row count" assert await c.execute(all_types_query) == 1, "Invalid row count returned" @@ -112,6 +113,8 @@ async def test_long_query( assert len(data) == 1, "Invalid data size returned by fetchall" +# Not compatible with core +@mark.parametrize("connection", ["remote"], indirect=True) async def test_drop_create(connection: Connection) -> None: """Create and drop table/index queries are handled properly.""" @@ -334,9 +337,12 @@ async def test_multi_statement_query(connection: Connection) -> None: async def test_set_invalid_parameter(connection: Connection): async with connection.cursor() as c: assert len(c._set_parameters) == 0 - with raises(OperationalError): + with raises((OperationalError, FireboltStructuredError)) as e: await c.execute("SET some_invalid_parameter = 1") + assert "Unknown setting" in str(e.value) or "query param not allowed" in str( + e.value + ), "Invalid error message" assert len(c._set_parameters) == 0 @@ -424,7 +430,7 @@ async def test_connection_with_mixed_case_db_and_engine( engine_name=test_engine_name, ) as connection: cursor = connection.cursor() - await cursor.execute('CREATE TABLE "test_table" (id int)') + await cursor.execute('CREATE TABLE IF NOT EXISTS "test_table" (id int)') # This fails if we're not running on a user engine await cursor.execute('INSERT INTO "test_table" VALUES (1)') diff --git a/tests/integration/dbapi/async/V2/test_server_async.py b/tests/integration/dbapi/async/V2/test_server_async.py index cb7f2b1e27e..5ab03e2c470 100644 --- a/tests/integration/dbapi/async/V2/test_server_async.py +++ b/tests/integration/dbapi/async/V2/test_server_async.py @@ -2,13 +2,17 @@ from random import randint from typing import Callable -from pytest import raises +from pytest import mark, raises from firebolt.async_db import Connection from firebolt.utils.exception import FireboltError, FireboltStructuredError LONG_SELECT = "SELECT checksum(*) FROM GENERATE_SERIES(1, 2500000000)" # approx 3 sec +# Async is only supported in remote engines, no core support yet +# Mark all the tests in this module to run only with the "remote" connection +pytestmark = mark.parametrize("connection", ["remote"], indirect=True) + async def test_insert_async(connection: Connection) -> None: cursor = connection.cursor() @@ -65,6 +69,7 @@ async def test_insert_async_running(connection: Connection) -> None: async def test_check_async_execution_from_another_connection( + connection: Connection, connection_factory: Callable[..., Connection], ) -> None: connection_1 = await connection_factory() diff --git a/tests/integration/dbapi/async/V2/test_streaming.py b/tests/integration/dbapi/async/V2/test_streaming.py index e67e6817170..c17c9330dda 100644 --- a/tests/integration/dbapi/async/V2/test_streaming.py +++ b/tests/integration/dbapi/async/V2/test_streaming.py @@ -22,7 +22,7 @@ async def test_streaming_select( async with connection.cursor() as c: # For timestamptz test assert ( - await c.execute(f"SET time_zone={timezone_name}") == -1 + await c.execute(f"SET timezone={timezone_name}") == -1 ), "Invalid set statment row count" await c.execute_stream(all_types_query) diff --git a/tests/integration/dbapi/conftest.py b/tests/integration/dbapi/conftest.py index 112749684d8..9d2b53c98ee 100644 --- a/tests/integration/dbapi/conftest.py +++ b/tests/integration/dbapi/conftest.py @@ -215,7 +215,6 @@ def select_geography_response() -> List[ColType]: def setup_struct_query() -> str: return """ SET advanced_mode=1; - SET enable_struct=true; SET enable_create_table_v2=true; SET enable_struct_syntax=true; SET prevent_create_on_information_schema=true; diff --git a/tests/integration/dbapi/sync/V2/conftest.py b/tests/integration/dbapi/sync/V2/conftest.py index a116657d4f1..c6b70f66de0 100644 --- a/tests/integration/dbapi/sync/V2/conftest.py +++ b/tests/integration/dbapi/sync/V2/conftest.py @@ -10,20 +10,36 @@ from tests.integration.conftest import Secret -@fixture +@fixture(params=["remote", "core"]) def connection( engine_name: str, database_name: str, auth: Auth, + core_auth: Auth, account_name: str, api_endpoint: str, + core_url: str, + request: Any, ) -> Connection: + if request.param == "core": + kwargs = { + "engine_name": None, + "database": "firebolt", + "auth": core_auth, + "url": core_url, + "account_name": None, + "api_endpoint": None, + } + else: + kwargs = { + "engine_name": engine_name, + "database": database_name, + "auth": auth, + "account_name": account_name, + "api_endpoint": api_endpoint, + } with connect( - engine_name=engine_name, - database=database_name, - auth=auth, - account_name=account_name, - api_endpoint=api_endpoint, + **kwargs, ) as connection: yield connection diff --git a/tests/integration/dbapi/sync/V2/test_queries.py b/tests/integration/dbapi/sync/V2/test_queries.py index 20653ba6c8b..25e72c66719 100644 --- a/tests/integration/dbapi/sync/V2/test_queries.py +++ b/tests/integration/dbapi/sync/V2/test_queries.py @@ -11,6 +11,7 @@ from firebolt.common._types import ColType from firebolt.common.row_set.types import Column from firebolt.db import Binary, Connection, Cursor, OperationalError, connect +from firebolt.utils.exception import FireboltStructuredError from tests.integration.dbapi.utils import assert_deep_eq VALS_TO_INSERT = ",".join([f"({i},'{val}')" for (i, val) in enumerate(range(1, 360))]) @@ -56,7 +57,7 @@ def test_select( with connection.cursor() as c: # For timestamptz test assert ( - c.execute(f"SET time_zone={timezone_name}") == -1 + c.execute(f"SET timezone={timezone_name}") == -1 ), "Invalid set statment row count" assert c.execute(all_types_query) == 1, "Invalid row count returned" @@ -116,6 +117,8 @@ def test_long_query( assert len(data) == 1, "Invalid data size returned by fetchall" +# Not compatible with core +@mark.parametrize("connection", ["remote"], indirect=True) def test_drop_create(connection: Connection) -> None: """Create and drop table/index queries are handled properly.""" @@ -334,9 +337,11 @@ def test_multi_statement_query(connection: Connection) -> None: def test_set_invalid_parameter(connection: Connection): with connection.cursor() as c: assert len(c._set_parameters) == 0 - with raises(OperationalError): + with raises((OperationalError, FireboltStructuredError)) as e: c.execute("set some_invalid_parameter = 1") - + assert "Unknown setting" in str(e.value) or "query param not allowed" in str( + e.value + ), "Invalid error message" assert len(c._set_parameters) == 0 @@ -507,7 +512,7 @@ def test_connection_with_mixed_case_db_and_engine( api_endpoint=api_endpoint, ) as connection: cursor = connection.cursor() - cursor.execute('CREATE TABLE "test_table" (id int)') + cursor.execute('CREATE TABLE IF NOT EXISTS "test_table" (id int)') # This fails if we're not running on a user engine cursor.execute('INSERT INTO "test_table" VALUES (1)') diff --git a/tests/integration/dbapi/sync/V2/test_server_async.py b/tests/integration/dbapi/sync/V2/test_server_async.py index 993b0923d20..1a5573d01f9 100644 --- a/tests/integration/dbapi/sync/V2/test_server_async.py +++ b/tests/integration/dbapi/sync/V2/test_server_async.py @@ -2,13 +2,17 @@ from random import randint from typing import Callable -from pytest import raises +from pytest import mark, raises from firebolt.db import Connection from firebolt.utils.exception import FireboltError, FireboltStructuredError LONG_SELECT = "SELECT checksum(*) FROM GENERATE_SERIES(1, 2500000000)" # approx 3 sec +# Async is only supported in remote engines, no core support yet +# Mark all the tests in this module to run only with the "remote" connection +pytestmark = mark.parametrize("connection", ["remote"], indirect=True) + def test_insert_async(connection: Connection) -> None: cursor = connection.cursor() @@ -63,6 +67,7 @@ def test_insert_async_running(connection: Connection) -> None: def test_check_async_execution_from_another_connection( + connection: Connection, connection_factory: Callable[..., Connection], ) -> None: connection_1 = connection_factory() diff --git a/tests/integration/dbapi/sync/V2/test_streaming.py b/tests/integration/dbapi/sync/V2/test_streaming.py index fb6a0b8f058..35b070f9bf2 100644 --- a/tests/integration/dbapi/sync/V2/test_streaming.py +++ b/tests/integration/dbapi/sync/V2/test_streaming.py @@ -22,7 +22,7 @@ def test_streaming_select( with connection.cursor() as c: # For timestamptz test assert ( - c.execute(f"SET time_zone={timezone_name}") == -1 + c.execute(f"SET timezone={timezone_name}") == -1 ), "Invalid set statment row count" c.execute_stream(all_types_query) diff --git a/tests/unit/async_db/core/__init__.py b/tests/unit/async_db/core/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/async_db/core/conftest.py b/tests/unit/async_db/core/conftest.py new file mode 100644 index 00000000000..94b583e6022 --- /dev/null +++ b/tests/unit/async_db/core/conftest.py @@ -0,0 +1,11 @@ +"""Fixtures for Async Firebolt Core unit tests.""" + +from pytest import fixture + +from firebolt.client.auth import FireboltCore + + +@fixture +def core_auth() -> FireboltCore: + """FireboltCore auth object for unit tests.""" + return FireboltCore() diff --git a/tests/unit/async_db/core/test_connection_validation.py b/tests/unit/async_db/core/test_connection_validation.py new file mode 100644 index 00000000000..0676e15b0e6 --- /dev/null +++ b/tests/unit/async_db/core/test_connection_validation.py @@ -0,0 +1,43 @@ +from unittest.mock import MagicMock, patch + +import pytest +from pytest import raises + +from firebolt.async_db import connect +from firebolt.client.auth.base import FireboltAuthVersion +from firebolt.utils.exception import ConfigurationError + + +@pytest.mark.asyncio +async def test_async_connect_with_incompatible_params(): + """Test that async connect rejects incompatible parameters with FireboltCore.""" + with patch("firebolt.async_db.connection.connect_core") as mock_connect_core: + + # Create a mock FireboltCore auth that returns the correct version + mock_auth = MagicMock() + mock_auth.get_firebolt_version.return_value = FireboltAuthVersion.CORE + + # Test with account_name + with raises(ConfigurationError, match="'account_name' are not compatible"): + await connect(auth=mock_auth, account_name="test_account") + + # Test with engine_name + with raises(ConfigurationError, match="'engine_name' are not compatible"): + await connect(auth=mock_auth, engine_name="test_engine") + + # Test with engine_url + with raises(ConfigurationError, match="'engine_url' are not compatible"): + await connect(auth=mock_auth, engine_url="https://example.com") + + # Test with multiple incompatible parameters + with raises(ConfigurationError, match="'account_name', 'engine_name'"): + await connect( + auth=mock_auth, account_name="test_account", engine_name="test_engine" + ) + + # Verify connect_core is not called in any of these cases + mock_connect_core.assert_not_called() + + # Test with compatible parameters + await connect(auth=mock_auth, database="test_db") + mock_connect_core.assert_called_once() diff --git a/tests/unit/db/core/__init__.py b/tests/unit/db/core/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/db/core/conftest.py b/tests/unit/db/core/conftest.py new file mode 100644 index 00000000000..17f91a5aa8c --- /dev/null +++ b/tests/unit/db/core/conftest.py @@ -0,0 +1,11 @@ +"""Fixtures for Firebolt Core unit tests.""" + +from pytest import fixture + +from firebolt.client.auth import FireboltCore + + +@fixture +def core_auth() -> FireboltCore: + """FireboltCore auth object for unit tests.""" + return FireboltCore() diff --git a/tests/unit/db/core/test_connection_core.py b/tests/unit/db/core/test_connection_core.py new file mode 100644 index 00000000000..0028c563858 --- /dev/null +++ b/tests/unit/db/core/test_connection_core.py @@ -0,0 +1,78 @@ +from typing import Callable +from unittest.mock import MagicMock, patch + +from pytest import raises +from pytest_httpx import HTTPXMock + +from firebolt.client.auth import FireboltCore +from firebolt.client.auth.base import FireboltAuthVersion +from firebolt.db import connect +from firebolt.utils.exception import ConfigurationError + + +def test_sync_connect_with_incompatible_params(): + """Test that sync connect rejects incompatible parameters with FireboltCore.""" + with patch("firebolt.db.connection.connect_core") as mock_connect_core: + + # Create a mock FireboltCore auth that returns the correct version + mock_auth = MagicMock() + mock_auth.get_firebolt_version.return_value = FireboltAuthVersion.CORE + + # Test with account_name + with raises(ConfigurationError, match="'account_name' are not compatible"): + connect(auth=mock_auth, account_name="test_account") + + # Test with engine_name + with raises(ConfigurationError, match="'engine_name' are not compatible"): + connect(auth=mock_auth, engine_name="test_engine") + + # Test with engine_url + with raises(ConfigurationError, match="'engine_url' are not compatible"): + connect(auth=mock_auth, engine_url="https://example.com") + + # Test with multiple incompatible parameters + with raises(ConfigurationError, match="'account_name', 'engine_name'"): + connect( + auth=mock_auth, account_name="test_account", engine_name="test_engine" + ) + + # Verify connect_core is not called in any of these cases + mock_connect_core.assert_not_called() + + # Test with compatible parameters + connect(auth=mock_auth, database="test_db") + mock_connect_core.assert_called_once() + + +def test_firebolt_core_no_requests(httpx_mock: HTTPXMock): + """Test that FireboltCore auth class doesn't send any requests during initialization.""" + # Create FireboltCore auth, no requests should be sent + FireboltCore() + + # Verify no requests were made + assert len(httpx_mock.get_requests()) == 0 + + +def test_core_connection_single_query_request( + httpx_mock: HTTPXMock, select_one_query_callback: Callable +): + """Test that a FireboltCore connection only makes a single request when running a query.""" + + httpx_mock.add_callback(select_one_query_callback) + + # Create auth and connection + auth = FireboltCore() + + # Connect and run a query + with connect(auth=auth, database="test_db") as connection: + cursor = connection.cursor() + cursor.execute("SELECT 1") + + # Verify exactly one request was made + requests = httpx_mock.get_requests() + assert len(requests) == 1 + + # Verify the request was for the query execution + request = requests[0] + assert request.method == "POST" + assert "SELECT 1" in request.content.decode() diff --git a/tests/unit/utils/test_firebolt_core.py b/tests/unit/utils/test_firebolt_core.py new file mode 100644 index 00000000000..2b49905a9d8 --- /dev/null +++ b/tests/unit/utils/test_firebolt_core.py @@ -0,0 +1,106 @@ +"""Unit tests for Firebolt Core utility functions.""" + +from pytest import raises + +from firebolt.utils.exception import ConfigurationError +from firebolt.utils.firebolt_core import ( + parse_firebolt_core_url, + validate_firebolt_core_parameters, +) + + +def test_ipv6_url_parsing(): + """Test IPv6 URL parsing for FireboltCore.""" + # Test with different IPv6 address formats + ipv6_urls = [ + ("http://[::1]:3473", "http", "::1", 3473), + ("http://[2001:db8::1]:8080", "http", "2001:db8::1", 8080), + ( + "https://[2001:db8:85a3:8d3:1319:8a2e:370:7348]:443", + "https", + "2001:db8:85a3:8d3:1319:8a2e:370:7348", + 443, + ), + ] + + for url, expected_protocol, expected_host, expected_port in ipv6_urls: + parsed = parse_firebolt_core_url(url) + + assert parsed.scheme == expected_protocol + assert parsed.hostname == expected_host + assert parsed.port == expected_port + + +def test_url_parsing(): + """Test URL parsing for FireboltCore.""" + # Test with different URL formats + urls = [ + # Default values when no URL is provided + (None, "http", "localhost", 3473), + # Full URL + ("http://example.com:8080", "http", "example.com", 8080), + # Protocol only + ("https://localhost:3473", "https", "localhost", 3473), + # Host only + ("http://custom-host:3473", "http", "custom-host", 3473), + # Port only + ("http://localhost:9999", "http", "localhost", 9999), + ] + + for url, expected_protocol, expected_host, expected_port in urls: + parsed = parse_firebolt_core_url(url) + + assert parsed.scheme == expected_protocol + assert parsed.hostname == expected_host + assert parsed.port == expected_port + + +def test_invalid_url_parsing(): + """Test that invalid URLs with forbidden protocols raise exceptions.""" + # Based on the implementation, only invalid protocols are actually checked + invalid_urls = [ + "ftp://localhost:3473", # Invalid protocol + "ssh://example.com:22", # Invalid protocol + ] + + for url in invalid_urls: + try: + parse_firebolt_core_url(url) + assert False, f"Expected ConfigurationError for URL: {url}" + except ConfigurationError: + pass # Expected exception + + # Test invalid port (too large) + with raises(ConfigurationError, match="Invalid URL format"): + parse_firebolt_core_url("http://localhost:70000") + + +def test_validate_firebolt_core_parameters(): + """Test validation of parameters for FireboltCore authentication.""" + # Test with no parameters (should pass) + validate_firebolt_core_parameters() # No exception should be raised + + # Test with individual forbidden parameters + with raises(ConfigurationError, match="'account_name' are not compatible"): + validate_firebolt_core_parameters(account_name="test_account") + + with raises(ConfigurationError, match="'engine_name' are not compatible"): + validate_firebolt_core_parameters(engine_name="test_engine") + + with raises(ConfigurationError, match="'engine_url' are not compatible"): + validate_firebolt_core_parameters(engine_url="https://example.com") + + # Test with multiple forbidden parameters + with raises(ConfigurationError, match="'account_name', 'engine_name'"): + validate_firebolt_core_parameters( + account_name="test_account", engine_name="test_engine" + ) + + with raises( + ConfigurationError, match="'account_name', 'engine_name', 'engine_url'" + ): + validate_firebolt_core_parameters( + account_name="test_account", + engine_name="test_engine", + engine_url="https://example.com", + )