AI Agent工具调用实战：如何高效获取数据库Schema与表结构

共计 2669 个字符，预计需要花费 7 分钟才能阅读完成。

1. 背景与痛点

在开发 AI Agent 时，经常需要让 Agent 理解数据库结构以便生成准确的 SQL 查询或进行数据分析。获取数据库 Schema 和表结构是最基础也最频繁的操作之一，但实际操作中会遇到几个典型问题：

• 连接复杂性 ：不同数据库(MySQL/PostgreSQL 等) 的元数据查询语法差异大
• 权限问题：生产环境需要严格控制数据库访问权限
• 性能瓶颈：当数据库包含大量表时，全量获取 Schema 可能超时
• 安全风险：直接拼接 SQL 字符串可能导致注入漏洞

2. 技术方案对比

直接查询方案

# 直接执行原生 SQL 查询（不推荐）cursor.execute("SHOW TABLES FROM mydb")

• 优点：简单直接
• 缺点：
• 需要处理不同数据库方言
• 缺乏错误处理和超时机制
• 存在 SQL 注入风险

封装工具方案

# 使用 ORM 工具封装（推荐）from db_toolkit import get_schema
schema = get_schema(db_url='postgresql://user:pass@host/db', timeout=10)

• 优点：
• 统一接口适配多种数据库
• 内置安全防护和超时控制
• 便于添加缓存等扩展功能
• 缺点：需要额外开发封装层

3. 核心实现

以下是使用 SQLAlchemy 实现的完整示例：

from sqlalchemy import create_engine, MetaData
from sqlalchemy.exc import SQLAlchemyError
import contextlib
from typing import Dict, List

@contextlib.contextmanager
def db_session(db_url: str):
    """数据库连接上下文管理器"""
    engine = None
    try:
        engine = create_engine(db_url, pool_pre_ping=True)
        conn = engine.connect()
        yield conn
    finally:
        if engine:
            engine.dispose()

def get_db_schema(db_url: str, timeout: int = 30) -> Dict[str, List[str]]:
    """
    获取数据库 schema 和表结构

    参数：db_url: 数据库连接字符串
        timeout: 超时时间(秒)

    返回：{'schemas': [], 'tables': {schema: [table1, table2]}}
    """result = {'schemas': [],'tables': {}}

    try:
        with db_session(db_url) as conn:
            # 设置执行超时
            conn.execute(f"SET statement_timeout = {timeout * 1000}")

            # 获取所有 schema
            schemas = conn.execute("SELECT schema_name FROM information_schema.schemata")
            result['schemas'] = [row[0] for row in schemas if not row[0].startswith('pg_')]

            # 获取各 schema 下的表
            for schema in result['schemas']:
                tables = conn.execute(
                    "SELECT table_name FROM information_schema.tables WHERE table_schema = %s",
                    (schema,)
                )
                result['tables'][schema] = [row[0] for row in tables]
    except SQLAlchemyError as e:
        raise RuntimeError(f"获取数据库结构失败: {str(e)}")

    return result

关键设计点：

1. 使用上下文管理器确保连接正确关闭
2. 参数化查询防止 SQL 注入
3. 显式设置 statement_timeout 避免长时间阻塞
4. 过滤系统 schema(pg_开头)

4. 性能与安全

性能优化

• 分页查询：当表数量超过 1000 时建议分页获取

  # 分页查询示例
  LIMIT = 100
  offset = 0
  while True:
      tables = conn.execute(
          "SELECT table_name FROM information_schema.tables"
          "WHERE table_schema = %s LIMIT %s OFFSET %s",
          (schema, LIMIT, offset)
      )
      if not tables.rowcount:
          break
      offset += LIMIT

• 并行查询：对多个 schema 可以并行查询

• 缓存结果：对不常变动的 schema 可以缓存 24 小时

安全措施

1. 永远不要拼接 SQL 字符串
2. 使用最小权限账户，只授予 information_schema 的读权限
3. 验证输入参数：

   # 检查 schema 名称合法性
   if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', schema):
       raise ValueError("Invalid schema name")

5. 生产环境避坑指南

权限管理

• 创建专用只读账号：

  CREATE ROLE schema_reader LOGIN PASSWORD 'secure_pwd';
  GRANT USAGE ON SCHEMA public TO schema_reader;
  GRANT SELECT ON ALL TABLES IN SCHEMA information_schema TO schema_reader;

大数据库处理

• 添加采样机制：当表超过 5000 个时只返回部分示例
• 支持按名称过滤：get_schema(filter='customer%')

缓存策略

from functools import lru_cache

@lru_cache(maxsize=32, ttl=3600)
def get_cached_schema(db_url: str) -> dict:
    return get_db_schema(db_url)

6. 总结与延伸

本文实现的工具已经可以满足基本需求，但还可以进一步扩展：

• 添加列详细信息获取功能
• 支持 NoSQL 数据库的 schema 发现
• 集成到 CI/CD 流程中自动校验数据库变更

建议在实际项目中根据具体需求逐步完善，核心是要保持接口的简单性和可靠性。完整的工具类实现可以参考 GitHub 上的 db-schema-toolkit 开源项目。

最后提醒：数据库结构获取虽然看起来简单，但在生产环境中要特别注意性能和安全性，建议所有查询都添加明确的超时限制，并且定期审计工具的使用情况。