适用于 Python 的 Databricks SQL 连接器

用于 Python 的 Databricks SQL 连接器是一个 Python 库，可用于使用 Python 代码在 Azure Databricks 全用途计算和 Databricks SQL 仓库上运行 SQL 命令。 相比类似的 Python 库（如 pyodbc），适用于 Python 的 Databricks SQL 连接器更易于设置和使用。 此库遵循 PEP 249 - Python 数据库 API 规范 v2.0。

用于 Python 的 Databricks SQL 连接器还支持 Azure Databricks SQLAlchemy 方言，但必须安装它才能使用这些功能。 请参阅将 SQLAlchemy 与 Azure Databricks 配合使用。

要求

• 运行 Python 3.8 及更高版本的开发计算机。
• Databricks 建议使用 Python 虚拟环境，例如 python 随附的 venv 提供的环境。 虚拟环境有助于确保同时使用正确版本的 Python 和适用于 Python 的 Databricks SQL 连接器。 设置和使用虚拟环境不在本文的讨论范围之内。 有关详细信息，请参阅创建虚拟环境。
• 现有的 全用途计算 或 SQL 仓库。

开始

1. 安装用于 Python 的 Databricks SQL 连接器。 PyArrow 是用于 Python 的 Databricks SQL 连接器的可选依赖项，默认情况下在连接器版本 4.0.0 及更高版本中未安装。 如果未安装 PyArrow，CloudFetch 和其他 Apache Arrow 功能等功能不可用，这可能会影响大量数据的性能。

   • 若要安装精简连接器，请使用：

     pip install databricks-sql-connector
     

   • 若要安装完整的连接器（包括 PyArrow），请使用：

     pip install databricks-sql-connector[pyarrow]
     

2. 为要使用的通用计算或 SQL 仓库收集以下信息：

   通用计算

   • 所有用途计算的服务器主机名。 可以从 选项卡中的服务器>值获取此值，以便进行全用途计算。

   • 全用途计算的 HTTP 路径。 可以从 选项卡中的 >值获取此值，以便进行全用途计算。

   注意

   SQL 连接器不支持连接到作业计算集群。

   SQL 仓库

   • SQL 仓库的服务器主机名。 从 SQL 仓库的“连接详细信息”选项卡的“服务器主机名”值中可以获取此主机名。
   • SQL 仓库的 HTTP 路径。 从 SQL 仓库的“连接详细信息”选项卡的“HTTP 路径”值中可以获取此路径。

身份验证

适用于 Python 的 Databricks SQL 连接器支持以下 Azure Databricks 身份验证类型：

• Databricks 个人访问令牌身份验证

• OAuth 计算机到计算机 (M2M) 身份验证

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

适用于 Python 的 Databricks SQL 连接器尚不支持以下 Azure Databricks 身份验证类型：

• 使用 Azure 托管标识进行身份验证
• 使用 MS Entra 服务主体进行身份验证
• 使用 Azure CLI 进行身份验证

Databricks 个人访问令牌身份验证

要将适用于 Python 的 Databricks SQL 连接器与 Azure Databricks 个人访问令牌身份验证配合使用，你必须先创建一个 Azure Databricks 个人访问令牌。 为此，请遵循 为工作区用户创建个人访问令牌中的步骤。

要对适用于 Python 的 Databricks SQL 连接器进行身份验证，可使用以下代码片段。 此代码片段假定你已设置以下环境变量：

• DATABRICKS_SERVER_HOSTNAME将其设置为用于通用计算或 SQL 仓库的服务器主机名值。
• DATABRICKS_HTTP_PATH，设置为用于所有用途的计算或 SQL 仓库的 HTTP 路径 值。
• DATABRICKS_TOKEN，设置为 Azure Databricks 个人访问令牌。

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

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:
# ...

OAuth 计算机到计算机 (M2M) 身份验证

适用于 Python 的 Databricks SQL 连接器版本 2.7.0 及更高版本支持 OAuth 计算机到计算机 (M2M) 身份验证。 还必须安装用于 Python 0.18.0 或更高版本的 Databricks SDK（例如通过运行 pip install databricks-sdk 或 python -m pip install databricks-sdk）。

要将 Databricks SQL Connector for Python 与 OAuth M2M 身份验证配合使用，必须执行以下操作：

1. 在 Azure Databricks 工作区中创建 Azure Databricks 服务主体，并为该服务主体创建 OAuth 机密。

   若要创建服务主体及其 OAuth 机密，请参阅 使用 OAuth 授权服务主体访问 Azure Databricks。 记下服务主体的 UUID 或应用程序 ID 值，以及服务主体的 OAuth 机密的机密值。

2. 为该服务主体提供对通用计算或仓库的访问权限。

   若要向服务主体授予对通用计算或仓库的访问权限，请参阅 “计算权限 ”或 “管理 SQL 仓库”。

要对适用于 Python 的 Databricks SQL 连接器进行身份验证，可使用以下代码片段。 此代码片段假定你已设置以下环境变量：

• DATABRICKS_SERVER_HOSTNAME 设置为用于所有用途计算或 SQL 仓库 的服务器主机名 值。
• DATABRICKS_HTTP_PATH，设置为用于所有用途的计算或 SQL 仓库的 HTTP 路径 值。
• DATABRICKS_CLIENT_ID，设置为服务主体的 UUID 或应用程序 ID 值。
• DATABRICKS_CLIENT_SECRET，设置为服务主体的 OAuth 机密的“机密”值。

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

from databricks.sdk.core import Config, oauth_service_principal
from databricks import sql
import os

server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME")

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)

with sql.connect(server_hostname      = server_hostname,
                 http_path            = os.getenv("DATABRICKS_HTTP_PATH"),
                 credentials_provider = credential_provider) as connection:
# ...

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

适用于 Python 的 Databricks SQL 连接器版本 2.7.0 及更高版本支持 OAuth 用户到计算机 (U2M) 身份验证。 还必须安装用于 Python 0.19.0 或更高版本的 Databricks SDK（例如通过运行 pip install databricks-sdk 或 python -m pip install databricks-sdk）。

要使用 OAuth U2M 身份验证对适用于 Python 的 Databricks SQL 连接器进行身份验证，可使用以下代码片段。 OAuth U2M 身份验证使用实时人工登录和同意对目标 Azure Databricks 用户帐户进行身份验证。 此代码片段假定你已设置以下环境变量：

• 将 DATABRICKS_SERVER_HOSTNAME 设置为计算或 SQL 仓库的 服务器主机名 值，用于所有用途。
• 将 DATABRICKS_HTTP_PATH 设置为用于通用计算或 SQL 仓库的 HTTP 路径 值。

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

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 auth_type       = "databricks-oauth") as connection:
# ...

示例

以下代码示例演示如何使用用于 Python 的 Databricks SQL 连接器来查询和插入数据、查询元数据、管理游标和连接、管理 Unity 目录中的文件以及配置日志记录。

这些代码示例从以下环境变量中检索它们的 server_hostname、http_path 和 access_token 连接变量值：

• DATABRICKS_SERVER_HOSTNAME，表示要求中的“服务器主机名”值。
• DATABRICKS_HTTP_PATH，表示要求中的“HTTP 路径”值。
• DATABRICKS_TOKEN，表示要求中的访问令牌。

设置用户代理

下面的代码示例演示如何设置 User-Agent 应用程序 product_name 进行使用情况跟踪。

from databricks import sql
import os

with sql.connect(server_hostname   = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path         = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token      = os.getenv("DATABRICKS_TOKEN"),
                 user_agent_entry = "product_name") as connection:
  with connection.cursor() as cursor:
    cursor.execute("SELECT 1 + 1")
    result = cursor.fetchall()

    for row in result:
      print(row)

查询数据

下面的代码示例演示如何调用用于 Python 的 Databricks SQL 连接器，以便在通用计算或 SQL 仓库上运行基本 SQL 命令。 此命令返回 trips 目录的 samples 架构中的nyctaxi 表中的前两行。

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:

  with connection.cursor() as cursor:
    cursor.execute("SELECT * FROM samples.nyctaxi.trips LIMIT 2")
    result = cursor.fetchall()

    for row in result:
      print(row)

查询标记

以下示例演示如何将键值标记附加到 SQL 查询以用于跟踪和分析目的。 查询标记显示在 system.query.history 表中。

from databricks import sql
import os

with sql.connect(
    server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
    http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
    access_token    = os.getenv("DATABRICKS_TOKEN"),
    session_configuration = {
        'query_tags': 'team:engineering,dashboard:abc123,env:prod'
    }
) as connection:
    with connection.cursor() as cursor:
        cursor.execute("SELECT * FROM samples.nyctaxi.trips LIMIT 2")
        result = cursor.fetchall()
        # Query is now tagged and trackable in system.query.history
        for row in result:
            print(row)

插入数据

以下示例演示如何插入少量数据（数千行）：

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:

  with connection.cursor() as cursor:
    cursor.execute("CREATE TABLE IF NOT EXISTS squares (x int, x_squared int)")

    squares = [(i, i * i) for i in range(100)]
    values = ",".join([f"({x}, {y})" for (x, y) in squares])

    cursor.execute(f"INSERT INTO squares VALUES {values}")

    cursor.execute("SELECT * FROM squares LIMIT 10")

    result = cursor.fetchall()

    for row in result:
      print(row)

要插入大量数据，请先将数据上传到云存储，然后执行 COPY INTO 命令。

查询元数据

可通过一些专用方法检索元数据。 以下示例检索有关示例表中的列的元数据：

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:

  with connection.cursor() as cursor:
    cursor.columns(schema_name="default", table_name="squares")
    print(cursor.fetchall())

管理游标和连接

最好关闭不再使用的任何连接和游标。 这会释放 Azure Databricks 全用途计算和 Databricks SQL 仓库上的资源。

可以使用上下文管理器（在前面示例中使用的 with 语法）来管理资源，或显式调用 close：

from databricks import sql
import os

connection = sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                         http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                         access_token    = os.getenv("DATABRICKS_TOKEN"))

cursor = connection.cursor()

cursor.execute("SELECT * from range(10)")
print(cursor.fetchall())

cursor.close()
connection.close()

管理 Unity Catalog 卷中的文件

Databricks SQL 连接器可让你将本地文件写入 Unity Catalog 卷、从卷下载文件以及从卷中删除文件，如下例所示：

from databricks import sql
import os

# For writing local files to volumes and downloading files from volumes,
# you must set the staging_allowed_local_path argument to the path to the
# local folder that contains the files to be written or downloaded.
# For deleting files in volumes, you must also specify the
# staging_allowed_local_path argument, but its value is ignored,
# so in that case its value can be set for example to an empty string.
with sql.connect(server_hostname            = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path                  = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token               = os.getenv("DATABRICKS_TOKEN"),
                 staging_allowed_local_path = "/tmp/") as connection:

  with connection.cursor() as cursor:

    # Write a local file to the specified path in a volume.
    # Specify OVERWRITE to overwrite any existing file in that path.
    cursor.execute(
      "PUT '/tmp/my-data.csv' INTO '/Volumes/main/default/my-volume/my-data.csv' OVERWRITE"
    )

    # Download a file from the specified path in a volume.
    cursor.execute(
      "GET '/Volumes/main/default/my-volume/my-data.csv' TO '/tmp/my-downloaded-data.csv'"
    )

    # Delete a file from the specified path in a volume.
    cursor.execute(
      "REMOVE '/Volumes/main/default/my-volume/my-data.csv'"
    )

配置日志记录

Databricks SQL 连接器使用 Python 的标准日志记录模块。 以下示例配置日志记录级别并生成调试日志：

from databricks import sql
import os, logging

logging.getLogger("databricks.sql").setLevel(logging.DEBUG)
logging.basicConfig(filename = "results.log",
                    level    = logging.DEBUG)

connection = sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                         http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                         access_token    = os.getenv("DATABRICKS_TOKEN"))

cursor = connection.cursor()

cursor.execute("SELECT * from range(10)")

result = cursor.fetchall()

for row in result:
   logging.debug(row)

cursor.close()
connection.close()

测试

若要测试代码，请使用 Python 测试框架，例如 pytest。 若要在模拟条件下测试代码，而不调用 Azure Databricks REST API 终结点或更改 Azure Databricks 帐户或工作区的状态，可以使用 Python 模拟库（如 unittest.mock）。

例如，如果以下名为“helpers.py”的文件，其中包含使用 Azure Databricks 个人访问令牌返回到 Azure Databricks 工作区的连接的 get_connection_personal_access_token 函数，以及使用连接从 select_nyctaxi_trips 目录的 trips 架构中的 samples 表中获取指定数目的数据行的 nyctaxi 函数：

# helpers.py

from databricks import sql
from databricks.sql.client import Connection, List, Row, Cursor

def get_connection_personal_access_token(
  server_hostname: str,
  http_path: str,
  access_token: str
) -> Connection:
  return sql.connect(
    server_hostname = server_hostname,
    http_path = http_path,
    access_token = access_token
  )

def select_nyctaxi_trips(
  connection: Connection,
  num_rows: int
) -> List[Row]:
  cursor: Cursor = connection.cursor()
  cursor.execute(f"SELECT * FROM samples.nyctaxi.trips LIMIT {num_rows}")
  result: List[Row] = cursor.fetchall()
  return result

另外，假设以下名为 main.py 的文件调用 get_connection_personal_access_token 和 select_nyctaxi_trips 函数：

# main.py

from databricks.sql.client import Connection, List, Row
import os
from helpers import get_connection_personal_access_token, select_nyctaxi_trips

connection: Connection = get_connection_personal_access_token(
  server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
  http_path = os.getenv("DATABRICKS_HTTP_PATH"),
  access_token = os.getenv("DATABRICKS_TOKEN")
)

rows: List[Row] = select_nyctaxi_trips(
  connection = connection,
  num_rows = 2
)

for row in rows:
  print(row)

以下名为 test_helpers.py 的文件测试 select_nyctaxi_trips 函数是否返回预期的响应。 此测试将模拟 Connection 对象，而不是创建与目标工作区的真实连接。 该测试还模拟一些符合真实数据中的架构和值的数据。 该测试通过模拟连接返回模拟数据，然后检查其中一个模拟数据行的值是否与预期值匹配。

# test_helpers.py

import pytest
from databricks.sql.client import Connection, List, Row
from datetime import datetime
from helpers import select_nyctaxi_trips
from unittest.mock import create_autospec

@pytest.fixture
def mock_data() -> List[Row]:
  return [
    Row(
      tpep_pickup_datetime = datetime(2016, 2, 14, 16, 52, 13),
      tpep_dropoff_datetime = datetime(2016, 2, 14, 17, 16, 4),
      trip_distance = 4.94,
      fare_amount = 19.0,
      pickup_zip = 10282,
      dropoff_zip = 10171
    ),
    Row(
      tpep_pickup_datetime = datetime(2016, 2, 4, 18, 44, 19),
      tpep_dropoff_datetime = datetime(2016, 2, 4, 18, 46),
      trip_distance = 0.28,
      fare_amount = 3.5,
      pickup_zip = 10110,
      dropoff_zip = 10110
    )
  ]

def test_select_nyctaxi_trips(mock_data: List[Row]):
  # Create a mock Connection.
  mock_connection = create_autospec(Connection)

  # Set the mock Connection's cursor().fetchall() to the mock data.
  mock_connection.cursor().fetchall.return_value = mock_data

  # Call the real function with the mock Connection.
  response: List[Row] = select_nyctaxi_trips(
    connection = mock_connection,
    num_rows = 2)

  # Check the value of one of the mocked data row's columns.
  assert response[1].fare_amount == 3.5

select_nyctaxi_trips 函数包含 SELECT 语句，因此不会更改 trips 表的状态，在此示例中并不是一定需要模拟。 但是，模拟让你能够快速运行测试，而无需等待与工作区建立实际连接。 此外，通过模拟，可以多次针对可能更改表状态的函数运行模拟测试，例如 INSERT INTO、UPDATE 和 DELETE FROM。

API 参考

本部分包含 databricks-sql-connector 软件包的 API 参考。 请参阅 Python 包索引 （PyPI） 中的 databricks-sql-connector 。

模块

包 databricks.sql 的 databricks-sql-connector 模块包含初始化与 SQL 仓库的连接的方法。

connect 方法

初始化与 SQL 仓库的连接。 返回 Connection 对象。

Connection 类

表示与计算或 SQL 仓库的连接。

方法

该 Connection 类提供以下方法。

Cursor 类

表示一种遍历数据记录的机制。

若要创建 Cursor 对象，请调用 cursorConnection'类的方法。

特性

选定的 Cursor 属性包括：

方法

选定的 Cursor 方法包括：

Row 类

行类是类似于元组的数据结构，表示 SQL 查询结果中的单个结果行。 如果行包含名为 "my_column" 的列，则你可以通过 "my_column" 访问 row 的 row.my_column 字段。 还可以使用数字索引来访问字段，例如 row[0]。 如果不允许将列名称用作属性方法名称（例如，它以数字开头），则可以将字段作为 row["1_my_column"] 来访问。

从版本 1.0 开始

选定的 Row 方法包括：

方法

类型转换

下表将 Apache Spark SQL 数据类型映射到其 Python 数据类型等效项。

疑难解答

tokenAuthWrapperInvalidAccessToken: Invalid access token 消息

问题：运行代码时看到类似于 的消息。

可能的原因：传递给 的值不是有效的 Azure Databricks 个人访问令牌。

建议的解决方法：检查传递给 的值是否正确，然后重试。

gaierror(8, 'nodename nor servname provided, or not known') 消息

问题：运行代码时看到类似于 的消息。

可能的原因：传递给 的值不是正确的主机名。

建议的解决方法：检查传递给 的值是否正确，然后重试。

有关查找服务器主机名的详细信息，请参阅获取 Azure Databricks 计算资源的连接详细信息。

IpAclError 消息

问题：在运行代码期间，尝试使用 Azure Databricks 笔记本上的连接器时看到消息 。

可能的原因：可能为 Azure Databricks 工作区启用了 IP 允许列表。 使用 IP 允许列表时，默认不允许从 Spark 群集连接回到控制平面。

建议的解决方法：请管理员将计算平面子网添加到 IP 允许列表。

其他资源

有关详细信息，请参阅：

• GitHub 上的适用于 Python 的 Databricks SQL 连接器存储库
• 数据类型
• Python 网站上的内置类型（适用于 bool、bytearray、float、int 和 str）
• Python 网站上的日期/时间（适用于 datetime.date 和 datatime.datetime）
• Python 网站上的十进制（适用于 decimal.Decimal）
• Python 网站上的内置常数（适用于 NoneType）