# SQL Tests Gen Pro

**Repository Path**: wydhhh/sql-tests-gen-pro

## Basic Information

- **Project Name**: SQL Tests Gen Pro
- **Description**: 在大模型辅助下，为 postgresql 内核代码 commit diff 生成覆盖 test.
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-06-10
- **Last Updated**: 2026-06-24

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# PostgreSQL 覆盖率测试用例自动生成

## 背景

在 PostgreSQL 内核开发中，每个 commit 修改了 C 源码，需要编写 SQL 测试用例使改动的代码行能被 gcov 标记为"已执行"。人工做这件事需要逐个 commit 阅读 diff、理解代码路径、设计触发场景——commit 一多就很费时间。

这个项目最初就是一个简单脚本：读取每个 commit 的 diff，丢给 AI，让 AI 返回 SQL。跑起来之后，实际效果不理想，几个明显的问题暴露出来：

1. **AI 不知道具体要覆盖哪些行。** 把整段 diff 给 AI，它可能生成 SQL 命中了 diff 中无关的代码行，却没命中真正需要覆盖的行。
2. **生成的 SQL 语句数很可能超出 psql 单 session 的执行上限**，又没有机制在超限时做取舍。
3. **一次生成很难全覆盖**，但没有反馈渠道告诉 AI 哪里没覆盖到、应该怎么补。

这些问题在"把 diff 丢给 AI"这个简单模式下是没法解决的。后续的改动都是围绕这三个问题展开的。

## 项目做了什么

核心任务：给定一组 PostgreSQL 内核 commit，每个 commit 附带 `match_info` 标注了需要覆盖的 C 源码行（靶心行），为每个 commit 生成 SQL 用例，使这些靶心行在 gcov 评测中尽可能被命中。

### 整体流程

```
main.py
  │
  ├─ 加载 RAG 引擎（从 train.json 构建索引）
  ├─ 加载测试集（test_v3.json）
  │
  ├─ 主阶段：逐 commit 生成
  │   ├─ diff 特征提取 + 模块映射
  │   ├─ 靶心行解析（从 match_info 提取目标行）
  │   ├─ 代码分组与上下文分析（调用链、分支路径、源码上下文）
  │   ├─ RAG 检索（找到相似 commit 作为 few-shot）
  │   ├─ AI 生成（索要 case 数 = 靶心行数 + slack）
  │   └─ 保存 → 下一个 commit
  │
  ├─ 后处理
  │   ├─ fix_submission：清洗危险 SQL
  │   ├─ 空 commit 兜底（SELECT 1）
  │   └─ 优先级裁剪：超限时按"超出靶心行数"优先裁剪
  │
  └─ 补充阶段（supplement_rounds > 0）
      ├─ 每轮：评测 → 找未覆盖行 → 按 commit 分组 → AI 补写 → 保存
      ├─ 补充完成后：fix_submission 清洗
      └─ 最终评测
```

### 数据格式

**输入（test_v3.json / train.json）**是一个 commit 数组，每个 commit 的关键字段：

| 字段 | 说明 |
|------|------|
| `id` | commit 编号 |
| `subject` | commit 标题 |
| `email_body` | commit message 全文 |
| `patches` | 数组，每个元素含 `file`（源文件路径）、`function`（函数名）、`raw_diff`（unified diff 原文）、`before_code`（修改前代码片段） |
| `match_info` | 结构化标注，含 `patches[].blocks[].matched` 标记哪些 diff 行是靶心行，以及对应的 `line_no`、`code` |

**输出（submission.json）**是一个数组，每个元素只有两个字段：`id` 和 `generated_sql_tests`（XML 格式的 SQL 用例）。

## 项目结构

```
sql-tests-gen-pro/
├── main.py                    # 主入口：流程编排、后处理、裁剪、CLI
├── supplement.py              # 多轮补充：评测 → 找缺口 → AI 补充 → 保存
├── run_fix_submission.py      # 危险 SQL 清洗与过滤
├── config.py                  # 全参数配置，支持 CLI 覆盖
├── run_local_eval.sh          # 本地评测脚本（编译 PG + 执行 SQL + gcov）
├── main.log                   # 运行日志
├── data/
│   ├── train.json             # 训练集（RAG 检索来源）
│   ├── test_v3.json           # 测试集（待生成用例的 commit 列表）
│   └── example_submission.json# 输出格式示例
├── core/
│   ├── case_generator/
│   │   ├── ai_based.py        # AI 生成器（主阶段 + 补充阶段双 prompt）
│   │   ├── rule_based.py      # 规则模板兜底生成器
│   │   └── base.py            # 用例基类定义
│   ├── code_analyzer/
│   │   ├── diff_parser.py     # diff 特征提取（模块、CTE、JOIN 等）
│   │   └── module_mapper.py   # 文件 → 模块映射
│   ├── match_analyzer.py      # 靶心行解析（与评测脚本对齐）
│   └── rag_engine.py          # RAG 相似度检索（关键词加权）
├── scripts/
│   ├── evaluate.py            # 评测分析（gcov 匹配、行过滤、统计）
│   └── evaluate_coverage.sh   # 覆盖率评测辅助脚本
├── utils/
│   ├── io.py                  # JSON 读写、进度管理
│   └── xml_formatter.py       # 用例 XML 格式化
├── outputs/
    ├── submission.json        # 最终生成的用例文件
    └── local_eval/            # 各轮评测结果（按时间戳分目录）
├── zzz_temp_submission/       # 主阶段/补充阶段临时快照（gitignore）
├── zzz_fix_submission/        # 清洗后的临时 submission（gitignore）
└── zzz_history_submission/    # 历史 submission 存档，内含我们跑出来的高覆盖率结果（名字含覆盖行数）
```

## 各模块做了什么

下面逐个介绍每个模块的职责和实现方式。

### 靶心行解析 (`match_analyzer.py`)

init 脚本直接把 diff 原文丢给 AI，AI 不知道哪些行是目标。`match_info` 是每个 commit 附带的结构化标注，标出了 diff 中哪些行需要覆盖。`match_analyzer.py` 做了三层过滤：

1. **去重**：同一个 (文件, 行号) 只保留一次。
2. **跳过注释**：`/* ... */` 注释块和 `*` 开头的行不是可执行代码。
3. **跳过函数声明**：通过 C 类型关键字和函数名规则判断，排除纯声明行。

过滤逻辑与 `scripts/evaluate.py` 中的 `collect_matched_lines()` 保持一致，保证生成阶段看到的靶心行和评测阶段统计的是同一套行。

### 代码分组与上下文分析 (`match_analyzer.py` 后半部分)

靶心行在传递给 AI 之前做了以下加工：

- **按临近度分组**：同一文件中行号跨度不超过窗口大小的靶心行归为一组（窗口大小可在 config 中调整）。
- **查找外层函数**：从靶心行向上扫描，找到包含它的函数签名。
- **调用链追溯**：从目标函数出发，在当前文件和跨文件中搜索调用者，构建从外层入口到目标函数的调用链。
- **分支路径还原**：从函数入口正向扫描到靶心行，还原每一层 if/else/switch/for/while 分支条件，标注目标行在哪条分支内。
- **提取函数/文件注释**：把 PG 源码自带的文档注释提取出来，帮助 AI 理解函数用途。
- **附带源码上下文**：每组目标行前后各附带若干行源码，让 AI 看到完整上下文。

这些信息不做预分析结论，直接提供给 AI。

### diff 特征提取与模块映射 (`diff_parser.py` + `module_mapper.py`)

两个轻量辅助模块，用于日志输出：`diff_parser.py` 对 diff 做正则匹配，提取 9 个布尔特征（错误处理、CTE、窗口函数、子查询、JOIN、聚合、事务、索引、分区）；`module_mapper.py` 通过文件名匹配把源文件映射到 PG 内核模块（optimizer/executor/rewrite/commands/parser/access/utils）。不直接参与 AI 生成。

### RAG 相似度检索 (`rag_engine.py`)

init 脚本没有参考样例，AI 每次都是"从零开始"。实际上，相同模块的 commit 往往有类似的代码改动模式，对应的 SQL 写法也有参考价值。

`rag_engine.py` 对每个待生成 commit，从训练集（`train.json`，包含已完成的 commit 及其生成用例）中检索最相似的 top-k 个 commit，将它们的 diff、靶心行、已生成用例作为 few-shot 示例注入 prompt。

相似度计算方式：
- 将 commit 的 subject、email_body、diff、match_info 代码行拼接成文本。
- 取两个文本的单词交集，对 PG 内核关键词（如 join、hash、index、planner 等）赋予 5 倍权重，对模块名（optimizer、executor 等）赋予 8 倍权重，对包含 `(` 的词（可能是函数名）赋予 7 倍权重。
- 文件路径精确匹配额外加分（同一源文件的 commit 高度相关，加 25 分）。

### AI 生成 (`ai_based.py`)

这是负责构造 prompt 并调用 AI 的模块。主阶段和补充阶段各有一套独立的系统提示词和用户提示词。

**主阶段的 prompt 设计思路：**
- 系统提示词告诉 AI 它的角色、任务、运行环境、禁止使用的 SQL 类型、输出格式。
- 用户提示词包含：commit 元信息（subject、message、文件、函数）、完整 diff、靶心行列表、按代码分组的靶心行分析、参考样例。
- 向 AI 索要的 case 数量 = 靶心行数 + slack（或按路径分组数 + slack），多出来的余量是考虑到 AI 一次生成可能覆盖不全。

**补充阶段的 prompt 设计思路：**
- 系统提示词强调"已有用例覆盖了一部分，你补缺口"的基调。
- 用户提示词在信息结构上做了区分：先给全部靶心行和代码分组让 AI 看全貌，再给剩余未覆盖行让 AI 聚焦缺口，最后给已有用例（`existing_sql`）避免重复和冲突。
- 索要 case 数 = 剩余未覆盖行数 + slack。

模型使用 DeepSeek V4 Pro 的 thinking 模式（`reasoning_effort=max`，`temperature=0.2`），模型在生成 SQL 前会先通过思维链分析代码路径。

AI 的输出解析（`_parse` 方法）：
- 先剥离 markdown 代码块标记。
- 按空行切块，每块的首行作为描述，其余为 SQL。
- 对描述行剥离 `**bold**` 标记，不动 SQL 内容。

### 规则模板兜底 (`rule_based.py`) + 用例基类 (`base.py`)

`base.py` 定义了 `TestCase` 数据类（description + sql + source）和 `BaseGenerator` 抽象基类。

`rule_based.py` 是 AI 调用失败时的兜底生成器。它内置了按 PG 模块（optimizer/executor/rewrite/parser/commands/access/utils）和按 diff 特征（窗口函数、CTE、聚合、子查询、错误路径、索引、JOIN）分类的预制 SQL 模板。当 AI 返回空结果时，`main.py` 会调用 `rule_based.py` 生成一组模板用例，确保每个 commit 至少有一些 SQL 可以执行。

### 工具模块 (`utils/`)

- `io.py`：`load_json`/`save_json` 封装 JSON 读写；`load_progress` 和 `save_submission` 实现断点续跑，已完成的 commit 不会重复生成。
- `xml_formatter.py`：`XMLFormatter.format()` 把 `TestCase` 列表序列化为评测脚本能解析的 XML 格式（`<test_cases><test_case id="N"><description>...</description><sql>...</sql></test_case></test_cases>`）。

### 多轮补充 (`supplement.py`)

主阶段生成后，部分靶心行可能没被覆盖。补充阶段负责找出这些缺口并让 AI 补写用例。

流程（设 rounds = x，共跑 x+1 次评测）：
1. 对当前 submission 跑评测，得到每个 commit 的覆盖情况。
2. 找出未覆盖的靶心行，按 commit 分组。
3. 对每个仍有未覆盖行的 commit，调用 AI 生成补充用例（索要 case 数 = 未覆盖行数 + slack），追加到已有用例后面。
4. 保存快照，进入下一轮。
5. 所有轮次结束后，跑 fix_submission 清洗，再跑最终评测。

每轮补充和最终评测都会保存 submission 快照和评测结果到 `outputs/local_eval/` 下按时间戳命名的目录中。

### 语句数管理与裁剪 (`main.py` 后半部分)

项目的评测环境是单个 psql session 执行所有 SQL，实践中约 3000 条是稳定上限。超限会导致 session 崩溃、后续所有 case 无法执行。

裁剪策略是**按优先级**的：
- `priority = 当前 case 数 − 靶心行数`。超出靶心行数越多的 commit 优先级越高，越先被裁剪——因为多出来的 case 是"冗余"的，裁掉对覆盖率影响最小。
- 当所有 commit 优先级相同（都已经 ≤ 靶心行数），改为按 SQL 数量降序均摊裁剪。
- 每次裁剪后重新计算总数，超限就继续。

### 危险 SQL 过滤 (`run_fix_submission.py`)

AI 生成的 SQL 有时会包含不该出现的语句：shell 转义（`\!`）、文件 I/O（`\i`、`COPY ... FROM PROGRAM`）、权限提升（`SET ROLE`）、扩展加载（`CREATE EXTENSION`）、`pg_sleep`、`DO` 块、`VACUUM`/`ANALYZE` 等维护命令。这些在评测环境中要么无法执行（导致 psql 报错），要么有安全风险。

`run_fix_submission.py` 用 20+ 条正则表达式匹配这些危险模式，逐条检查每个 test_case 的 `<sql>` 内容，移除包含危险 SQL 的整条 case。清洗在 main.py 的主阶段结束后、补充阶段结束后各执行一次。

### 容错与工程化

- **断点续跑**：从已有 `submission.json` 恢复进度，已完成的 commit 跳过。
- **AI 失败回退**：AI 返回空时调用 `rule_based.py` 的模板生成器兜底。
- **空 commit 兜底**：清洗后无 SQL 的 commit 填入 `SELECT 1`。
- **评测失败不中断**：补充阶段评测脚本退出码非零时继续处理已有结果。
- **快照管理**：补充阶段每轮评测前复制 submission 快照到 `zzz_temp_submission/`，评测结果保存到 `outputs/local_eval/` 下按时间戳命名的目录。

## 使用方式

**前置条件：**
- 本地有 PostgreSQL 13.23 源码，放在 `postgresql-13.23/` 目录下，已用 `--enable-coverage` 编译
- Python 3.10+，`pip install openai`
- DeepSeek API key（设环境变量 `DEEPSEEK_API_KEY` 或在 `config.py` 中填写）

**运行：**

```bash
# 基本运行（主阶段 + 补充阶段）
python main.py

# 仅主阶段，不跑补充
python main.py --supplement-rounds 0

# 调整参数
python main.py --temperature 0.1 --supplement-cases-slack 3 --rag-top-k 3

# 从已有 submission 恢复继续（断点续跑）
python main.py --resume

# 只处理前 10 个 commit（调试用）
python main.py --limit 10
```

**评测脚本：**
- `run_local_eval.sh`：本地评测入口，负责编译 PG、创建 psql session、执行 SQL、收集 gcov 数据。由 `supplement.py` 在补充阶段自动调用。
- `scripts/evaluate.py`：评测分析，读取 gcov 输出，对每条靶心行做覆盖匹配（精确匹配 + 偏移容忍），计算 recall 和 PrecNF。
- `scripts/evaluate_coverage.sh`：`run_local_eval.sh` 内部调用的辅助脚本，处理 gcov 报告的生成。

## 评测指标

评测由 `scripts/evaluate.py` 完成，核心逻辑：

1. 从 submission JSON 中提取所有 `<sql>` 块，拼接成单个 SQL 文件。
2. 在 psql session 中执行该 SQL 文件。
3. 运行 gcov 收集覆盖率数据。
4. 对每个 commit 的每条靶心行，在 gcov 输出中查找匹配。匹配策略：
   - **精确匹配**：源码行号与 gcov 行号一致。
   - **偏移容忍**：源码行号在 gcov 中找不到时，在 ±5 行范围内搜索相同代码内容。
   - 都找不到则标记为 `gcov_not_found`。
5. 计算两个核心指标：
   - **recall**：已覆盖的靶心行数 / 总靶心行数。越高越好。
   - **PrecNF**（precision excluding not-found）：已覆盖的靶心行数 / (已覆盖 + gcov 中可执行但未覆盖的靶心行数)。分母排除了 gcov 中找不到的行（这些行无法判断覆盖状态）。

每条靶心行的匹配结果会记录在 `eval_result.json` 的 `line_details` 中，包含 `match_type`、`gcov_count`、`gcov_executable`、`covered` 等字段，方便排查未覆盖原因。

## 关键配置

`config.py` 中所有参数支持 CLI 覆盖（`--参数名 值`）。主要参数：

| 参数 | 默认值 | 说明 |
|------|--------|------|
| `MAX_SQL_STATEMENTS` | 3000 | 全局 SQL 语句数上限，超限触发优先级裁剪 |
| `CASE_SLACK` | 3 | 主阶段索要 case 数 = 靶心行数 + slack |
| `SLOT_BY_PATH_GROUP` | True | True=按路径分组数分配槽位，False=按靶心行数分配 |
| `CASE_SLACK_GROUP` | 5 | 按路径分组时的 slack 值 |
| `SUPPLEMENT_ROUNDS` | 1 | 补充轮次，x 轮 = 跑 x+1 次评测 |
| `SUPPLEMENT_CASES_SLACK` | 2 | 补充阶段索要 case 数 = 剩余行数 + slack |
| `RAG_TOP_K` | 2 | RAG 检索返回 top-k 个相似 commit |
| `CONTEXT_LINES` | 25 | 靶心行上下文 ± 行数 |
| `GROUP_WINDOW_SIZE` | 20 | 路径分组窗口大小 |
| `AI_MODE` | thinking | thinking=思维链模式，standard=标准模式 |
| `AI_TEMPERATURE` | 0.2 | 生成温度 |
| `AI_REASONING_EFFORT` | max | 思考强度（high/max），仅 thinking 模式 |

补充阶段可独立配置模型、温度、模式等参数（`SUPPLEMENT_MODEL` 等），设为 `None` 时自动回退到主阶段配置。

## 最终结果
最终在当前 config 默认值配置下，

初始阶段生成得分：
```log
[INFO] 从 submission 提取 SQL: /home/wyd-wsl/sql-tests-gen-pro/outputs/submission.json
{
  "input_records": 50,
  "sql_cases": 316,
  "sql_statements": 1999,
  "output_dir": "/home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/parsed_sql/submission",
  "merged_sql": "/home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/all_gen_sql.sql"
}
[INFO] 开始本地评测: suppl_r0_20260624_160036
[INFO] === Step 0: 清理环境 ===
[INFO] 删除旧 gcda（不要删 gcno！）
[INFO] 详细日志写入: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/coverage_workspace/run.log
Local evaluation summary
result_file: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/eval_result.json
n_items: 50
total_meaningful_added: 460
total_matched: 200
total_not_found: 2
total_covered: 113
global_recall: 0.4348
global_precision: 0.565
global_precision_excl_ctrl: 0.5873
global_precision_excl_not_found: 0.5707

[INFO] 覆盖率 HTML: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/coverage_workspace/report/index.html
[INFO] SQL 输出日志: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/coverage_workspace/psql_output.log
[INFO] 运行日志: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/coverage_workspace/run.log
[INFO] 生成分数文件: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/scores.txt
PrecNF: 0.570700
efficiency: 0.250671
result: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_r0_20260624_160036/eval_result.json
========================================
  最终分数 PrecNF: 0.570700
  效率分数 efficiency: 0.250671
========================================
```

补充阶段一轮后，生成得分：
```log
[INFO] 从 submission 提取 SQL: /home/wyd-wsl/sql-tests-gen-pro/outputs/submission.json
{
  "input_records": 50,
  "sql_cases": 408,
  "sql_statements": 2673,
  "output_dir": "/home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/parsed_sql/submission",
  "merged_sql": "/home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/all_gen_sql.sql"
}
[INFO] 开始本地评测: suppl_final_20260624_191046
[INFO] === Step 0: 清理环境 ===
[INFO] 删除旧 gcda（不要删 gcno！）
[INFO] 详细日志写入: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/coverage_workspace/run.log
Local evaluation summary
result_file: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/eval_result.json
n_items: 50
total_meaningful_added: 460
total_matched: 200
total_not_found: 2
total_covered: 121
global_recall: 0.4348
global_precision: 0.605
global_precision_excl_ctrl: 0.6243
global_precision_excl_not_found: 0.6111

[INFO] 覆盖率 HTML: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/coverage_workspace/report/index.html
[INFO] SQL 输出日志: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/coverage_workspace/psql_output.log
[INFO] 运行日志: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/coverage_workspace/run.log
[INFO] 生成分数文件: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/scores.txt
PrecNF: 0.611100
efficiency: 0.259124
result: /home/wyd-wsl/sql-tests-gen-pro/outputs/local_eval/suppl_final_20260624_191046/eval_result.json
========================================
  最终分数 PrecNF: 0.611100
  效率分数 efficiency: 0.259124
========================================
```

token 共计消费人民币约二十元（深度思考模式导致的开销较高）。
结果保留在 zzz_history_submission/ 目录下，文件名含对应覆盖行数。


## 已知局限

1. **单 session 执行限制**：所有 SQL 在单个 psql session 中顺序执行，一个 case 的 side effect（如残留表、锁、GUC 设置）可能影响后续 case。目前通过 prompt 要求每个 case 自包含（CREATE 前 DROP）来缓解，但无法完全消除。
2. **gcov 行号漂移**：编译优化可能导致 gcov 报告的行号与源码行号不完全对齐，部分靶心行在 gcov 输出中找不到（`match_type=gcov_not_found`），这部分行无法判断是否覆盖。
3. **AI 生成质量不稳定**：同一 prompt 多次调用可能得到不同质量的输出，尤其在复杂代码路径上。thinking 模式能提升稳定性但不能完全消除波动。且后续补充阶段补充能力乏力，难以突破当前生成能力的瓶颈。
4. **不生成预期输出验证**：当前只关注代码行是否被执行，不验证执行结果的正确性。一条 SQL 可能触发了目标代码行但产生了错误结果，这在当前评测中不会被发现。
5. **RAG 依赖训练集质量**：相似度检索完全依赖 `train.json` 中的已有 commit 和用例，训练集覆盖不全时参考价值有限。