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

Question

我设法使用以下代码片段从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()

此代码片段来自官方documentation，如您所见，它需要server_hostname、http_path和access_token。我这里的问题是，我可以在没有access_token的情况下进行身份验证吗？也许可以使用托管身份，因为这两种技术都来自Microsoft？

Answer 1

您可以使用AAD令牌而不是个人访问令牌。生成它的最简单的方法是使用微软的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令牌过期(有多快-这取决于您使用的是服务主体还是托管身份)，因此您可能需要定期重新创建连接。