如何在Databricks中为VectorSearchClient动态生成JWT以避免过期？

针对你遇到的Token 1小时过期导致服务端点无法查询的问题，以下是两种适配client credentials模式的动态Token生成方案：

方案一：自定义带Token自动刷新的VectorSearchClient包装类 #

通过封装原客户端，实现Token的延迟加载、过期自动刷新，并加入缓存避免频繁调用Token接口：

import constants as con
import requests
from databricks.vector_search.client import VectorSearchClient
import time

class AutoRefreshVectorSearchClient:
    def __init__(self, workspace_url):
        self.workspace_url = workspace_url
        self._client = None
        self._access_token = None
        self._token_expiry_time = 0  # 记录Token过期时间戳

    def _get_valid_token(self):
        # 提前5分钟刷新Token，避免网络延迟导致过期
        current_time = time.time()
        if self._access_token is None or current_time >= self._token_expiry_time - 300:
            CLIENT_ID = con.CLIENT_ID
            CLIENT_SECRET = con.CLIENT_SECRET
            token_endpoint_url = con.TOKEN_ENDPOINT_URL

            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:
                raise Exception(f"获取Token失败: {response.status_code} - {response.text}")
            
            token_info = response.json()
            self._access_token = token_info.get("access_token")
            # 计算过期时间（当前时间+接口返回的有效期秒数）
            expires_in = token_info.get("expires_in", 3600)
            self._token_expiry_time = current_time + expires_in
        return self._access_token

    def get_client(self):
        # Token刷新后重新初始化客户端
        token = self._get_valid_token()
        if self._client is None or self._client._config.personal_access_token != token:
            self._client = VectorSearchClient(
                workspace_url=self.workspace_url,
                personal_access_token=token
            )
        return self._client

# 使用方式
vsc_client = AutoRefreshVectorSearchClient(DATABRICKS_HOST)
# 每次调用业务方法前先获取有效客户端
vsc = vsc_client.get_client()

方案二：在MLflow模型的预测逻辑中动态管理Token #

如果是通过MLflow部署模型，可在模型的predict方法内每次请求前检查Token有效性，确保使用未过期的Token：

import mlflow
import constants as con
import requests
from databricks.vector_search.client import VectorSearchClient
import time

class VectorSearchModel(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.workspace_url = con.DATABRICKS_HOST
        self._client = None
        self._access_token = None
        self._token_expiry_time = 0

    def _refresh_token(self):
        CLIENT_ID = con.CLIENT_ID
        CLIENT_SECRET = con.CLIENT_SECRET
        token_endpoint_url = con.TOKEN_ENDPOINT_URL

        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:
            raise Exception(f"Token刷新失败: {response.status_code} - {response.text}")
        
        token_info = response.json()
        self._access_token = token_info["access_token"]
        self._token_expiry_time = time.time() + token_info.get("expires_in", 3600)
        # 重新初始化客户端
        self._client = VectorSearchClient(
            workspace_url=self.workspace_url,
            personal_access_token=self._access_token
        )

    def predict(self, context, model_input):
        # 检查Token是否需要刷新
        if time.time() >= self._token_expiry_time - 300 or self._client is None:
            self._refresh_token()
        
        # 替换为实际的向量检索逻辑
        results = self._client.get_index(index_name="your_index_name").similarity_search(query_text=model_input["query"])
        return results

# 注册模型时使用该类
mlflow.pyfunc.log_model(
    artifact_path="vector_search_model",
    python_model=VectorSearchModel()
)

关键注意点 #

• 加入提前5分钟刷新的逻辑，避免因网络延迟或Token刚好在请求过程中过期导致失败
• 缓存Token和过期时间，减少对Token端点的请求次数，避免触发限流
• 异常处理要明确，确保Token获取失败时能抛出清晰的错误信息

内容的提问来源于stack exchange，提问作者Agostino