如何将 JWT 传递到 Databricks 中的 VectorSearchClient 而不会过期？

Question

我希望使用以下逻辑在 databricks 中创建一个搜索向量检索器

vsc = VectorSearchClient(
    workspace_url=DATABRICKS_HOST,
    personal_access_token=get_access_token(),
    )

get_access_token() 函数如下所示

import constants as con

def get_access_token():
    """
    Creates a databricks access token from a client id and client secret.

    :return: databricks access token
    """
    CLIENT_ID = con.CLIENT_ID
    CLIENT_SECRET = con.CLIENT_SECRET

    # Token endpoint URL
    token_endpoint_url = con.TOKEN_ENDPOINT_URL

    # Prepare the data payload
    data = {"grant_type": "client_credentials", "scope": "all-apis"}

    # Make the POST request with Basic Authentication
    response = requests.post(
    token_endpoint_url, data=data, auth=(CLIENT_ID, CLIENT_SECRET)
    )

    # Check if the request was successful
    if response.status_code == 200:
        # Parse the JSON response
        token_info = response.json()
        access_token = token_info.get("access_token")
        return access_token
    else:
        print("Error:", response.status_code)
        print("Response:", response.text)
        return None

但是，当我将检索器传递到 MLFlow 实验并在 Databricks 中创建服务端点时，令牌将在 1 小时后过期，并且我无法再查询该端点。

有没有动态创建token的解决方案？对于这个项目，不幸的是，我需要这样的令牌设置，并且不能使用密钥。

Answer 1

根据这个文档，令牌将在一小时后过期，过期后需要创建另一个令牌。

因此，另一种方法是以编程方式在全局存储过期时间并在每次使用时进行检查，如果过期则生成新令牌，否则使用当前令牌。

代码

import requests
import time

_access_token = None
_token_expiration = 0

def get_access_token():
    """
    Creates a databricks access token from a client id and client secret.

    :return: databricks access token
    """

    global _access_token, _token_expiration

    if _access_token and time.time() < _token_expiration:
        print("using current token")
        return _access_token

    CLIENT_ID = "<client_id>"
    CLIENT_SECRET = "<seceret>"
    workspace_url = "<url>"

    
    token_endpoint_url = f'{workspace_url}/oidc/v1/token'

    
    data = {"grant_type": "client_credentials", "scope": "all-apis"}

    
    response = requests.post(
    token_endpoint_url, data=data, auth=(CLIENT_ID, CLIENT_SECRET)
    )

    if response.status_code == 200:
        token_info = response.json()
        _access_token = token_info.get("access_token")
        expires_in = token_info.get("expires_in", 3600)
        _token_expiration = time.time() + expires_in - 60
        print("Generated new token")
        return _access_token
    else:
        print("Error:", response.status_code)
        print("Response:", response.text)
        return None

通过上面的代码，你可以在每次token过期时动态生成token。