Use SQLAlchemy with Azure Databricks

Azure Databricks provides a SQLAlchemy dialect (the system SQLAlchemy uses to communicate with various types of database API implementations and databases) for Azure Databricks. SQLAlchemy is a Python SQL toolkit and Object Relational Mapper (ORM). SQLAlchemy provides a suite of well known enterprise-level persistence patterns, designed for efficient and high-performing database access, adapted into a simple and Pythonic domain language. See Features and Philosophy.

The SQLAlchemy dialect for Azure Databricks needs to be installed to use the SQLAlchemy features with Azure Databricks. This article covers SQLAlchemy dialect for Azure Databricks version 1.0 and 2.0, which will be based on Databricks SQL Connector for Python version 4.0.0 or above.

Requirements

  • A development machine running Python >=3.8 and <=3.11.
  • Databricks recommends that you use Python virtual environments, such as those provided by venv that are included with Python. Virtual environments help to ensure that you are using the correct versions of Python and the Databricks SQL Connector for Python together. Setting up and using virtual environments is outside of the scope of this article. For more information, see Creating virtual environments.
  • An existing cluster or SQL warehouse.

Get started

  • Install databricks-sqlalchemy for SQLAlchemy v1 using pip install databricks-sqlalchemy~=1.0 or for SQLAlchemy v2 using pip install databricks-sqlalchemy. For version information, see the databricks-sqlalchemy release history.

  • Gather the following information for the cluster or SQL warehouse that you want to use:

    Cluster

    SQL warehouse

    • The server hostname of the SQL warehouse. You can get this from the Server Hostname value in the Connection Details tab for your SQL warehouse.
    • The HTTP path of the SQL warehouse. You can get this from the HTTP Path value in the Connection Details tab for your SQL warehouse.

Authentication

The SQLAlchemy dialect for Azure Databricks uses the Databricks SQL Connector for Python to connect, and passes the values in connect_args to the connector. It supports the following authentication types:

Personal access token authentication

To create an Azure Databricks personal access token, follow the steps in Create personal access tokens for workspace users.

To authenticate the SQLAlchemy dialect, use the following code snippet. This snippet assumes that you have set the following environment variables:

  • DATABRICKS_TOKEN, set to the Azure Databricks personal access token.
  • DATABRICKS_SERVER_HOSTNAMEset to the Server Hostname value for your cluster or SQL warehouse.
  • DATABRICKS_HTTP_PATH, set to HTTP Path value for your cluster or SQL warehouse.
  • DATABRICKS_CATALOG, set to the target catalog in Unity Catalog.
  • DATABRICKS_SCHEMA, set to the target schema (also known as a database) in Unity Catalog.

To set environment variables, see your operating system's documentation.

import os
from sqlalchemy import create_engine

access_token    = os.getenv("DATABRICKS_TOKEN")
server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME")
http_path       = os.getenv("DATABRICKS_HTTP_PATH")
catalog         = os.getenv("DATABRICKS_CATALOG")
schema          = os.getenv("DATABRICKS_SCHEMA")

engine = create_engine(
  url = f"databricks://token:{access_token}@{server_hostname}?" +
        f"http_path={http_path}&catalog={catalog}&schema={schema}"
)

# ...

You use the preceding engine variable to connect to your specified catalog and schema through your Azure Databricks compute resource.

OAuth user-to-machine (U2M) authentication

OAuth U2M authentication uses real-time human sign-in and consent to authenticate the target Azure Databricks user account. For the Databricks SQL Connector for Python versions that support it on each cloud, see OAuth user-to-machine (U2M) authentication.

To authenticate the SQLAlchemy dialect with OAuth U2M, use the following code snippet. This snippet assumes that you have set the DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, DATABRICKS_CATALOG, and DATABRICKS_SCHEMA environment variables, as described in Personal access token authentication.

import os
from sqlalchemy import create_engine

server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME")
http_path       = os.getenv("DATABRICKS_HTTP_PATH")
catalog         = os.getenv("DATABRICKS_CATALOG")
schema          = os.getenv("DATABRICKS_SCHEMA")

engine = create_engine(
  url = f"databricks://{server_hostname}?" +
        f"http_path={http_path}&catalog={catalog}&schema={schema}",
  connect_args = {"auth_type": "databricks-oauth"}
)

# ...

OAuth machine-to-machine (M2M) authentication

OAuth M2M authentication authenticates as an Azure Databricks service principal. You must also install the Databricks SDK for Python (for example, by running pip install databricks-sdk). For the Databricks SQL Connector for Python and SDK versions that support it on each cloud, see OAuth machine-to-machine (M2M) authentication.

To use OAuth M2M authentication, do the following:

  1. Create an Azure Databricks service principal in your Azure Databricks workspace, and create an OAuth secret for that service principal. See Authorize service principal access to Azure Databricks with OAuth.
  2. Give that service principal access to your cluster or SQL warehouse. See Compute permissions or Manage a SQL warehouse.

To authenticate the SQLAlchemy dialect with OAuth M2M, use the following code snippet. In addition to the DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, DATABRICKS_CATALOG, and DATABRICKS_SCHEMA environment variables, this snippet assumes that you have set the following environment variables:

  • DATABRICKS_CLIENT_ID, set to the service principal's UUID or Application ID value.
  • DATABRICKS_CLIENT_SECRET, set to the Secret value for the service principal's OAuth secret.
import os
from databricks.sdk.core import Config, oauth_service_principal
from sqlalchemy import create_engine

server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME")
http_path       = os.getenv("DATABRICKS_HTTP_PATH")
catalog         = os.getenv("DATABRICKS_CATALOG")
schema          = os.getenv("DATABRICKS_SCHEMA")

def credential_provider():
  config = Config(
    host          = f"https://{server_hostname}",
    client_id     = os.getenv("DATABRICKS_CLIENT_ID"),
    client_secret = os.getenv("DATABRICKS_CLIENT_SECRET"))
  return oauth_service_principal(config)

engine = create_engine(
  url = f"databricks://{server_hostname}?" +
        f"http_path={http_path}&catalog={catalog}&schema={schema}",
  connect_args = {"credentials_provider": credential_provider}
)

# ...

SQLAlchemy v1

For connection examples, refer this example.py

SQLAlchemy v2

For connection examples, see the following section and the sqlalchemy_example.py file in GitHub.

DBAPI reference

Additional resources

  • Last updated on