DOCSEARCHbyAlgolia
Submit Docs

Databricks SQL Connector for Python

repository·main·Indexed 42 days ago

https://github.com/databricks/databricks-sql-python

A Thrift-based, DB API 2.0 compliant Python client for connecting to Databricks clusters and SQL warehouses. Version 4.5.0 supports high-performance data exchange via Apache Arrow and an optional Rust-based kernel backend for improved performance. It provides support for various authentication methods including PAT, OAuth (U2M and M2M), and Azure service principals, and integrates with SQLAlchemy 2.0 via the databricks-sqlalchemy library.

Tokens
26.2K
Snippets
67
Records
108
Agent score
53%

What's inside Databricks SQL Connector for Python

  1. Use native parameters for safe query execution

    main

    The Databricks SQL Connector supports native parameterized query execution. Instead of the connector escaping strings and rendering them inline (the legacy behavior), native parameters send variable markers and a separate collection of values to the Databricks Runtime. This prevents SQL injection and allows the server to cache prepared statements for better performance.

    Requirements:

    • databricks-sql-connector>=3.0.0
    • A SQL warehouse or all-purpose cluster running Databricks Runtime >=14.2

    Limitations:

    • Maximum of 255 parameter markers per query.
    • Total size of all parameterized values cannot exceed 1MB.
    • Native parameters are not supported for volume operations like PUT.

  2. When to use `use_inline_params=True` vs `False`

    main

    The connector supports two parameterization modes.

    Inline Parameters (use_inline_params=True)

    Use this mode only if:

    1. You need to pass more than 255 parameters in a single query.
    2. You need to pass parameter values greater than 1MB in a single query.
    3. Your code relies heavily on %s positional markers.
    4. Your code passes sequences (lists) as parameter values (e.g., for IN clauses).

    Native Parameters (use_inline_params=False)

    Native parameters are the recommended approach for most users. You should migrate to native parameters to:

    • Improve security/safety.
    • Enable the use of SQL LIKE wildcards (%) without causing exceptions.
    • Access future support for complex types like STRUCT, MAP, and ARRAY.

    Note: If you use use_inline_params=False, %s occurrences are passed directly to the database, allowing them to function correctly within LIKE statements.

  3. How the Bot-Prelude Composite works

    main

    To reduce boilerplate across multiple bot workflows (Reviewer, Engineer Author, Engineer Followup), a shared local composite action ./.github/actions/bot-prelude can be used. This composite encapsulates the common setup steps required before running the engine.

    Inputs

    • app-id: The ID of the relevant GitHub App (Reviewer or Engineer).
    • private-key: The private key for the specified App.
    • engine-ref: The specific SHA or reference of the databricks-bot-engine to use.

    Outputs

    • token: The minted bot token for the current repository.

    Workflow Pattern

    1. Checkout: Perform the repository checkout (specific to the flow).
    2. Setup Environment: Either Setup Python (for Reviewer) or setup-poetry (for Engineer, to ensure product dependencies are available).
    3. Use Prelude: Call ./.github/actions/bot-prelude with the required App credentials.
    4. Run Engine: Execute the engine entrypoint via python -m databricks_bot_engine.<bot_name>.run_<flow>.

  4. Manage transactions with manual commit/rollback

    main

    The connector supports multi-statement transactions. By default, the connection is in autocommit mode. To manage transactions manually, set connection.autocommit = False. You can then use connection.commit() to save changes or connection.rollback() to undo them.

    connection.autocommit = False
# ... execute statements ...
connection.commit()
# or
connection.rollback()

  5. Manage transaction autocommit behavior

    main

    By default, the Databricks SQL Connector operates in autocommit mode (connection.autocommit = True), where every SQL statement executes in its own transaction and commits immediately. To group multiple statements into an atomic multi-statement transaction (MST), you must explicitly disable autocommit by setting connection.autocommit = False.

    from databricks import sql

connection = sql.connect(
    server_hostname="your-workspace.cloud.databricks.com",
    http_path="/sql/1.0/warehouses/abc123"
)

# Default behavior: autocommit is True
print(connection.autocommit)  # True

# To use explicit transactions:
connection.autocommit = False
cursor = connection.cursor()
cursor.execute("INSERT INTO my_table VALUES (1, 'data')")
# Data is not yet visible to other connections until commit is called
connection.commit()

  6. How backend selection works in the Databricks SQL Connector

    main

    The Databricks SQL Python connector supports multiple backends for executing queries. By default, it uses the Thrift backend. You can opt-in to the Kernel backend, which is a native Rust core exposed via PyO3, by setting use_kernel=True in the connect() call.

    Backend Comparison

    • Thrift (Default): The standard backend. It supports the widest range of parameters.
    • Kernel (Early Access): A high-performance Rust-based backend.
      • Requires Python ≥ 3.10.
      • Requires installing the databricks-sql-connector[kernel] extra.
      • Note: The Kernel backend only accepts a curated subset of parameters. Many Thrift-specific options are ignored or rejected when using the Kernel backend.
    • SEA (Deprecated): A pure-Python backend selected via use_sea=True. It is being deprecated and is mutually exclusive with the Kernel backend.

    Selection Parameters

    OptionTypeThriftKernelDefaultNote
    use_kernelbool——FalseSelects Kernel backend. Mutually exclusive with use_sea.
    use_seabool——FalseSelects deprecated SEA backend. Mutually exclusive with use_kernel.
    use_hybrid_dispositionbool❌❌FalseSEA-only option; ignored by Thrift and Kernel.
    max_connectionsint❌❌10SEA-only HTTP-pool limit; ignored by Thrift and Kernel.
    # To use the high-performance Rust kernel backend
from databricks import sql

connection = sql.connect(
    server_hostname="dbc-xxxx.cloud.databricks.com",
    http_path="/sql/1.0/endpoints/xxxx",
    access_token="your_token",
    use_kernel=True
)

  7. Configure OAuth U2M redirect ports

    main

    When using a custom oauth_client_id for User-to-Machine (U2M) authentication, you can specify a specific local port for the browser redirect flow using oauth_redirect_port.

    Important Behavior:

    • On both backends, this option is only honored if a custom oauth_client_id is supplied.
    • If you use the built-in client ID (or omit the option), the connector uses the registered range 8020–8024 and binds the first available port. In this case, providing a oauth_redirect_port has no effect.

  8. Understand backend behavioral differences (Thrift vs Kernel)

    main

    When choosing between the Thrift and Kernel backends, be aware of the following functional gaps and behavioral divergences:

    Features available on Thrift but NOT on Kernel

    • Custom credentials_provider.
    • experimental_oauth_persistence (custom OAuth token store).
    • TLS-client-cert authentication (_use_cert_as_auth).
    • Result-transport tuning: use_cloud_fetch, max_download_threads, enable_query_result_lz4_compression.
    • Arrow-native rendering for decimals and timestamps.
    • Unity Catalog Volume staging (staging_allowed_local_path).
    • Advanced retry tuning (e.g., _retry_delay_default, _retry_dangerous_codes).
    • force_enable_telemetry.
    • Specific connection parameters: _port, _connection_uri, _pool_connections, and _proxy_auth_method.

    Features available on Kernel but NOT on Thrift

    • Kernel-managed OAuth M2M: oauth_client_secret, oauth_jwt_key_file, oauth_jwt_kid, oauth_jwt_passphrase, oauth_jwt_algorithm, and token_url.
    • Custom oauth_scopes.
    • Kernel U2M encrypted token storage (oauth_token_cache_enabled).

    Key Behavioral Divergences

    • Authentication Defaults: A credential-less Thrift connection defaults to Databricks OAuth U2M. The Kernel requires an explicit U2M auth_type or another supported credential flow.
    • Connection Pooling: _pool_connections only configures the shared Python client. However, _pool_maxsize configures the kernel's Rust HTTP pool if positive and using kernel $\ge$ 1.1.0. A value of 0 treats it as unset (Python client keeps 20, Kernel keeps 100).
    • Custom Headers: The Kernel drops Authorization and x-databricks-org-id headers but forwards other http_headers.
    • Telemetry: The Kernel owns its telemetry lifecycle and does not honor force_enable_telemetry.

  9. Quickstart: Connect and execute a query

    main

    To connect to Databricks, you need the server_hostname (the Databricks instance host name) and the http_path (the path to a SQL endpoint or a Databricks Runtime interactive cluster).

    This example demonstrates a basic connection using environment variables and executing a parameterized query. Note that this specific example uses Databricks OAuth U2M, which requires a browser for authentication and is intended for use on a user's local machine.

    import os
from databricks import sql

host = os.getenv("DATABRICKS_HOST")
http_path = os.getenv("DATABRICKS_HTTP_PATH")

connection = sql.connect(
  server_hostname=host,
  http_path=http_path)

cursor = connection.cursor()
cursor.execute('SELECT :param `p`, * FROM RANGE(10)', {"param": "foo"})
result = cursor.fetchall()
for row in result:
  print(row)

cursor.close()
connection.close()

  10. Use SQLAlchemy with Databricks

    main

    As of version 4.0.0, SQLAlchemy support is provided via a separate library called databricks-sqlalchemy. You can choose between SQLAlchemy v1 or v2 dialects.

    Installation:

    • For SQLAlchemy v1: pip install databricks-sqlalchemy~=1.0
    • For SQLAlchemy v2: pip install databricks-sqlalchemy

  11. Configure Engineer-bot for bug-fixes

    main

    The engineer-bot is used to automate bug-fixes and tasks. It requires specific configuration files and environment setup to function correctly within the repository.

    Required Configuration Files

    • .bot/config.yaml: Defines the flow (e.g., flow: bug-fix) and an bash_allowlist containing permitted commands (e.g., [poetry, run, python, -m, pytest]).
    • Prompt Files: Located in the .bot/ directory (e.g., engineer/system.md, engineer/user.md, engineer-followup/system.md). These provide the LLM with context about the codebase (e.g., src/databricks/sql/) and testing constraints.

    Environment and Triggers

    • Environment: Use environment: azure-prod on jobs to ensure the bot can access DATABRICKS_HOST and DATABRICKS_TOKEN from env-scoped secrets.
    • Triggers: The bot is typically triggered by applying the engineer-bot label to an issue (for the author flow) or a PR (for the followup flow).
    • Dependency Management: You must run a dependency installation step (like poetry install) before the engine runs to ensure the connector's dependencies are present for self-verification.
    # .bot/config.yaml example
flow: bug-fix
bash_allowlist:
  - poetry
  - run
  - python
  - -m
  - pytest

  12. Run Databricks SQL Connector example scripts

    main

    You can run the provided examples by either cloning the entire repository or fetching individual scripts using curl.

    Option 1: Clone the repository

    1. Clone the repository to your local system.
    2. Install the connector and obtain credentials as described in the README.
    3. Navigate to the examples directory: cd examples/.
    4. Run a script using the Python CLI, for example: python query_execute.py.

    Option 2: Fetch individual scripts with curl

    1. Install the connector and obtain credentials.
    2. Find the Raw URL of the script on GitHub (e.g., https://raw.githubusercontent.com/databricks/databricks-sql-python/main/examples/query_execute.py).
    3. Download the script: curl <RAW_URL> > <FILENAME>.py.
    4. Run the script: python <FILENAME>.py.
    # Example: Fetching and running a specific script
curl https://raw.githubusercontent.com/databricks/databricks-sql-python/main/examples/query_execute.py > query_execute.py
python query_execute.py