通过托管标识从 python 连接到 Databricks

Question

我设法使用以下代码 sn-p 从 python 连接到 Databricks：

from databricks import sql

connection = sql.connect(
  server_hostname='<server-hostname>',
  http_path='<http-path>',
  access_token='<personal-access-token>')

cursor = connection.cursor()

cursor.execute('SELECT * FROM <database-name>.<table-name> LIMIT 2')

result = cursor.fetchall()

for row in result:
  print(row)

cursor.close()

这个sn-p来自官方documentation，如你所见，它需要server_hostname、http_path和access_token。我的问题是，我可以在没有access_token 的情况下验证自己的身份吗？也许使用托管标识，因为这两种技术都来自 Microsoft？

Answer 1

您可以使用 AAD 令牌代替个人访问令牌。生成它的最简单方法是使用 Microsoft 的 azure-identity library。

以下是为服务主体生成 AAD 令牌的代码示例：

from databricks import sql
from azure.identity import ClientSecretCredential
import os

tenant_id = '...'
client_id = '...'
client_secret = os.environ['SP_SECRET']
csc = ClientSecretCredential(tenant_id, client_id, client_secret)

# important!
dbx_scope = '2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default'
# this will be used as access token
aad_token = csc.get_token(dbx_scope).token

也应该可以使用托管标识 - 只需使用 ManagedIdentityCredential class 代替。我相信您也可以使用该库生成用户的 AAD 令牌。如果您不想有额外的依赖，您可以通过 REST API 生成它，例如 code。

管理员需要使用corresponding REST API 或Databricks Terraform provider 将要使用的标识添加到Azure Databricks 工作区。并且还授予访问 Databricks 集群或 SQL 端点的权限 - 这也可以通过 REST API 或 Terraform 实现。

需要考虑的一件事是 AAD 令牌过期（多快 - 这取决于您使用的是什么：服务主体或托管标识），因此您可能需要定期重新创建连接。