ddlkit 0.1.0__py3-none-any.whl

ddlkit/__init__.py

@@ -0,0 +1,24 @@
+"""ddlkit —— 多方言 DDL 文本解析器（无损保真）。
+
+只做一件事：**DDL 文本 -> 结构化数据**。不含规则引擎。
+"""
+from __future__ import annotations
+
+from .api import DIALECT_ALIASES, parse_ddl, parse_file, resolve_dialect
+from .model import Column, Constraint, Index, ParseResult, SourceRef, Table
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Column",
+    "Constraint",
+    "DIALECT_ALIASES",
+    "Index",
+    "ParseResult",
+    "SourceRef",
+    "Table",
+    "__version__",
+    "parse_ddl",
+    "parse_file",
+    "resolve_dialect",
+]

ddlkit/api.py

@@ -0,0 +1,110 @@
+"""对外入口。
+
+调用方只需要知道两件事：**文本从哪来** 和 **它是哪个方言**。
+包内部负责：编码探测 -> 导出格式识别 -> 语句切分 -> 词法 -> 结构提取 -> 注释回填。
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from . import comments as _comments
+from . import encoding as _encoding
+from . import source as _source
+from .extract import extract_index, extract_table
+from .lexer import lex
+from .model import ParseResult, SourceRef, Table
+
+#: 业务库名 -> sqlglot 方言名。新增库只需在这里加一行。
+DIALECT_ALIASES: dict[str, str] = {
+    "dm": "oracle",
+    "dameng": "oracle",
+    "达梦": "oracle",
+    "ck": "clickhouse",
+    "clickhouse": "clickhouse",
+    "ob": "mysql",
+    "ob_mysql": "mysql",
+    "oceanbase": "mysql",
+    "oceanbase_mysql": "mysql",
+    "hive": "hive",
+    "hive2": "hive",
+}
+
+
+def resolve_dialect(name: str) -> str:
+    """把业务库名映射成 sqlglot 方言名；已是方言名则原样返回。"""
+    key = name.strip().lower()
+    return DIALECT_ALIASES.get(key, key)
+
+
+def parse_ddl(
+    text: str | bytes,
+    dialect: str,
+    *,
+    filename: str | None = None,
+    encoding: str | None = None,
+    fmt: str | None = None,
+) -> ParseResult:
+    """解析一段 DDL 文本（可含多库/多语句）。"""
+    if isinstance(text, bytes):
+        text, used_encoding = _encoding.decode(text, encoding)
+    else:
+        used_encoding = encoding or "utf-8"
+
+    resolved = resolve_dialect(dialect)
+    fmt = fmt or _source.detect(text)
+
+    result = ParseResult()
+    statements = _source.split(text, fmt)
+
+    tables: list[Table] = []
+    by_name: dict[str, Table] = {}
+    pending_indexes: list[tuple[str, object, int]] = []
+    comment_on_texts: list[str] = []
+
+    for stmt in statements:
+        if stmt.kind == "other":
+            continue
+        lexed = lex(stmt.sql, resolved)
+        src = SourceRef(path=filename, line=stmt.line, encoding=used_encoding, fmt=fmt)
+
+        if stmt.kind == "create_table":
+            table = extract_table(lexed, resolved, src)
+            if table is None:
+                result.unsupported.append(f"line {stmt.line}: {stmt.sql[:80]}")
+                continue
+            tables.append(table)
+            if table.schema:
+                by_name[f"{table.schema}.{table.name}".upper()] = table
+            by_name.setdefault(table.name.upper(), table)
+        elif stmt.kind == "create_index":
+            parsed = extract_index(lexed, src)
+            if parsed is not None:
+                pending_indexes.append((parsed[0], parsed[1], stmt.line))
+        elif stmt.kind == "comment_on":
+            comment_on_texts.append(stmt.sql)
+
+    for target, index, line in pending_indexes:
+        table = by_name.get(target.upper())
+        if table is None:
+            result.warnings.append(f"line {line}: 索引 {index.name!r} 找不到目标表 {target!r}")
+            continue
+        table.indexes.append(index)
+
+    if comment_on_texts:
+        _comments.backfill(comment_on_texts, tables)
+
+    result.tables = tables
+    return result
+
+
+def parse_file(
+    path: str | Path,
+    dialect: str,
+    *,
+    encoding: str | None = None,
+    fmt: str | None = None,
+) -> ParseResult:
+    """解析一个 DDL 导出文件。"""
+    p = Path(path)
+    text, used = _encoding.decode(p.read_bytes(), encoding)
+    return parse_ddl(text, dialect, filename=p.name, encoding=used, fmt=fmt)

ddlkit/comments.py

@@ -0,0 +1,72 @@
+"""``COMMENT ON`` 语句回填。
+
+Oracle 系的库（达梦）把注释写成**独立语句**：
+
+    COMMENT ON TABLE  "S"."T"   IS '表注释'
+    COMMENT ON COLUMN "S"."T"."C" IS '列注释'
+
+只解析单条 CREATE TABLE 会漏掉全部注释，导致"注释必填"类规则大面积误报。
+该语法极规整，正则可 100% 覆盖，比赌解析器对 COMMENT 语句的支持更稳。
+"""
+from __future__ import annotations
+
+import re
+
+from .model import Table
+
+COMMENT_ON = re.compile(
+    r"""COMMENT\s+ON\s+(?P<kind>TABLE|COLUMN)\s+
+        (?P<target>
+            "[^"]+"(?:\s*\.\s*"[^"]+")*
+          | `[^`]+`(?:\s*\.\s*`[^`]+`)*
+          | [A-Za-z_][\w$#]*(?:\s*\.\s*[A-Za-z_][\w$#]*)*
+        )
+        \s+IS\s+'(?P<text>(?:[^']|'')*)'""",
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_NAME_SPLIT = re.compile(r"\s*\.\s*")
+
+
+def _parts(target: str) -> list[str]:
+    return [p.strip('"`').upper() for p in _NAME_SPLIT.split(target.strip())]
+
+
+def _build_index(tables: list[Table]) -> dict[tuple[str, ...], Table]:
+    index: dict[tuple[str, ...], Table] = {}
+    for t in tables:
+        if t.schema:
+            index[(t.schema.upper(), t.name.upper())] = t
+        index[(t.name.upper(),)] = t
+    return index
+
+
+def backfill(statements: list[str], tables: list[Table]) -> int:
+    """把 COMMENT ON 回填到表/列上，返回成功回填的条数。"""
+    index = _build_index(tables)
+    hits = 0
+    for stmt in statements:
+        for m in COMMENT_ON.finditer(stmt):
+            kind = m.group("kind").upper()
+            segs = _parts(m.group("target"))
+            text = m.group("text").replace("''", "'")
+            if not segs:
+                continue
+
+            if kind == "TABLE":
+                table = index.get(tuple(segs[-2:])) or index.get((segs[-1],))
+                if table is not None:
+                    table.comment = text
+                    hits += 1
+            else:
+                if len(segs) < 2:
+                    continue
+                table = index.get(tuple(segs[-3:-1])) or index.get((segs[-2],))
+                if table is None:
+                    continue
+                for col in table.columns:
+                    if col.name.upper() == segs[-1]:
+                        col.comment = text
+                        hits += 1
+                        break
+    return hits

ddlkit/encoding.py

@@ -0,0 +1,34 @@
+"""编码探测。
+
+实测：同一批导出文件里 UTF-8 与 GB18030 混存（OB 的 test.sql 是 GB18030，
+其余 88 个文件是 UTF-8）。按单一编码读取会让中文注释整体乱码，
+而"注释必填"类规则会因此大面积误报。
+"""
+from __future__ import annotations
+
+BOM_UTF8 = b"\xef\xbb\xbf"
+CANDIDATES: tuple[str, ...] = ("utf-8", "gb18030")
+
+
+def sniff(raw: bytes) -> str:
+    """返回最可能的编码名。
+
+    顺序敏感：UTF-8 先试，失败再试 GB18030。
+    局限：字节序列恰好同时合法于两种编码时会误判（中文文本极少见）。
+    需要更高准确率时可换 charset-normalizer，本包默认零依赖。
+    """
+    if raw.startswith(BOM_UTF8):
+        return "utf-8-sig"
+    for enc in CANDIDATES:
+        try:
+            raw.decode(enc)
+            return enc
+        except UnicodeDecodeError:
+            continue
+    return "utf-8"  # 兜底，配合 errors="replace"
+
+
+def decode(raw: bytes, encoding: str | None = None) -> tuple[str, str]:
+    """返回 (文本, 实际使用的编码名)。"""
+    enc = encoding or sniff(raw)
+    return raw.decode(enc, errors="strict" if encoding else "replace"), enc

ddlkit/extract.py

@@ -0,0 +1,398 @@
+"""结构提取：token 流 -> 无损 UIM。
+
+不调用 sqlglot 的 parser，因此：
+* 不会挂死（达梦 ``NOT CLUSTER PRIMARY KEY`` 的无限回溯不存在）；
+* 不会改写（``TINYINT`` 不会被变成 ``SMALLINT``）。
+所有取值优先用 ``lex.raw()`` / ``lex.span()`` 从原串切片，保证原文可追溯。
+"""
+from __future__ import annotations
+
+import re
+
+from .keywords import (
+    CLAUSE_STARTERS,
+    COLUMN_ATTRS,
+    CONSTRAINT_STARTERS,
+    DECLARATION_MODIFIERS,
+    DEFAULT_TERMINATORS,
+    TYPE_TRAILERS,
+)
+from .lexer import Lexed, find_matching, split_top, unquote
+from .model import Column, Constraint, Index, SourceRef, Table
+from .source import Statement
+
+_STRING_LITERAL = re.compile(r"'((?:[^']|'')*)'")
+_QUOTE_CHARS = ('"', "`", "'", "[")
+
+_PARTITION_HEADS = frozenset({"PARTITION", "PARTITION BY", "SUBPARTITION"})
+
+
+def _string_literal(raw: str) -> str | None:
+    m = _STRING_LITERAL.search(raw)
+    return m.group(1).replace("''", "'") if m else None
+
+
+def _is_quoted(raw: str) -> bool:
+    return len(raw) >= 2 and raw[0] in _QUOTE_CHARS and raw[-1] in ('"', "`", "'", "]")
+
+
+# ------------------------------------------------------------------ 表名
+
+def _resolve_name(lex: Lexed, stop: int) -> tuple[list[tuple[str, str]], bool]:
+    """解析 ``CREATE ... TABLE`` 之后、表体 ``(`` 之前的名字。
+
+    返回 ([(原文, 去引号)], 是否含临时表修饰)。
+    """
+    idx = [i for i in range(stop) if lex.norm(i) not in DECLARATION_MODIFIERS]
+    segs: list[list[int]] = []
+    cur: list[int] = []
+    for i in idx:
+        if lex.text(i) == ".":
+            if cur:
+                segs.append(cur)
+            cur = []
+        else:
+            cur.append(i)
+    if cur:
+        segs.append(cur)
+
+    parts: list[tuple[str, str]] = []
+    for s in segs:
+        raw = lex.span(s[0], s[-1])
+        parts.append((raw, unquote(raw)))
+
+    # Hive beeline 里 ``ods.tbl`` 整体被一对反引号包住，需要再拆一层
+    if len(parts) == 1 and "." in parts[0][1]:
+        parts = [(x, x) for x in parts[0][1].split(".")]
+    return parts, False
+
+
+# ------------------------------------------------------------------ 列
+
+def _scan_default(lex: Lexed, k: int, end: int) -> tuple[str, int]:
+    """``DEFAULT`` 之后的值原文。
+
+    终止标记 **不含 NULL**：``DEFAULT NULL`` 必须保留成 ``"NULL"``。
+    括号内的逗号/关键字不终止（depth 感知）。
+    """
+    if k + 1 > end:
+        return "", k + 1
+    j = k + 1
+    depth = 0
+    while j <= end:
+        txt = lex.text(j)
+        up = lex.norm(j)
+        if txt == "(":
+            depth += 1
+        elif txt == ")":
+            depth -= 1
+        elif depth == 0 and j > k + 1 and up in DEFAULT_TERMINATORS:
+            break
+        j += 1
+    return lex.span(k + 1, j - 1).strip(), j
+
+
+def _parse_column(lex: Lexed, start: int, end: int) -> Column:
+    col = Column(line=lex.line_of(start), col=lex.col_of(start))
+
+    name_raw = lex.raw(start)
+    col.name_raw = name_raw
+    col.name = unquote(name_raw)
+    col.quoted = name_raw != col.name
+
+    i = start + 1
+    if i <= end:
+        type_start = i
+        i += 1
+        while i <= end:
+            if lex.text(i) == "(":
+                close = find_matching(lex.tokens, i, end)
+                if close < 0:
+                    break
+                i = close + 1
+                continue
+            if lex.norm(i) in TYPE_TRAILERS:
+                i += 1
+                continue
+            break
+        col.type_raw = lex.span(type_start, i - 1)
+        col.type_name = lex.norm(type_start)
+        args = re.search(r"\(.*\)", col.type_raw, re.S)
+        if args:
+            col.type_args_raw = args.group(0)
+
+    k = i
+    while k <= end:
+        up = lex.norm(k)
+        if up == "NOT" and k + 1 <= end and lex.norm(k + 1) == "NULL":
+            col.nullable = False
+            k += 2
+            continue
+        if up == "NULL":
+            col.nullable = True
+            k += 1
+            continue
+        if up == "DEFAULT":
+            col.default_raw, k = _scan_default(lex, k, end)
+            continue
+        if up == "COMMENT" and k + 1 <= end:
+            col.comment = unquote(lex.raw(k + 1))
+            k += 2
+            continue
+        if up in ("IDENTITY", "AUTO_INCREMENT", "AUTOINCREMENT"):
+            key = up.lower()
+            if k + 1 <= end and lex.text(k + 1) == "(":
+                close = find_matching(lex.tokens, k + 1, end)
+                if close > 0:
+                    col.extras[key] = lex.span(k + 1, close)
+                    k = close + 1
+                    continue
+            col.extras[key] = "true"
+            k += 1
+            continue
+        if up == "PRIMARY KEY" or up == "PRIMARY":
+            col.extras["primary_key"] = "true"
+            k += 1 if up == "PRIMARY KEY" else 2
+            continue
+        if up == "UNIQUE":
+            col.extras["unique"] = "true"
+            k += 1
+            continue
+        if up == "COLLATE" and k + 1 <= end:
+            col.extras["collate"] = lex.raw(k + 1)
+            k += 2
+            continue
+        if up in ("ENCODE", "CODEC") and k + 1 <= end:
+            col.extras[up.lower()] = lex.raw(k + 1)
+            k += 2
+            continue
+        if up == "ON":
+            col.extras["on"] = lex.span(k, end)
+            break
+        k += 1
+
+    # 派生：DEFAULT NULL 蕴含可空（仅在没有显式 NOT NULL 时）
+    if col.nullable is None and col.default_raw and col.default_raw.strip().upper() == "NULL":
+        col.nullable = True
+    return col
+
+
+# ------------------------------------------------------------------ 约束
+
+def _parse_constraint(lex: Lexed, start: int, end: int) -> Constraint:
+    c = Constraint(line=lex.line_of(start), raw=lex.span(start, end))
+    i = start
+
+    def eat_cluster(pos: int) -> int:
+        if pos <= end and lex.norm(pos) == "NOT" and pos + 1 <= end and lex.norm(pos + 1) == "CLUSTER":
+            c.clustered = False
+            return pos + 2
+        if pos <= end and lex.norm(pos) == "CLUSTER":
+            c.clustered = True
+            return pos + 1
+        return pos
+
+    i = eat_cluster(i)
+    if i <= end and lex.norm(i) == "CONSTRAINT":
+        if i + 1 <= end:
+            c.name_raw = lex.raw(i + 1)
+            c.name = unquote(c.name_raw)
+        i += 2
+        i = eat_cluster(i)
+
+    head = lex.norm(i) if i <= end else ""
+    if head in ("PRIMARY KEY", "PRIMARY"):
+        c.kind = "PRIMARY KEY"
+        i += 1 if head == "PRIMARY KEY" else 2
+    elif head in ("FOREIGN KEY", "FOREIGN"):
+        c.kind = "FOREIGN KEY"
+        i += 1 if head == "FOREIGN KEY" else 2
+    elif head == "UNIQUE":
+        c.kind = "UNIQUE"
+        i += 1
+        if i <= end and lex.norm(i) == "KEY":
+            i += 1
+    elif head == "CHECK":
+        c.kind = "CHECK"
+        i += 1
+    elif head in ("KEY", "INDEX"):
+        c.kind = "INDEX"
+        i += 1
+        if i <= end and lex.text(i) != "(":
+            c.name_raw = lex.raw(i)
+            c.name = unquote(c.name_raw)
+            i += 1
+    else:
+        c.kind = head or "UNKNOWN"
+        i += 1
+
+    while i <= end and lex.text(i) != "(":
+        if c.name is None:
+            c.name_raw = lex.raw(i)
+            c.name = unquote(c.name_raw)
+        i += 1
+
+    if i <= end and lex.text(i) == "(":
+        close = find_matching(lex.tokens, i, end)
+        if close > 0:
+            if c.kind == "CHECK":
+                c.extras["expression"] = lex.span(i, close)
+            else:
+                for s, t in split_top(lex.tokens, i + 1, close - 1):
+                    c.columns.append(unquote(lex.span(s, t)))
+            i = close + 1
+
+    if i <= end:
+        c.extras["tail"] = lex.span(i, end)
+    return c
+
+
+# ------------------------------------------------------------------ 尾部子句
+
+def _is_boundary(lex: Lexed, i: int, start: int, end: int) -> bool:
+    """判断下标 i 处的关键字是否真的开启一个新子句，而不是子句内部的值。
+
+    注意 ``)`` **不在**排除列表里：``ENGINE = MergeTree() ORDER BY (a, b)``
+    这种写法里，子句正是紧跟在右括号之后的。
+    """
+    if i == start:
+        return True
+    prev = lex.text(i - 1)
+    if prev in (",", "(", "=", "+", "-", "*", "/", "|"):
+        return False
+    # ``DEFAULT CHARSET = utf8mb4`` 应作为一个子句，不要把 DEFAULT 单独切出来
+    if lex.norm(i) == "DEFAULT" and i + 1 <= end and lex.norm(i + 1) in ("CHARSET", "CHARACTER SET"):
+        return False
+    return True
+
+
+def _scan_clauses(lex: Lexed, start: int, end: int) -> dict[str, str]:
+    """把表体括号之后的尾部子句原样归集到 extras。
+
+    只在 depth == 0 处判定边界——``STORAGE(ON "X", CLUSTERBTR)`` 里的
+    ``ON`` 在括号内，不会被误判成新子句。
+    """
+    extras: dict[str, str] = {}
+    cur_key: str | None = None
+    cur_start = 0
+    depth = 0
+    for i in range(start, end + 1):
+        txt = lex.text(i)
+        if txt == "(":
+            depth += 1
+            continue
+        if txt == ")":
+            depth -= 1
+            continue
+        if depth == 0:
+            up = lex.norm(i)
+            if up in CLAUSE_STARTERS and _is_boundary(lex, i, start, end):
+                if cur_key is not None:
+                    extras[cur_key] = lex.span(cur_start, i - 1).strip()
+                cur_key = up
+                cur_start = i
+    if cur_key is not None:
+        extras[cur_key] = lex.span(cur_start, end).strip()
+    return extras
+
+
+# ------------------------------------------------------------------ 入口
+
+def extract_table(lex: Lexed, dialect: str, src: SourceRef) -> Table | None:
+    paren = next((i for i in range(len(lex)) if lex.text(i) == "("), -1)
+    if paren <= 0:
+        return None
+    # CTAS（CREATE TABLE ... AS SELECT ...）没有列定义区
+    if any(lex.norm(i) == "AS" for i in range(paren)):
+        return None
+
+    tbl = Table(dialect=dialect, source=src)
+    tbl.temporary = any(
+        lex.norm(i) in ("GLOBAL", "LOCAL", "TEMPORARY", "TEMP", "VOLATILE", "TRANSIENT")
+        for i in range(min(paren, 8))
+    )
+
+    parts, _ = _resolve_name(lex, paren)
+    # parts 元素是 (原文, 去引号后) 二元组
+    if len(parts) >= 3:
+        _, tbl.catalog = parts[-3]
+        _, tbl.schema = parts[-2]
+        tbl.name_raw, tbl.name = parts[-1]
+    elif len(parts) == 2:
+        _, tbl.schema = parts[0]
+        tbl.name_raw, tbl.name = parts[1]
+    elif len(parts) == 1:
+        tbl.name_raw, tbl.name = parts[0]
+    tbl.quoted = _is_quoted(tbl.name_raw)
+
+    close = find_matching(lex.tokens, paren, len(lex) - 1)
+    if close < 0:
+        tbl.warnings.append("unclosed column list")
+        return tbl
+
+    for s, e in split_top(lex.tokens, paren + 1, close - 1):
+        head = lex.norm(s)
+        if head in _PARTITION_HEADS:
+            key = head if head in ("PARTITION BY",) else "PARTITION BY"
+            tbl.extras[key] = lex.span(s, e)
+        elif head in CONSTRAINT_STARTERS:
+            tbl.constraints.append(_parse_constraint(lex, s, e))
+        else:
+            tbl.columns.append(_parse_column(lex, s, e))
+
+    if close + 1 < len(lex):
+        tbl.extras.update(_scan_clauses(lex, close + 1, len(lex) - 1))
+
+    raw_comment = tbl.extras.get("COMMENT")
+    if raw_comment:
+        tbl.comment = _string_literal(raw_comment)
+    return tbl
+
+
+def extract_index(lex: Lexed, src: SourceRef) -> tuple[str, Index] | None:
+    """``CREATE [UNIQUE] INDEX name ON table (cols) [clauses]``"""
+    idx = Index(line=lex.line_of(0), raw=lex.span(0, len(lex) - 1))
+    i = 0
+    while i < len(lex) and lex.norm(i) in ("CREATE", "OR", "REPLACE", "GLOBAL", "LOCAL"):
+        i += 1
+    kinds: list[str] = []
+    while i < len(lex) and lex.norm(i) in (
+        "UNIQUE", "BITMAP", "CLUSTERED", "NONCLUSTERED", "FULLTEXT", "SPATIAL"
+    ):
+        kinds.append(lex.norm(i))
+        i += 1
+    if i >= len(lex) or lex.norm(i) not in ("INDEX", "KEY"):
+        return None
+    i += 1
+    idx.kind = (" ".join(kinds) + " INDEX").strip() if kinds else "INDEX"
+
+    if i < len(lex) and lex.text(i) != "ON":
+        idx.name_raw = lex.raw(i)
+        idx.name = unquote(idx.name_raw)
+        i += 1
+
+    target = ""
+    if i < len(lex) and lex.norm(i) == "ON":
+        i += 1
+        seg: list[int] = []
+        while i < len(lex) and lex.text(i) != "(":
+            seg.append(i)
+            i += 1
+        if seg:
+            target = unquote(lex.span(seg[0], seg[-1])).split(".")[-1].strip('"`')
+
+    if i < len(lex) and lex.text(i) == "(":
+        close = find_matching(lex.tokens, i, len(lex) - 1)
+        if close > 0:
+            for s, t in split_top(lex.tokens, i + 1, close - 1):
+                names = [lex.raw(x) for x in range(s, t + 1)
+                         if lex.norm(x) not in ("ASC", "DESC", "NULLS")]
+                if names:
+                    idx.columns.append(unquote(names[0]))
+            i = close + 1
+    if i < len(lex):
+        idx.extras["tail"] = lex.span(i, len(lex) - 1)
+    return target, idx
+
+
+__all__ = ["Statement", "extract_index", "extract_table"]