iAOP/core/data-bus/postgres_schema.py

# -*- coding: utf-8 -*-
"""PostgreSQL 关系表 schema（模板 / 模型 / 用户 / 权限）—— 子任务 #30。

按模板自动生成隔离 schema（`tpl_{tpl}`），内含四张基础关系表：
- templates   模板注册表（行业模板配置资产）；
- models      模型注册表（模型版本 / 算法 / 超参包）；
- users       用户表（角色 + 细粒度权限 JSONB）；
- permissions 角色-资源-动作权限矩阵。

另生成角色授权语句（GRANT），按模板配置声明 databus_ro / databus_rw 等角色。
不依赖 psycopg2：仅产出 DDL 文本，供模板配置台预览 / 运维执行。
"""
from __future__ import annotations

from typing import List, Optional

from .templating import TemplateNaming

# 默认基础表清单（对应子任务 #30「模板/模型/用户/权限」）
DEFAULT_TABLES: List[str] = ["templates", "models", "users", "permissions"]

# 表 → 列定义（可被模板 YAML 的 schema_tables 覆盖为子集/扩展）
TABLE_COLUMNS: dict = {
    "templates": [
        "id BIGSERIAL PRIMARY KEY",
        "template VARCHAR(64) NOT NULL UNIQUE",
        "version VARCHAR(32) NOT NULL DEFAULT '0.1.0'",
        "config JSONB NOT NULL DEFAULT '{{}}'::jsonb",  # format() 转义 {}
        "status VARCHAR(16) NOT NULL DEFAULT 'draft'",
        "created_at TIMESTAMPTZ NOT NULL DEFAULT now()",
    ],
    "models": [
        "id BIGSERIAL PRIMARY KEY",
        "model_id VARCHAR(128) NOT NULL UNIQUE",
        "template VARCHAR(64) REFERENCES {schema}.templates(template) ON DELETE CASCADE",
        "algorithm VARCHAR(32) NOT NULL DEFAULT 'xgboost'",
        "hyperparams JSONB NOT NULL DEFAULT '{{}}'::jsonb",  # format() 转义 {}
        "version VARCHAR(32) NOT NULL DEFAULT '0.1.0'",
        "status VARCHAR(16) NOT NULL DEFAULT 'staging'",
        "created_at TIMESTAMPTZ NOT NULL DEFAULT now()",
    ],
    "users": [
        "id BIGSERIAL PRIMARY KEY",
        "username VARCHAR(64) NOT NULL UNIQUE",
        "role VARCHAR(32) NOT NULL DEFAULT 'viewer'",
        "permissions JSONB NOT NULL DEFAULT '[]'::jsonb",
        "created_at TIMESTAMPTZ NOT NULL DEFAULT now()",
    ],
    "permissions": [
        "id BIGSERIAL PRIMARY KEY",
        "role VARCHAR(32) NOT NULL",
        "resource VARCHAR(64) NOT NULL",
        "action VARCHAR(16) NOT NULL",
        "UNIQUE (role, resource, action)",
    ],
}


def generate_schema_ddl(
    naming: TemplateNaming,
    tables: Optional[List[str]] = None,
    table_columns: Optional[dict] = None,
) -> str:
    """生成 schema 与基础关系表的完整 DDL（幂等 IF NOT EXISTS）。

    Args:
        naming: 模板命名器（决定 schema 名）。
        tables: 表清单（默认 DEFAULT_TABLES）。
        table_columns: 表→列定义覆盖（默认 TABLE_COLUMNS）。

    Returns:
        多语句 DDL 文本（含 CREATE SCHEMA / CREATE TABLE）。
    """
    names = tables or DEFAULT_TABLES
    cols = table_columns or TABLE_COLUMNS
    schema = naming.pg_schema()

    lines = [f"CREATE SCHEMA IF NOT EXISTS {schema};", ""]
    for table in names:
        columns = cols.get(table)
        if not columns:
            continue
        lines.append(f"CREATE TABLE IF NOT EXISTS {schema}.{table} (")
        lines.append("  " + ",\n  ".join(c.format(schema=schema) for c in columns))
        lines.append(");")
        lines.append("")
    return "\n".join(lines).rstrip() + "\n"


def generate_grant_ddl(
    naming: TemplateNaming,
    roles: Optional[List[str]] = None,
) -> List[str]:
    """按模板声明的关系角色生成授权语句。

    角色约定：`{role}_ro` 只读、`{role}_rw` 读写（可扩展）。
    Args:
        naming: 模板命名器。
        roles: 角色清单（如 ["databus"] → databus_ro / databus_rw）。
    Returns:
        GRANT 语句列表。
    """
    schema = naming.pg_schema()
    grants: List[str] = []
    for role in (roles or ["databus"]):
        for suffix, privileges in (("ro", "SELECT"), ("rw", "SELECT, INSERT, UPDATE, DELETE")):
            role_name = f"{role}_{suffix}"
            grants.append(f"GRANT USAGE ON SCHEMA {schema} TO {role_name};")
            grants.append(
                f"GRANT {privileges} ON ALL TABLES IN SCHEMA {schema} TO {role_name};"
            )
            if suffix == "rw":
                grants.append(
                    f"GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA {schema} TO {role_name};"
                )
    return grants