zsjz-ai/docs/db/export_pg_schema.py

# -*- coding: utf-8 -*-
"""导出 PostgreSQL 的建表结构（无数据）。

本机没有 pg_dump，因此用系统目录 + pg_get_*def 生成等价 DDL。
用法：python docs/db/export_pg_schema.py [输出文件]
"""
import datetime
import os
import sys

import psycopg2

DB = dict(host="192.168.0.106", port=5432, dbname="zsjz-ai",
          user="postgres", password="postgres")
DEFAULT_OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           "zsjz-ai-schema-20261028.sql")


def ident(name):
    return '"' + str(name).replace('"', '""') + '"'


def lit(value):
    return "'" + str(value).replace("'", "''") + "'"


def seq_opts(cur, seq_oid):
    cur.execute("""select seqtypid::regtype::text, seqstart, seqincrement, seqmin, seqmax,
                          seqcache, seqcycle
                   from pg_sequence where seqrelid=%s""", (seq_oid,))
    row = cur.fetchone()
    if not row:
        return ""
    typ, start, inc, mn, mx, cache, cycle = row
    parts = []
    if typ not in ("bigint", "integer", "smallint"):
        parts.append("AS " + typ)
    parts.append("START WITH %s" % start)
    parts.append("INCREMENT BY %s" % inc)
    if (inc > 0 and mn == 1) or (inc < 0 and mn == -(2 ** 63)):
        parts.append("NO MINVALUE")
    else:
        parts.append("MINVALUE %s" % mn)
    if (inc > 0 and mx == 2 ** 63 - 1) or (inc < 0 and mx == -1):
        parts.append("NO MAXVALUE")
    else:
        parts.append("MAXVALUE %s" % mx)
    parts.append("CACHE %s" % cache)
    if cycle:
        parts.append("CYCLE")
    return " ".join(parts)


def main():
    out_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_OUT
    conn = psycopg2.connect(**DB)
    cur = conn.cursor()

    def q(sql, args=None):
        cur.execute(sql, args)
        return cur.fetchall()

    lines = []
    lines.append("-- PostgreSQL database schema dump (structure only, no data)")
    lines.append("-- Database: " + DB["dbname"])
    lines.append("-- Server: " + q("select version()")[0][0])
    lines.append("-- Generated: " + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    lines.append("-- Note: generated from PostgreSQL catalogs because pg_dump is not installed on this host.")
    lines.append("")
    lines.append("SET client_encoding = 'UTF8';")
    lines.append("SET standard_conforming_strings = on;")
    lines.append("SET search_path = public, pg_catalog;")
    lines.append("")

    lines.append("-- Extensions")
    for ext, _ver in q("select extname, extversion from pg_extension "
                       "where extname <> 'plpgsql' order by 1"):
        lines.append("CREATE EXTENSION IF NOT EXISTS %s;" % ident(ext))
    lines.append("")

    # 序列：identity 序列由列内联创建，其余（serial）先创建再建表
    seqs = q("""select c.relname, c.oid from pg_class c
                join pg_namespace n on n.oid=c.relnamespace
                where n.nspname='public' and c.relkind='S' order by c.relname""")
    identity_seq_oids = {r[0] for r in q("""select d.objid from pg_depend d
        where d.classid='pg_class'::regclass and d.refclassid='pg_class'::regclass
          and d.deptype='i'""")}
    serial_seq_names = []
    if seqs:
        lines.append("-- Sequences (serial columns; identity sequences are created inline)")
        for name, oid in seqs:
            if oid in identity_seq_oids:
                continue
            serial_seq_names.append(name)
            lines.append("CREATE SEQUENCE public.%s" % ident(name))
            lines.append("    %s;" % seq_opts(cur, oid))
        lines.append("")

    tables = [r[0] for r in q("""select c.relname from pg_class c
        join pg_namespace n on n.oid=c.relnamespace
        where n.nspname='public' and c.relkind='r' order by c.relname""")]

    lines.append("-- Tables")
    for tname in tables:
        cols = q("""select a.attname, format_type(a.atttypid,a.atttypmod), a.attnotnull,
                           a.attidentity, a.attgenerated,
                           pg_get_expr(ad.adbin,ad.adrelid) as def,
                           a.attnum
                    from pg_attribute a
                    left join pg_attrdef ad on ad.adrelid=a.attrelid and ad.adnum=a.attnum
                    where a.attrelid=%s::regclass and a.attnum>0 and not a.attisdropped
                    order by a.attnum""", ("public." + ident(tname),))
        lines.append("CREATE TABLE public.%s (" % ident(tname))
        col_lines = []
        for attname, typ, notnull, identity, generated, default, attnum in cols:
            s = "    %s %s" % (ident(attname), typ)
            if generated == "s":
                s += " GENERATED ALWAYS AS (%s) STORED" % default
            elif identity in ("d", "a"):
                seq = q("""select s.relname, s.oid from pg_depend d
                           join pg_class s on s.oid=d.objid
                           where d.classid='pg_class'::regclass and d.refclassid='pg_class'::regclass
                             and d.refobjid=%s::regclass and d.refobjsubid=%s and d.deptype='i'""",
                        ("public." + ident(tname), attnum))
                if seq:
                    sname, soid = seq[0]
                    s += " GENERATED %s AS IDENTITY (SEQUENCE NAME public.%s %s)" % (
                        "ALWAYS" if identity == "a" else "BY DEFAULT",
                        ident(sname), seq_opts(cur, soid))
                else:
                    s += " GENERATED %s AS IDENTITY" % (
                        "ALWAYS" if identity == "a" else "BY DEFAULT")
            elif default is not None:
                s += " DEFAULT " + default
            if notnull:
                s += " NOT NULL"
            col_lines.append(s)
        lines.append(",\n".join(col_lines))
        lines.append(");")
        lines.append("")

    lines.append("-- Constraints")
    cons = q("""select c.relname, con.conname, con.contype,
                       pg_get_constraintdef(con.oid), con.convalidated
                from pg_constraint con join pg_class c on c.oid=con.conrelid
                join pg_namespace n on n.oid=c.relnamespace
                where n.nspname='public' and con.contype in ('p','u','f','c')
                order by c.relname, con.conname""")
    # FK 依赖被引用表的主键/唯一约束：先建 PK/UNIQUE，再建 FK/CHECK
    cons.sort(key=lambda x: (0 if x[2] in ("p", "u") else 1, x[0], x[1]))
    for tname, cname, ctype, cdef, valid in cons:
        s = "ALTER TABLE ONLY public.%s ADD CONSTRAINT %s %s" % (
            ident(tname), ident(cname), cdef)
        if not valid:
            s += " NOT VALID"
        lines.append(s + ";")
    lines.append("")

    lines.append("-- Indexes (excluding constraint-backed indexes)")
    for ddl, _t, _i in q("""select pg_get_indexdef(i.indexrelid), c.relname, ic.relname
                            from pg_index i join pg_class c on c.oid=i.indrelid
                            join pg_class ic on ic.oid=i.indexrelid
                            join pg_namespace n on n.oid=c.relnamespace
                            left join pg_constraint con on con.conindid=i.indexrelid
                            where n.nspname='public' and con.oid is null
                            order by c.relname, ic.relname"""):
        lines.append(ddl + ";")
    lines.append("")

    lines.append("-- Sequence ownership")
    for name in serial_seq_names:
        row = q("""select d.refobjid::regclass::text, a.attname from pg_depend d
                   join pg_class s on s.oid=d.objid
                   join pg_attribute a on a.attrelid=d.refobjid and a.attnum=d.refobjsubid
                   where d.classid='pg_class'::regclass and d.refclassid='pg_class'::regclass
                     and d.deptype='a' and s.relname=%s""", (name,))
        if row:
            tbl, col = row[0]
            lines.append("ALTER SEQUENCE public.%s OWNED BY %s.%s;" % (
                ident(name), tbl, ident(col)))
    lines.append("")

    lines.append("-- Comments")
    for tname in tables:
        row = q("select obj_description(%s::regclass,'pg_class')", ("public." + ident(tname),))
        if row and row[0][0]:
            lines.append("COMMENT ON TABLE public.%s IS %s;" % (ident(tname), lit(row[0][0])))
        for attname, comment in q("""select a.attname, col_description(a.attrelid,a.attnum)
                                     from pg_attribute a
                                     where a.attrelid=%s::regclass and a.attnum>0
                                       and not a.attisdropped
                                       and col_description(a.attrelid,a.attnum) is not null
                                     order by a.attnum""", ("public." + ident(tname),)):
            lines.append("COMMENT ON COLUMN public.%s.%s IS %s;" % (
                ident(tname), ident(attname), lit(comment)))
    lines.append("")

    with open(out_path, "w", encoding="utf-8", newline="\n") as f:
        f.write("\n".join(lines).rstrip() + "\n")
    conn.close()
    print("written %s (%d bytes, %d lines)" % (out_path, os.path.getsize(out_path), len(lines)))


if __name__ == "__main__":
    main()