机制解析:从设计提案到 parser 源码实现)
TiDB 中的 AST 还原 SQL 文本Restore机制解析从设计提案到 parser 源码实现【免费下载链接】tidbTiDB is built for agentic workloads that grow unpredictably, with ACID guarantees and native support for transactions, analytics, and vector search. No data silos. No noisy neighbors. No infrastructure ceiling.项目地址: https://gitcode.com/GitHub_Trending/ti/tidb导读当 TiDB 在 SQL 解析层拿到一棵 AST抽象语法树后许多高级特性如CREATE VIEW的列展开、Plan Cache、SQL Binding 归一化等都要求能把 AST 节点重新拼接回等价的 SQL 文本。本文以 TiDB 官方设计提案 docs/design/2018-11-29-ast-to-sql-text.md 为主线结合当前仓库pkg/parser中真实落地的实现代码系统讲解ast.Node.Restore(ctx)接口、RestoreFlags/RestoreCtx设计、Writer 辅助方法与典型调用链。读完你既能理解这套还原机制的演进脉络也能在自己的 SQL 工具链中复用其接口设计与格式控制方法。为什么需要从 AST 还原 SQL 文本TiDB 首先使用 parser 将 SQL 语句解析为一棵ast.Node组成的语法树随后优化器、执行器会在树上做变换。很多功能需要把可能被改写过的AST 节点重新还原成 SQL 文本最典型的场景正是提案中给出的例子CREATE VIEW v AS SELECT * FROM t;在真正落库前TiDB 需要把视图定义中的SELECT *展开成显式列名例如SELECT test.t.col0, test.t.col1 FROM test.t再重新生成完整的、可执行的CREATE VIEW语句文本。要做到这一点就要求任何 AST 节点都能被还原为 SQL 文本——即select子节点被展开替换后其父级、祖级节点乃至整个语句仍可重新序列化。在引入 Restore 机制之前ast.Node上只存在Text()/SetText()方法parser 在解析过程中通过SetText()记录节点对应的原始文本片段之后可通过Text()取回。但该机制并不完整——只有根节点的Text()能可靠工作一旦 AST 被修改过节点记录的原始文本便与新树不一致。设计文档因此明确指出实现本提案时不应依赖node.Text()同理也不应依赖exprNode.Format()。每个 AST 节点本质是一棵嵌套树其子节点与 SQL 片段一一对应这正是 Restore 递归还原的基础。以CREATE USER语句为例CreateUserStmt是一个语句级节点它聚合了UserSpec、用户身份User、认证选项AuthOpt等子节点整体构成如图所示的层次结构提案核心给 ast.Node 增加 Restore() 方法设计提案给出的接口定义如下该方法现已实际进入pkg/parser/ast包type Node interface { // Restore AST to SQL text and append them to ctx. // return error when the AST is invalid. Restore(ctx *RestoreCtx) error // ... }在 pkg/parser/ast/ast.go#L28 中可以看到当前Node接口的完整形态——Restore已是每个 AST 节点必须实现的能力与AcceptVisitor 访问、Text、SetText等方法并列type Node interface { // Restore returns the sql text from ast tree Restore(ctx *format.RestoreCtx) error Accept(v Visitor) (node Node, proceed bool) AcceptInPlace(v InPlaceVisitor) (proceed bool) Text() string OriginalText() string SetText(enc charset.Encoding, text string) // ... }注意Restore并不直接返回字符串而是接收一个*format.RestoreCtx并将还原出的文本追加写入其内部的 writer。这样调用方可以自由选择输出目标bytes.Buffer、strings.Builder、os.Stdout等也便于在遍历过程中控制格式。RestoreFlags控制输出格式的九面旗标为了让还原出的 SQL 能以不同格式输出提案引入RestoreFlags。最初的九面旗标构成如下互斥分组同一分组内靠左的旗标优先级更高互斥分组作用RestoreStringSingleQuotes/RestoreStringDoubleQuotes字符串字面量用单引号或双引号包裹RestoreStringEscapeBackslash是否对反斜杠做转义处理RestoreKeyWordUppercase/RestoreKeyWordLowercase关键字SELECT、CREATE等大写或小写RestoreNameUppercase/RestoreNameLowercase标识符库名、表名等大写或小写RestoreNameDoubleQuotes/RestoreNameBackQuotes标识符用双引号或反引号包裹这些定义可以在 pkg/parser/format/format.go#L209-L250 中找到其枚举声明与设计文档完全一致文档中左侧位置旗标优先级更高的注释也原样保留在源码中。例如在字符串引号分组中RestoreStringSingleQuotes排在前因此当两个旗标同时被置位时按单引号处理。RestoreCtx携带旗标与写入目标还原过程的上下文由RestoreCtx承载。提案中的定义是// RestoreCtx is Restore context to hold flags and writer type RestoreCtx struct { Flags RestoreFlags In io.Writer } const DefaultRestoreFlags RestoreStringSingleQuotes | RestoreKeyWordUppercase | RestoreNameBackQuotes默认旗标组合意味着字符串用单引号、关键字用大写、标识符用反引号。当前实现中RestoreCtx的结构已扩展得更丰富见 pkg/parser/format/format.go#L378-L407type RestoreCtx struct { Flags RestoreFlags In RestoreWriter // io.Writer io.StringWriter DefaultDB string ParentBinaryOp int // 供表达式还原判断优先级、决定括号去留 ParentBinarySide int InUnaryOperation bool CTERestorer // 记录/查询 CTE 名称 }建议始终通过构造函数创建实例它会把DefaultDB初始化为空串func NewRestoreCtx(flags RestoreFlags, in RestoreWriter) *RestoreCtx { return RestoreCtx{Flags: flags, In: in, DefaultDB: } }五个 Writer 辅助方法还原输出的基本原语Restore实现体本身不直接操作字符串拼接细节而是调用RestoreCtx提供的语义化 Writer 方法让格式控制集中收敛WriteKeyWord(keyWord)按旗标将关键字转为大写或小写后写入见 format.go#L416-L424WriteString(str)按旗标选择单/双引号包裹字符串字面量并处理//\的转义见 format.go#L453-L469WriteName(name)按旗标对标识符做大小写转换、并用反引号或双引号包裹内部对反引号做 双重转义见 format.go#L473-L494WritePlain(text)/WritePlainf(format, args...)原样写入不做任何转换见 format.go#L497-L504此外还有用于 TiDB 扩展语法的WriteWithSpecialComments/WriteKeyWordWithSpecialComments可在还原结果中生成/*T![feature] ... */形式的特殊注释。之所以把关键字/标识符/字符串分开处理是为了让同一棵 AST 在不同输出需求下例如大写规范输出、不带任何引号的迁移脚本无需改动节点实现。还原原理自顶向下递归、按层拼接AST 是一棵子节点即 SQL 片段的树因此还原算法非常朴素从根节点开始逐层调用各子节点的Restore()按语法次序把子节点输出与字面量关键字、符号、空格拼接起来。设计文档以如下语句为例SELECT column0 FROM table0 UNION SELECT column1 FROM table1 WHERE a 1其语法树顶层是一个UnionStmt下面挂两个UnionSel子 SELECT第二个 SELECT 内部又细分出FieldList列名、TableRefs表引用、ExprNodeWHERE a 1表达式等子结构。还原过程正是沿着这棵树的边逐层下行拼接每个节点只负责自己那一层的拼装子句细节交给子节点递归完成。这种子节点可插拔的设计让改写子节点后重新还原整句成为可能——这正是视图列展开、SQL 归一化等场景的前提。兼容性约束还原结果只需AST 等价SQL 文本与 AST 是一对多关系同一棵 AST 可以对应无数种写法关键字大小写、空白、括号冗余、字符串引号风格不同等因此不可能也不必还原出与原始输入逐字符相同的文本。提案确立的验收标准是由原始 SQL 语句解析出的 AST与由还原出的 SQL 语句再次解析得到的 AST两者应当相等。这也解释了为何需要RestoreFlags提供多种输出风格——不同消费方对格式的要求不同但对AST 等价的要求完全一致。以默认旗标还原时CREATE DATABASE db1这类语句会被稳定输出为规范形式从而可作为 SQL 指纹/归一化文本参与缓存命中判断。实现示例从提案代码到今天仓库中的真实实现设计文档选取了ast.CreateDatabaseStmt与ast.DropDatabaseStmt作为首批示例。回到 2018 年时该代码位于独立的pingcap/parser仓库如今 parser 已经作为子模块内嵌在 TiDB 主仓库的pkg/parser目录中这些Restore实现可以在 pkg/parser/ast/ddl.go 里直接读到。先看子节点DatabaseOption.Restore它按选项类型分发用WriteKeyWord/WritePlain还原CHARACTER SET、COLLATE等数据库选项ddl.go#L102-L140其核心逻辑与提案代码保持一致// Restore implements Node interface. func (n *DatabaseOption) Restore(ctx *format.RestoreCtx) error { switch n.Tp { case DatabaseOptionCharset: ctx.WriteKeyWord(CHARACTER SET) ctx.WritePlain( ) ctx.WritePlain(n.Value) case DatabaseOptionCollate: ctx.WriteKeyWord(COLLATE) ctx.WritePlain( ) ctx.WritePlain(n.Value) case DatabaseOptionEncryption: ctx.WriteKeyWord(ENCRYPTION) ctx.WritePlain( ) ctx.WriteString(n.Value) case DatabaseOptionPlacementPolicy: placementOpt : PlacementOption{ Tp: PlacementOptionPolicy, UintValue: n.UintValue, StrValue: n.Value, } return placementOpt.Restore(ctx) // ... default: return errors.Errorf(invalid DatabaseOptionType: %d, n.Tp) } return nil }可以看到随着 TiDB 演进源码在提案示例基础上又补充了ENCRYPTION、PLACEMENT POLICY、SET TIFLASH REPLICA等新选项分支。这也说明Restore是一个随语法能力持续增长的方法集合。父节点CreateDatabaseStmt.Restore则展示了关键字 条件分支 子节点递归 错误标注的通用骨架ddl.go#L153-L167// Restore implements Node interface. func (n *CreateDatabaseStmt) Restore(ctx *format.RestoreCtx) error { ctx.WriteKeyWord(CREATE DATABASE ) if n.IfNotExists { ctx.WriteKeyWord(IF NOT EXISTS ) } ctx.WriteName(n.Name.O) for i, option : range n.Options { ctx.WritePlain( ) err : option.Restore(ctx) if err ! nil { return errors.Annotatef(err, An error occurred while splicing CreateDatabaseStmt DatabaseOption: [%v], i) } } return nil }设计文档同时给出了DropDatabaseStmt的还原写法作为对照同样遵循先关键字、再可选子句、后名称的顺序// Restore implements Node interface. func (n *DropDatabaseStmt) Restore(ctx *RestoreCtx) error { ctx.WriteKeyWord(DROP DATABASE ) if n.IfExists { ctx.WriteKeyWord(IF EXISTS ) } ctx.WriteName(n.Name) return nil }与视图场景直接相关的 CreateViewStmt回到文章开头那个CREATE VIEW场景在 pkg/parser/ast/ddl.go#L1674-L1743 中CreateViewStmt结构体用Select StmtNode字段保存视图查询的 AST其Restore()在输出ALGORITHM、DEFINER、SQL SECURITY、视图名与列清单后通过n.Select.Restore(ctx)递归还原查询体最后追加WITH ... CHECK OPTION。视图定义里SELECT 子节点可被展开替换后整体再序列化的能力正是由这套递归接口提供的。提案给出的实现注意点不要依赖exprNode.Format()旧 Format 机制并非为可逆还原设计不要依赖node.Text()它只能忠实反映解析时记录的原文本无法表达被改写后的 AST。这些约定在今天的Node接口文档注释中仍能看到影响——Restore从追加写入的上下文出发天然避免了这两个陷阱。落地现状Restore 在 TiDB 各模块中的典型调用Restore 机制如今贯穿 TiDB 的多个子系统。从调用方式看统一模式是先构造strings.Builder/bytes.Buffer再format.NewRestoreCtx(flags, buf)最后node.Restore(ctx)语义检查与报错信息构造在 pkg/planner/core/preprocess.go#L2106-L2124 中当CAST表达式的小数位/精度非法时会先用NewRestoreCtx(format.DefaultRestoreFlags, buf)把表达式还原成 SQL 文本再拼进ErrMBiggerThanD、ErrTooBigPrecision等错误消息让用户看到是哪段表达式写错了非预备计划缓存Non-Prepared Plan Cache在 pkg/planner/core/plan_cache_param.go#L44 中使用RestoreForNonPrepPlanCache | RestoreStringWithoutCharset | RestoreStringSingleQuotes | RestoreNameBackQuotes组合旗标对 SQL 归一化作为缓存键的一部分SQL Binding 归一化pkg/bindinfo相关代码同样依赖 Restore 生成规范化文本见 pkg/bindinfo/binding.go此外表达式还原还通过ParentBinaryOp、RestoreSkipRedundantParentheses等机制在保留语义的前提下智能省略冗余括号服务于规范化路径。作为对照仅用Text()无法支撑上述场景因为它们消费的都是经过改写/需要规范化的 AST 而非原始输入。RestoreFlags 的后续演进最初九面旗标之外当前 pkg/parser/format/format.go#L231-L249 已新增多面旗标读者在阅读代码或自行调用时可留意旗标用途RestoreSpacesAroundBinaryOperation/RestoreBracketAroundBinaryOperation二元运算两侧空格 / 强制加括号RestoreStringWithoutCharset/RestoreStringWithoutDefaultCharset省略字符串上的 charset 前缀或 DEFAULT charsetRestoreTiDBSpecialComment将 TiDB 扩展语法包进/*T!...*/特殊注释SkipPlacementRuleForRestore还原时跳过 placement 相关子句RestoreWithTTLEnableOff还原 TTL 表时强制TTL_ENABLEOFFRestoreWithoutSchemaName/RestoreWithoutTableName还原时省略库名 / 表名按标识符层级裁剪输出RestoreForNonPrepPlanCache非预备计划缓存场景的归一化还原RestoreBracketAroundBetweenExpr为BETWEEN表达式补括号RestoreSkipRedundantParentheses规范化路径中省略不影响语义的冗余括号工程实施与演进节奏设计文档将整个改造划分为四个阶段推进考虑到若干ast.Node之间存在依赖例如语句节点依赖表达式节点、DDL 节点依赖名称节点子任务按依赖顺序分批实现避免一次性全量改造。第一批落地了CreateDatabaseStmt/DropDatabaseStmt等基础 DDL 语句作为样板随后逐步覆盖到 DML、表达式、函数调用乃至 TTL、placement、CTE 等全部语法家族。小结从 2018 年这篇设计提案到今天的pkg/parser从 AST 还原 SQL 文本已经成为 TiDB 的一项基础设施能力接口层面Node.Restore(ctx *format.RestoreCtx) error让每类节点都能递归序列化自身格式层面RestoreFlagsRestoreCtx以互斥分组、优先级语义和 Writer 辅助方法把输出风格与语法拼装解耦正确性层面以还原文本二次解析出的 AST 与原 AST 相等为准则规避了 AST 与文本一对多带来的歧义应用层面视图定义生成、表达式报错回显、非预备计划缓存、SQL Binding 归一化等场景都建立在这一机制之上。对于想在自己的 SQL 工具链中做 AST 改写与回写的开发者这套设计互斥格式分组 上下文写入 递归还原 AST 等价校验是可直接借鉴的成熟范式。相关代码与文档索引设计提案原文docs/design/2018-11-29-ast-to-sql-text.mdAST 节点Node接口定义pkg/parser/ast/ast.go#L28RestoreFlags与RestoreCtx、Writer 辅助方法pkg/parser/format/format.go#L209CreateDatabaseStmt/DatabaseOption/CreateViewStmt的 Restore 实现pkg/parser/ast/ddl.go#L102还原调用示例报错回显 / 非预备计划缓存pkg/planner/core/preprocess.go#L2114、pkg/planner/core/plan_cache_param.go#L44【免费下载链接】tidbTiDB is built for agentic workloads that grow unpredictably, with ACID guarantees and native support for transactions, analytics, and vector search. No data silos. No noisy neighbors. No infrastructure ceiling.项目地址: https://gitcode.com/GitHub_Trending/ti/tidb创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考