CrewAI DatabricksQueryTool 详解：让 AI Agent 用 SQL 查询 Databricks 工作区表

本文以 CrewAI 工具库 crewai-tools 中的 DatabricksQueryTool 为对象，系统讲解该工具如何让 AI Agent 通过 SQL 查询 Databricks 工作区表并拿到结果集：包括两种认证方式的配置、实例化时的默认参数（default_catalog / default_schema / default_warehouse_id）、运行时五个输入参数（含 row_limit 自动注入 LIMIT 的行为），并结合 源码实现 剖析从语句提交、状态轮询到结果格式化的完整执行链路，读完即可在自己的 Crew 中安全落地该工具。

工具定位

DatabricksQueryTool 位于 databricks_query_tool 模块，继承自 crewai.tools.BaseTool，其设计目标是"让 Agent 能用 SQL 访问 Databricks 中存储的数据"：Agent 只需给出 SQL 语句（以及可选的 catalog / schema / warehouse 参数），工具就会通过 databricks-sdk 的 Statement Execution API 提交查询、等待执行完成、把结果集整理成对齐的文本表格返回给 Agent。

该工具已通过 crewai_tools 包导出（crewai_tools 顶层与 crewai_tools.tools 子包均导出了 DatabricksQueryTool），因此可以直接 from crewai_tools import DatabricksQueryTool 使用。

源码中工具自身的注册信息（见 databricks_query_tool.py）：

• name = "Databricks SQL Query"；
• description 为 "Execute SQL queries against Databricks workspace tables and return the results. Provide a 'query' parameter with the SQL query to execute."——这段描述会直接作为提示词上下文交给 LLM 决策是否调用该工具；
• args_schema = DatabricksQueryToolSchema，由 Pydantic 模型定义合法入参；
• package_dependencies = ["databricks-sdk"]，用于在依赖缺失时给出安装提示。

安装

README 给出的 pip 安装方式为：

pip install 'crewai_tools' 'databricks-sdk'

在 pyproject.toml 中，crewai-tools 提供了名为 databricks-sdk 的可选依赖组，要求 databricks-sdk>=0.46.0，因此也可以按 官方工具文档 推荐的方式一次性安装：

uv add 'crewai-tools[databricks-sdk]'

两种方式等价，区别只是是否把版本约束交给 crewai-tools 管理。注意 databricks-sdk 是运行时按需导入的（见下文 workspace_client 属性），如果环境中缺少该包，工具在真正发起查询前就会抛出带安装提示的 ImportError："`databricks-sdk` package not found, please run `uv add databricks-sdk"`。

认证：两种提供凭据的方式

工具不接收 host / token 等认证参数，而是完全依赖环境变量。源码 init 在构造时即调用 _validate_credentials() 做前置校验（L143-L154）：

1. 使用 Databricks CLI profile：设置 DATABRICKS_CONFIG_PROFILE 环境变量指向你的 profile 名称（适用于 ~/.databrickscfg 中已配置 profile 的场景）；
2. 直接使用凭据：同时设置 DATABRICKS_HOST 和 DATABRICKS_TOKEN。

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="dapi1234567890abcdef"

校验逻辑是"二选一"：只要环境变量中存在 DATABRICKS_CONFIG_PROFILE，或者 DATABRICKS_HOST 与 DATABRICKS_TOKEN 同时存在，即通过；否则在工具实例化阶段就抛出 ValueError，而不是等到查询时才失败。源码中 env_vars 字段（L102-L120）也以结构化方式声明了这三个变量及用途描述，供框架侧展示。

个人访问令牌（PAT）与 workspace 地址的获取方式：在 Databricks workspace 的 User Settings → Developer 中创建 PAT 并查看 host 信息。

实例化：三个默认参数

from crewai_tools import DatabricksQueryTool

# 基本用法（不指定默认值）
databricks_tool = DatabricksQueryTool()

# 指定 catalog、schema、warehouse 默认值
databricks_tool = DatabricksQueryTool(
    default_catalog="my_catalog",
    default_schema="my_schema",
    default_warehouse_id="warehouse_id"
)

三个构造参数（L95-L98）全部可选：

构造参数	类型	说明
default_catalog	str | None	默认 catalog（如 main）
default_schema	str | None	默认 schema（如 default）
default_warehouse_id	str | None	默认 SQL warehouse ID

它们的语义是"运行时的兜底值"：_run() 中按 kwargs.get("catalog") or self.default_catalog 的顺序取值（L226-L230），即每次调用显式传入的参数优先于实例化时的默认值。

需要特别注意：warehouse_id 虽然"可选"，但必须来自参数或默认值二者之一，否则 _run() 会直接返回错误信息 "SQL warehouse ID must be provided either as a parameter or as a default."（L245-L246）。这是因为 Statement Execution API 必须明确指定在哪个 SQL warehouse 上执行语句，Databricks 侧没有可推断的默认值。而 catalog / schema 可以缺失——缺失时 Databricks 会按会话自身的默认上下文解析表名。

输入参数与 row_limit 的自动 LIMIT 注入

工具对 LLM 暴露的入参由 Pydantic 模型 DatabricksQueryToolSchema（L37-L70）定义：

参数	必填	默认值	说明
query	是	—	要执行的 SQL；空字符串或纯空白会被校验器拒绝（Query cannot be empty）
catalog	否	实例 default_catalog	Databricks catalog 名
db_schema	否	实例 default_schema	Databricks schema 名
warehouse_id	否（但必须可解析）	实例 default_warehouse_id	SQL warehouse ID
row_limit	否	1000	返回的最大行数

这里有一个容易踩坑的命名差异：模块 README 将 schema 参数写作 schema，但源码中实际字段名是 db_schema（schema 是 Pydantic 模型的保留属性名，不能直接用）。官方工具文档 的 Parameters 一节同样写的是 db_schema，因此 Agent 调用时应使用 db_schema。

row_limit 的行为值得单独说明。校验器中的 validate_input（L59-L70）在参数进入执行逻辑之前做了一次查询改写：

if self.row_limit and "limit" not in self.query.lower():
    self.query = f"{self.query.rstrip(';')} LIMIT {self.row_limit};"

即：当 SQL 中不包含大小写不敏感的 limit 字样时，工具会自动去掉尾部分号并追加 LIMIT {row_limit}（默认 1000）。这意味着：

• 未写 LIMIT 的 SELECT * 类查询天然被限制在 1000 行以内，避免 Agent 一次拉回超大结果集；
• 如果你的 SQL 已经自带 LIMIT，则保持原样、不会被覆盖；
• 判断依据是字符串匹配 limit，因此查询文本中出现含该子串的其他内容也会跳过自动注入——对常规 DML/DDL 语句无影响。

执行流程：提交、轮询与状态机

_run()（L209 起）的完整执行链路可以从源码结构看分为四段：

1. 惰性创建 WorkspaceClient。 workspace_client 属性（L156-L168）在首次访问时才 from databricks.sdk import WorkspaceClient 并实例化，后续复用同一实例。WorkspaceClient() 不传任何参数，正是靠上文的三个环境变量完成认证。

2. 提交语句。 认证信息确认、参数校验通过后，工具构造一个 ExecutionContext（TypedDict，只含 catalog / schema 两个键），调用：

execution = self.workspace_client.statement_execution.execute_statement(
    warehouse_id=warehouse_id, statement=query, **context
)
statement_id = execution.statement_id

提交失败时返回 "Error starting query execution: ..."；拿不到 statement_id 时返回 "Failed to retrieve statement ID after execution."。

3. 轮询等待终态。 Databricks 语句执行是异步的，工具以固定节奏轮询 statement.get_statement(statement_id)：

• 超时窗口 timeout = 300 秒（5 分钟），每次间隔 time.sleep(2) 轮询一次；
• 状态判定基于 result.status.state，兼容字符串与枚举两种形态：包含 SUCCEEDED 则跳出轮询；包含 FAILED 则提取错误信息并返回 "Query execution failed: {error_info}"；包含 CANCELED 则返回 "Query was canceled"；
• 错误信息提取做了双重兼容，优先取 status.error.message，其次 status.error.error_message，再不行退化为 str(error)；
• 轮询请求本身抛异常时不会立即失败，连续 3 次以上才返回 "Error checking query status: ..."，容忍瞬时网络抖动；
• 若 5 分钟内未达终态，返回 "Query timed out after 5 minutes (last state: ...)" 并附上最后观察到的状态。

4. 结果解析。 拿到 SUCCEEDED 的 statement 后，工具从 result.manifest.schema.columns 取列名，从 result.result.data_array 按 chunk 迭代行数据。从源码结构看，这一段做了非常防御式的解析：兼容 data_array 与 data 两种载荷、处理 chunk 内行结构异常的情况（例如值被逐字符拆开时按启发式规则重建行边界）、把超出列名的多余值归入动态列 Column_{i}，最后统一归一化——保证每行都含有全部列、缺失值补 None。对 DDL 等无结果集的成功语句，则直接返回 "Query executed successfully (no results to display)"。

结果格式化：Agent 友好的文本表格

解析出的行数据最终由 _format_results()（L170-L207）渲染为固定宽度对齐的文本表格：

• 按每列数据的最大显示宽度计算列宽；
• 首行表头 + -+- 分隔线 + 数据行，列之间以 | 分隔；
• None 值统一显示为 NULL；
• 末尾附 (N rows returned) 行计数；
• 空结果的三种情况分别有明确文案：无行、"有行但无列"、"有列但无数据"。

这种纯文本表格对 LLM 是最友好的返回形态——结构清晰、token 开销可控、无需 Agent 再解析 JSON。所有异常路径也不会把原始堆栈抛给 Agent，而是返回带 traceback.format_exc() 详情的错误文本（如 "Error executing Databricks query: ..."），既保留了调试信息，又保证工具调用本身不中断 Crew 流程。

在 CrewAI Agent 中使用

README 的最小集成方式是给 Agent 挂载工具：

from crewai_tools import DatabricksQueryTool

databricks_tool = DatabricksQueryTool(
    default_catalog="my_catalog",
    default_schema="my_schema",
    default_warehouse_id="warehouse_id"
)
# 在 Agent 定义中挂载
# Agent(config=..., allow_delegation=False, tools=[databricks_tool])

官方工具文档 给出了一个完整的 Crew 示例，可直接复制运行（前提是环境变量已按认证一节配置）：

from crewai import Agent, Task, Crew
from crewai_tools import DatabricksQueryTool

tool = DatabricksQueryTool(
    default_catalog="main",
    default_schema="default",
)

agent = Agent(
    role="Data Analyst",
    goal="Query Databricks",
    tools=[tool],
    verbose=True,
)

task = Task(
    description="SELECT * FROM my_table LIMIT 10",
    expected_output="10 rows",
    agent=agent,
)

crew = Crew(
    agents=[agent],
    tasks=[task],
    verbose=True,
)
result = crew.kickoff()

print(result)

也可以不经过 Agent，直接调用工具本身验证连通性（源码 docstring 中的示例）：

>>> tool = DatabricksQueryTool(default_warehouse_id="<your_warehouse_id>")
>>> results = tool.run(query="SELECT * FROM my_table LIMIT 10")

错误处理与实战建议

结合 官方文档的 Error handling & tips 一节 与源码行为，整理如下：

• 认证错误：确认 DATABRICKS_HOST 以 https:// 开头且 token 有效；使用 CLI profile 时确认 DATABRICKS_CONFIG_PROFILE 指向的 profile 存在。
• 权限问题：确保 token 对目标 SQL warehouse 与 schema 有访问权限；warehouse_id 必须与 token 所在工作区匹配。
• 长查询限制：工具端硬编码了 5 分钟轮询超时，超时即返回 "Query timed out"。在 Agent 循环中应避免发起长耗时查询，建议在 SQL 中主动加过滤条件和 LIMIT（row_limit 的自动注入只是保底，不是性能手段）。
• db_schema 而非 schema：给 Agent 的任务描述或自定义提示中如涉及该参数，建议使用实际字段名 db_schema。
• 错误信息可读性：所有失败路径都返回自然语言错误串（SQL 执行失败、状态检查失败、结果处理失败均附详情），Agent 通常能据此自行修正 SQL 重试；如果你需要程序化判断失败，可对这些返回串做前缀匹配（如 Query execution failed、Error executing Databricks query）。

关键文件索引

内容	路径
工具模块 README（本篇主体文档）	lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/README.md
工具完整实现	lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py
官方工具文档（Crew 示例与排错建议）	docs/edge/en/tools/search-research/databricks-query-tool.mdx
databricks-sdk 可选依赖声明（>=0.46.0）	lib/crewai-tools/pyproject.toml
包级导出	lib/crewai-tools/src/crewai_tools/tools/init.py