diff --git a/backend/app/api/snapshots.py b/backend/app/api/snapshots.py index 69cf84389..f221aaaaa 100644 --- a/backend/app/api/snapshots.py +++ b/backend/app/api/snapshots.py @@ -44,6 +44,11 @@ from app.spec.index_redundancy import detect_index_redundancy from app.spec.data_dictionary import snapshot_to_data_dictionary_md from app.spec.naming_lint import lint_naming +from app.spec.orm_codegen import ( + generate_prisma_schema, + generate_sqlalchemy_models, + generate_typeorm_entities, +) from app.spec.relationship_inference import infer_relationships from app.spec.schema_stats import compute_schema_stats from app.spec.sensitive_columns import detect_sensitive_columns @@ -372,6 +377,31 @@ async def wide_tables( ) +@router.get("/{schema_snapshot_uuid}/orm-models", response_class=PlainTextResponse) +async def export_orm_models( + schema_snapshot_uuid: uuid.UUID, + flavor: str = Query("sqlalchemy", pattern="^(sqlalchemy|prisma|typeorm)$"), + user: CurrentUser = Depends(get_current_user), + session: AsyncSession = Depends(get_read_session), +) -> str: + """Generate ORM model code from a snapshot (SQLAlchemy 2.0 or Prisma). + + Forward engineering to application code: hand developers ready-to-edit + models instead of retyped tables. IDOR-safe (uniform not-found marker). + """ + snap = await _get_authorized_snapshot(session, schema_snapshot_uuid, user) + if snap is None: + return "-- snapshot not found\n" + data = await session.get(SchemaSnapshotData, schema_snapshot_uuid) + if data is None: + return "-- snapshot data not found\n" + if flavor == "prisma": + return generate_prisma_schema(data.snapshot_json) + if flavor == "typeorm": + return generate_typeorm_entities(data.snapshot_json) + return generate_sqlalchemy_models(data.snapshot_json) + + @router.get("/{schema_snapshot_uuid}/stats", response_model=SchemaStatsOut) async def schema_stats( schema_snapshot_uuid: uuid.UUID, diff --git a/backend/app/spec/orm_codegen.py b/backend/app/spec/orm_codegen.py new file mode 100644 index 000000000..a6ddc9157 --- /dev/null +++ b/backend/app/spec/orm_codegen.py @@ -0,0 +1,307 @@ +"""Generate ORM model code (SQLAlchemy 2.0 / Prisma) from a snapshot. + +Forward engineering all the way to application code: reverse-engineer a legacy +database, then hand developers ready-to-edit ORM models instead of making them +retype every table. Most diagram tools stop at DDL — this is a differentiator +(see docs/erd-tool-feature-research.md). + +Pure and dialect-agnostic. Type mapping is best-effort: unmapped SQL types fall +back to ``str``/``String`` with the original type kept in a comment, so the +output always imports/parses. +""" + +from __future__ import annotations + +import json +import keyword +import re +from typing import Any + +_PY_TYPES: list[tuple[re.Pattern[str], str]] = [ + (re.compile(r"^(big)?serial|^(small|big)?int|^integer"), "int"), + (re.compile(r"^bool"), "bool"), + (re.compile(r"^(timestamp|datetime)"), "dt.datetime"), + (re.compile(r"^date$"), "dt.date"), + (re.compile(r"^time($|\()"), "dt.time"), + (re.compile(r"^(numeric|decimal)"), "Decimal"), + (re.compile(r"^(real|float|double)"), "float"), + (re.compile(r"^uuid"), "uuid.UUID"), + (re.compile(r"^(json|jsonb)"), "dict"), + (re.compile(r"^bytea|^(var)?binary|^blob"), "bytes"), + (re.compile(r"^(text|(var)?char|character|enum)"), "str"), +] + +_PRISMA_TYPES: list[tuple[re.Pattern[str], str]] = [ + (re.compile(r"^bigserial|^bigint"), "BigInt"), + (re.compile(r"^serial|^(small)?int|^integer"), "Int"), + (re.compile(r"^bool"), "Boolean"), + (re.compile(r"^(timestamp|datetime|date$|time($|\())"), "DateTime"), + (re.compile(r"^(numeric|decimal)"), "Decimal"), + (re.compile(r"^(real|float|double)"), "Float"), + (re.compile(r"^(json|jsonb)"), "Json"), + (re.compile(r"^bytea|^(var)?binary|^blob"), "Bytes"), +] + + +def _map_type(data_type: str, table: list[tuple[re.Pattern[str], str]], default: str) -> str: + lowered = data_type.strip().lower() + for pattern, mapped in table: + if pattern.search(lowered): + return mapped + return default + + +def _snake_name(value: object, fallback: str) -> str: + name = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", str(value or "")) + name = re.sub(r"[^0-9a-zA-Z]+", "_", name).strip("_").lower() + name = name or fallback + if name[0].isdigit(): + name = f"{fallback}_{name}" + if keyword.iskeyword(name): + name = f"{name}_" + return name + + +def _class_name(table_name: str) -> str: + parts = re.split(r"[^0-9a-zA-Z]+", table_name) + name = "".join(p[:1].upper() + p[1:] for p in parts if p) + if not name or name[0].isdigit(): + name = f"Table{name}" + return f"{name}Model" if keyword.iskeyword(name) else name + + +def _column_names(cols: list[dict[str, Any]]) -> dict[str, str]: + return {str(col.get("column_name")): _snake_name(col.get("column_name"), "column") for col in cols} + + +def _index(snapshot: dict[str, Any] | None) -> dict[str, Any]: + snapshot = snapshot or {} + relations = [ + r for r in snapshot.get("relations") or [] + if (r.get("relation_kind") or "r") in ("r", "p") + ] + cols_by_oid: dict[Any, list[dict[str, Any]]] = {} + for c in snapshot.get("columns") or []: + cols_by_oid.setdefault(c.get("relation_oid"), []).append(c) + for cols in cols_by_oid.values(): + cols.sort(key=lambda c: c.get("column_position") or 0) + pk_by_oid: dict[Any, set[str]] = {} + for pk in snapshot.get("pk_columns") or []: + pk_by_oid.setdefault(pk.get("relation_oid"), set()).add(str(pk.get("column_name"))) + fk_by_child: dict[tuple[Any, str], dict[str, Any]] = {} + for e in snapshot.get("fk_edges") or []: + fk_by_child[(e.get("child_relation_oid"), str(e.get("child_column_name")))] = e + rel_by_oid = {r.get("relation_oid"): r for r in relations} + return { + "relations": relations, + "cols_by_oid": cols_by_oid, + "pk_by_oid": pk_by_oid, + "fk_by_child": fk_by_child, + "rel_by_oid": rel_by_oid, + } + + +def generate_sqlalchemy_models(snapshot: dict[str, Any] | None) -> str: + """Render SQLAlchemy 2.0 (Mapped/mapped_column) models.""" + ix = _index(snapshot) + lines = [ + "# Generated by pg-erd-cloud — SQLAlchemy 2.0 models (edit freely)", + "from __future__ import annotations", + "", + "import datetime as dt", + "import uuid", + "from decimal import Decimal", + "", + "from sqlalchemy import ForeignKey", + "from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column", + "", + "", + "class Base(DeclarativeBase):", + " pass", + ] + for rel in ix["relations"]: + oid = rel.get("relation_oid") + table = str(rel.get("relation_name")) + schema = str(rel.get("schema_name")) + lines += ["", "", f"class {_class_name(table)}(Base):"] + if rel.get("relation_comment"): + lines.append(f" {str(rel['relation_comment'])!r}") + lines.append("") + lines.append(f" __tablename__ = {table!r}") + if schema and schema != "public": + lines.append(f" __table_args__ = {{'schema': {schema!r}}}") + cols = ix["cols_by_oid"].get(oid, []) + col_names = _column_names(cols) + if not cols: + lines.append(" pass") + continue + for col in cols: + name = str(col.get("column_name")) + attr = col_names[name] + raw_type = str(col.get("data_type") or "") + py_type = _map_type(raw_type, _PY_TYPES, "str") + not_null = bool(col.get("is_not_null")) + annotated = py_type if not_null else f"{py_type} | None" + args = [repr(name)] if attr != name else [] + fk = ix["fk_by_child"].get((oid, name)) + if fk is not None: + parent = ix["rel_by_oid"].get(fk.get("parent_relation_oid")) + if parent is not None: + target = f'{parent.get("schema_name")}.{parent.get("relation_name")}.{fk.get("parent_column_name")}' + if parent.get("schema_name") == "public": + target = f'{parent.get("relation_name")}.{fk.get("parent_column_name")}' + args.append(f"ForeignKey({target!r})") + if name in ix["pk_by_oid"].get(oid, set()): + args.append("primary_key=True") + comment = "" if _map_type(raw_type, _PY_TYPES, "?") != "?" else f" # type: {raw_type}" + call = f"mapped_column({', '.join(args)})" if args else "mapped_column()" + lines.append(f" {attr}: Mapped[{annotated}] = {call}{comment}") + return "\n".join(lines) + "\n" + + +def generate_prisma_schema(snapshot: dict[str, Any] | None) -> str: + """Render a Prisma schema (models with @id/@map; relations as field pairs).""" + ix = _index(snapshot) + lines = [ + "// Generated by pg-erd-cloud — Prisma schema (edit freely)", + 'datasource db {', + ' provider = "postgresql"', + ' url = env("DATABASE_URL")', + '}', + ] + # child fk list per parent for reverse relation fields + children_of: dict[Any, list[dict[str, Any]]] = {} + for (child_oid, _), e in ix["fk_by_child"].items(): + children_of.setdefault(e.get("parent_relation_oid"), []).append(e) + + for rel in ix["relations"]: + oid = rel.get("relation_oid") + table = str(rel.get("relation_name")) + model = _class_name(table) + lines += ["", f"model {model} {{"] + cols = ix["cols_by_oid"].get(oid, []) + col_names = _column_names(cols) + pk_cols = ix["pk_by_oid"].get(oid, set()) + for col in cols: + name = str(col.get("column_name")) + field_name = col_names[name] + raw_type = str(col.get("data_type") or "") + p_type = _map_type(raw_type, _PRISMA_TYPES, "String") + optional = "" if col.get("is_not_null") else "?" + attrs = [] + if name in pk_cols and len(pk_cols) == 1: + attrs.append("@id") + if field_name != name: + attrs.append(f"@map({json.dumps(name)})") + comment = "" if _map_type(raw_type, _PRISMA_TYPES, "?") != "?" and p_type != "String" or raw_type.lower().startswith(("text", "varchar", "char")) else f" // {raw_type}" + lines.append(f" {field_name} {p_type}{optional}{' ' + ' '.join(attrs) if attrs else ''}{comment}") + # relation fields + for (child_oid, child_col), e in ix["fk_by_child"].items(): + if child_oid != oid: + continue + parent = ix["rel_by_oid"].get(e.get("parent_relation_oid")) + if parent is None: + continue + pmodel = _class_name(str(parent.get("relation_name"))) + child_field = col_names.get(child_col, _snake_name(child_col, "column")) + parent_cols = ix["cols_by_oid"].get(e.get("parent_relation_oid"), []) + parent_field = _column_names(parent_cols).get(str(e.get("parent_column_name")), _snake_name(e.get("parent_column_name"), "column")) + field = re.sub(r"_id$", "", child_field) or _snake_name(pmodel, "relation") + lines.append( + f" {field} {pmodel} @relation(fields: [{child_field}], references: [{parent_field}])" + ) + for e in children_of.get(oid, []): + child = ix["rel_by_oid"].get(e.get("child_relation_oid")) + if child is None: + continue + cmodel = _class_name(str(child.get("relation_name"))) + lines.append(f" {cmodel[0].lower() + cmodel[1:]}s {cmodel}[]") + if len(pk_cols) > 1: + ordered = [col_names[str(c.get("column_name"))] for c in cols if str(c.get("column_name")) in pk_cols] + lines.append(f" @@id([{', '.join(ordered)}])") + lines.append(f" @@map({json.dumps(table)})") + lines.append("}") + return "\n".join(lines) + "\n" + + +_TS_TYPES: list[tuple[re.Pattern[str], str]] = [ + (re.compile(r"^(big)?serial|^(small|big)?int|^integer|^(real|float|double|numeric|decimal)"), "number"), + (re.compile(r"^bool"), "boolean"), + (re.compile(r"^(timestamp|datetime|date$|time($|\())"), "Date"), + (re.compile(r"^(json|jsonb)"), "object"), + (re.compile(r"^bytea|^(var)?binary|^blob"), "Buffer"), +] + + +def generate_typeorm_entities(snapshot: dict[str, Any] | None) -> str: + """Render TypeORM entity classes (TypeScript decorators).""" + ix = _index(snapshot) + lines = [ + "// Generated by pg-erd-cloud — TypeORM entities (edit freely)", + "import {", + " Column,", + " Entity,", + " JoinColumn,", + " ManyToOne,", + " OneToMany,", + " PrimaryColumn,", + "} from 'typeorm';", + ] + # reverse-relation index: parent oid -> child fk edges + children_of: dict[Any, list[dict[str, Any]]] = {} + for (child_oid, _), e in ix["fk_by_child"].items(): + children_of.setdefault(e.get("parent_relation_oid"), []).append(e) + + for rel in ix["relations"]: + oid = rel.get("relation_oid") + table = str(rel.get("relation_name")) + schema = str(rel.get("schema_name")) + cls = _class_name(table) + entity_args = json.dumps(table) if schema == "public" else f"{{ name: {json.dumps(table)}, schema: {json.dumps(schema)} }}" + lines += ["", f"@Entity({entity_args})", f"export class {cls} {{"] + cols = ix["cols_by_oid"].get(oid, []) + col_names = _column_names(cols) + for col in cols: + name = str(col.get("column_name")) + field_name = col_names[name] + raw_type = str(col.get("data_type") or "") + ts_type = _map_type(raw_type, _TS_TYPES, "string") + nullable = not col.get("is_not_null") + is_pk = name in ix["pk_by_oid"].get(oid, set()) + options = [] + if field_name != name: + options.append(f"name: {json.dumps(name)}") + if nullable and not is_pk: + options.append("nullable: true") + decorator = "@PrimaryColumn" if is_pk else "@Column" + decorator = f"{decorator}({{{', '.join(options)}}})" if options else f"{decorator}()" + comment = "" if _map_type(raw_type, _TS_TYPES, "?") != "?" and ts_type != "string" or raw_type.lower().startswith(("text", "varchar", "char")) else f" // {raw_type}" + lines.append(f" {decorator}") + lines.append(f" {field_name}{'?' if nullable else '!'}: {ts_type}{' | null' if nullable else ''};{comment}") + lines.append("") + for (child_oid, child_col), e in ix["fk_by_child"].items(): + if child_oid != oid: + continue + parent = ix["rel_by_oid"].get(e.get("parent_relation_oid")) + if parent is None: + continue + pcls = _class_name(str(parent.get("relation_name"))) + child_field = col_names.get(child_col, _snake_name(child_col, "column")) + field = re.sub(r"_id$", "", child_field) or _snake_name(pcls, "relation") + lines.append(f" @ManyToOne(() => {pcls})") + lines.append(f" @JoinColumn({{ name: {json.dumps(child_col)} }})") + lines.append(f" {field}?: {pcls};") + lines.append("") + for e in children_of.get(oid, []): + child = ix["rel_by_oid"].get(e.get("child_relation_oid")) + if child is None: + continue + ccls = _class_name(str(child.get("relation_name"))) + field = ccls[0].lower() + ccls[1:] + "s" + lines.append(f" @OneToMany(() => {ccls}, (child) => child)") + lines.append(f" {field}?: {ccls}[];") + lines.append("") + while lines and lines[-1] == "": + lines.pop() + lines.append("}") + return "\n".join(lines) + "\n" diff --git a/backend/tests/test_orm_codegen.py b/backend/tests/test_orm_codegen.py new file mode 100644 index 000000000..51a50fadc --- /dev/null +++ b/backend/tests/test_orm_codegen.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import ast + +from app.spec.orm_codegen import generate_prisma_schema, generate_sqlalchemy_models + +SNAP = { + "relations": [ + {"relation_oid": 1, "relation_kind": "r", "schema_name": "public", "relation_name": "member", "relation_comment": "회원"}, + {"relation_oid": 2, "relation_kind": "r", "schema_name": "public", "relation_name": "orders", "relation_comment": None}, + {"relation_oid": 3, "relation_kind": "v", "schema_name": "public", "relation_name": "v_report"}, + ], + "columns": [ + {"relation_oid": 1, "column_name": "member_id", "column_position": 1, "data_type": "bigint", "is_not_null": True}, + {"relation_oid": 1, "column_name": "email", "column_position": 2, "data_type": "varchar(255)", "is_not_null": True}, + {"relation_oid": 1, "column_name": "joined_at", "column_position": 3, "data_type": "timestamp with time zone", "is_not_null": False}, + {"relation_oid": 2, "column_name": "order_id", "column_position": 1, "data_type": "bigint", "is_not_null": True}, + {"relation_oid": 2, "column_name": "member_id", "column_position": 2, "data_type": "bigint", "is_not_null": True}, + {"relation_oid": 2, "column_name": "total", "column_position": 3, "data_type": "numeric(10,2)", "is_not_null": False}, + ], + "pk_columns": [ + {"relation_oid": 1, "column_name": "member_id"}, + {"relation_oid": 2, "column_name": "order_id"}, + ], + "fk_edges": [ + {"fk_constraint_oid": 10, "fk_constraint_name": "fk_orders_member", + "child_relation_oid": 2, "parent_relation_oid": 1, + "child_column_name": "member_id", "parent_column_name": "member_id", "column_ordinal": 1}, + ], +} + + +def test_sqlalchemy_output_is_valid_python_with_expected_shapes(): + code = generate_sqlalchemy_models(SNAP) + ast.parse(code) # must always be syntactically valid Python + assert "class Member(Base):" in code + assert "class Orders(Base):" in code + assert "class VReport" not in code # views excluded + assert "member_id: Mapped[int] = mapped_column(primary_key=True)" in code + assert "joined_at: Mapped[dt.datetime | None] = mapped_column()" in code + assert "ForeignKey('member.member_id')" in code + assert "total: Mapped[Decimal | None]" in code + + +def test_sqlalchemy_unknown_type_falls_back_with_comment(): + snap = { + "relations": [{"relation_oid": 1, "relation_kind": "r", "schema_name": "public", "relation_name": "t"}], + "columns": [{"relation_oid": 1, "column_name": "shape", "column_position": 1, "data_type": "polygon", "is_not_null": False}], + "pk_columns": [], "fk_edges": [], + } + code = generate_sqlalchemy_models(snap) + ast.parse(code) + assert "shape: Mapped[str | None] = mapped_column() # type: polygon" in code + + +def test_prisma_models_relations_and_map(): + schema = generate_prisma_schema(SNAP) + assert "model Member {" in schema and "model Orders {" in schema + assert "member_id BigInt @id" in schema + assert "member Member @relation(fields: [member_id], references: [member_id])" in schema + assert "orderss Orders[]" in schema or "orders Orders[]" in schema # reverse side exists + assert '@@map("orders")' in schema + assert "total Decimal?" in schema + + +def test_composite_pk_and_empty_snapshot(): + snap = { + "relations": [{"relation_oid": 1, "relation_kind": "r", "schema_name": "public", "relation_name": "m2m"}], + "columns": [ + {"relation_oid": 1, "column_name": "a_id", "column_position": 1, "data_type": "bigint", "is_not_null": True}, + {"relation_oid": 1, "column_name": "b_id", "column_position": 2, "data_type": "bigint", "is_not_null": True}, + ], + "pk_columns": [ + {"relation_oid": 1, "column_name": "a_id"}, + {"relation_oid": 1, "column_name": "b_id"}, + ], + "fk_edges": [], + } + schema = generate_prisma_schema(snap) + assert "@@id([a_id, b_id])" in schema + assert "@id\n" not in schema.replace("@@id", "") # no single-column @id emitted + ast.parse(generate_sqlalchemy_models({})) # empty snapshot still valid + + +def test_typeorm_entities_decorators_and_relations(): + from app.spec.orm_codegen import generate_typeorm_entities + + code = generate_typeorm_entities(SNAP) + assert '@Entity("member")' in code + assert "export class Member {" in code + assert "@PrimaryColumn()" in code + assert "member_id!: number;" in code + assert "joined_at?: Date | null;" in code + assert "@ManyToOne(() => Member)" in code + assert '@JoinColumn({ name: "member_id" })' in code + assert "@OneToMany(() => Orders" in code + # balanced braces => structurally sound TS + assert code.count("{") == code.count("}") + + +def test_non_snake_db_names_are_mapped_to_safe_identifiers(): + from app.spec.orm_codegen import generate_typeorm_entities + + snap = { + "relations": [ + { + "relation_oid": 1, + "relation_kind": "r", + "schema_name": "public", + "relation_name": "Order Items", + "relation_comment": 'quote " is fine', + } + ], + "columns": [ + {"relation_oid": 1, "column_name": "Order ID", "column_position": 1, "data_type": "bigint", "is_not_null": True}, + {"relation_oid": 1, "column_name": "class", "column_position": 2, "data_type": "text", "is_not_null": False}, + {"relation_oid": 1, "column_name": "2FA Enabled", "column_position": 3, "data_type": "boolean", "is_not_null": True}, + ], + "pk_columns": [{"relation_oid": 1, "column_name": "Order ID"}], + "fk_edges": [], + } + + py_code = generate_sqlalchemy_models(snap) + ast.parse(py_code) + assert "class OrderItems(Base):" in py_code + assert "order_id: Mapped[int] = mapped_column('Order ID', primary_key=True)" in py_code + assert "class_: Mapped[str | None] = mapped_column('class')" in py_code + assert "column_2_fa_enabled: Mapped[bool] = mapped_column('2FA Enabled')" in py_code + + prisma = generate_prisma_schema(snap) + assert "model OrderItems {" in prisma + assert 'order_id BigInt @id @map("Order ID")' in prisma + assert 'class_ String? @map("class")' in prisma + assert 'column_2_fa_enabled Boolean @map("2FA Enabled")' in prisma + + ts_code = generate_typeorm_entities(snap) + assert '@Entity("Order Items")' in ts_code + assert '@PrimaryColumn({name: "Order ID"})' in ts_code + assert 'order_id!: number;' in ts_code + assert '@Column({name: "class", nullable: true})' in ts_code + assert 'class_?: string | null;' in ts_code