将 SQLAlchemy 与 Azure Databricks 配合使用

Azure Databricks 提供适用于 Azure Databricks 的 SQLAlchemy方言（SQLAlchemy 用于与各种类型的数据库 API 实现和数据库进行通信的系统）。 SQLAlchemy 是 Python SQL 工具包和对象关系映射器 (ORM)。 SQLAlchemy 提供了一整套众所周知的企业级持久性模式，旨在实现高效且高性能的数据库访问，并采用简单的 Pythonic 域语言。请参见功能和理念。

需要安装 Azure Databricks 的 SQLAlchemy 方言才能将 SQLAlchemy 功能与 Azure Databricks 配合使用。本文介绍适用于 Azure Databricks 版本 1.0 和 2.0 的 SQLAlchemy 方言，该方言将基于用于 Python 4.0.0 或更高版本的 Databricks SQL 连接器。

要求

运行 Python >=3.8 和<=3.11 的开发计算机。
Databricks 建议使用 Python 虚拟环境，例如 Python 自带的 venv 提供的虚拟环境。虚拟环境有助于确保同时使用正确版本的 Python 和适用于 Python 的 Databricks SQL 连接器。设置和使用虚拟环境不在本文的讨论范围之内。有关详细信息，请参阅创建虚拟环境。
现有群集或 SQL 仓库。

开始

请使用 pip install databricks-sqlalchemy~=1.0（针对 SQLAlchemy v1）或 pip install databricks-sqlalchemy（针对 SQLAlchemy v2）安装 databricks-sqlalchemy。有关版本信息，请参阅 databricks-sqlalchemy 发布历史记录。
收集想要使用的群集或 SQL 仓库的以下信息：

群集
- 群集的服务器主机名。可以从群集的高级选项 JDBC/ODBC 选项卡中的服务器主机名值获取此主机名。
- 群集的 HTTP 路径。您可以在您的群集的Advanced Options JDBC/ODBC选项卡中的HTTP Path值中获取此内容。
SQL 数据仓库
- SQL 仓库的服务器主机名。从 SQL 仓库的“连接详细信息”选项卡中的“服务器主机名”值中可以获取此主机名。
- SQL 仓库的 HTTP 路径。从 SQL 仓库的“连接详细信息”选项卡中的“HTTP 路径”值可以获取此路径。

身份验证

Azure Databricks 的 SQLAlchemy 方言使用 Python 版 Databricks SQL Connector 进行连接，并将 connect_args 中的值传递给连接器。它支持以下认证类型：

个人访问令牌认证
OAuth 用户与机器 (U2M) 身份验证
OAuth 机器对机器 (M2M) 身份验证

个人访问令牌认证

若要创建 Azure Databricks 个人访问令牌，请按照为工作区用户创建个人访问令牌中的步骤作。

要对 SQLAlchemy 方言进行身份验证，请使用以下代码片段。此代码片段假定你已设置以下环境变量：

DATABRICKS_TOKEN，设置为 Azure Databricks 个人访问令牌。
DATABRICKS_SERVER_HOSTNAME设置为你的群集或 SQL 仓库的服务器主机名值。
DATABRICKS_HTTP_PATH，设置为你的群集或 SQL 仓库的 HTTP 路径值。
DATABRICKS_CATALOG，设置为 Unity Catalog 中的目标目录。
DATABRICKS_SCHEMA，设置为 Unity Catalog 中的目标架构（也称为数据库）。

若要设置环境变量，请参阅操作系统的文档。

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}"
)

# ...

可使用上述 engine 变量通过 Azure Databricks 计算资源连接到指定的目录和模式。

OAuth 用户-计算机 (U2M) 身份验证

OAuth U2M 认证使用实时人工登录和同意来认证目标 Azure Databricks 用户账户。有关 Databricks SQL Connector for Python 在各云上支持该功能的版本，请参见 OAuth 用户对机器（U2M）身份验证。

要用 OAuth U2M 认证 SQLAlchemy 方言，请使用以下代码片段。本代码片段假设你已设置 DATABRICKS_SERVER_HOSTNAME、DATABRICKS_HTTP_PATH、DATABRICKS_CATALOG 和 DATABRICKS_SCHEMA 环境变量，如个人访问令牌认证中所述。

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 机器对机器 (M2M) 身份验证

OAuth M2M 身份验证以 Azure Databricks 服务主体的身份进行身份验证。你还必须安装 Python 的 Databricks SDK（例如，通过运行 pip install databricks-sdk）。关于在各云平台上支持此功能的 Databricks SQL Connector for Python 和 SDK 版本，请参见 OAuth 机器对机器（M2M）身份验证。

要使用OAuth M2M认证，请执行以下操作：

在 Azure Databricks 工作区中创建 Azure Databricks 服务主体，并为该服务主体创建 OAuth 机密。请参阅使用 OAuth 授权服务主体访问 Azure Databricks。
授予该服务主体对你的集群或 SQL 仓库的访问权限。请参阅计算权限或管理 SQL 仓库。

要用OAuth M2M认证SQLAlchemy方言，请使用以下代码片段。除了 DATABRICKS_SERVER_HOSTNAME、DATABRICKS_HTTP_PATH、DATABRICKS_CATALOG 和 DATABRICKS_SCHEMA 环境变量外，此代码片段还假设你已设置以下环境变量：

DATABRICKS_CLIENT_ID，设置为服务主体的 UUID 或应用程序 ID 值。
DATABRICKS_CLIENT_SECRET，设置为服务主体的 OAuth 机密的机密值。

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

有关连接示例，请参阅此 example.py

SQLAlchemy v2

有关连接示例，请参阅以下部分和 GitHub 中的 sqlalchemy_example.py 文件。