Build an AI Database Schema Designer with Python

pip install anthropic pydantic

export ANTHROPIC_API_KEY="sk-your-ezai-key"

import json
import anthropic
from pydantic import BaseModel

class SchemaDesigner:
    def __init__(self):
        self.client = anthropic.Anthropic(
            base_url="https://ezaiapi.com"
        )
        self.system_prompt = """You are a senior database architect.
Given app requirements in plain English, output a JSON object with:
- "tables": array of table definitions
- "indexes": array of CREATE INDEX statements
- "notes": array of design decision explanations

Each table has: "name", "columns" (array of {name, type, constraints}),
"primary_key", "foreign_keys" (array of {column, references_table,
references_column, on_delete}).

Rules:
- Use snake_case for all names
- Always include id (SERIAL/BIGSERIAL), created_at, updated_at
- Add soft delete (deleted_at) for user-facing entities
- Use TIMESTAMPTZ, never TIMESTAMP
- Add indexes on foreign keys and commonly queried columns
- Use DECIMAL for money, never FLOAT
- Add CHECK constraints where appropriate
- Include junction tables for many-to-many relationships
Output ONLY valid JSON. No markdown, no explanation outside the JSON."""

    def design(self, requirements: str, dialect: str = "postgresql") -> dict:
        msg = self.client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=4096,
            system=self.system_prompt,
            messages=[{
                "role": "user",
                "content": f"Design a {dialect} schema for: {requirements}"
            }]
        )
        return json.loads(msg.content[0].text)

def to_ddl(schema: dict) -> str:
    """Convert the AI-generated schema JSON to executable DDL."""
    statements = []

    for table in schema["tables"]:
        cols = []
        for col in table["columns"]:
            line = f"    {col['name']} {col['type']}"
            if col.get("constraints"):
                line += f" {col['constraints']}"
            cols.append(line)

        # Add foreign key constraints
        for fk in table.get("foreign_keys", []):
            ref = f"    FOREIGN KEY ({fk['column']}) "
            ref += f"REFERENCES {fk['references_table']}({fk['references_column']})"
            if fk.get("on_delete"):
                ref += f" ON DELETE {fk['on_delete']}"
            cols.append(ref)

        sql = f"CREATE TABLE {table['name']} (\n"
        sql += ",\n".join(cols)
        sql += "\n);"
        statements.append(sql)

    # Add indexes
    for idx in schema.get("indexes", []):
        statements.append(idx)

    return "\n\n".join(statements)

designer = SchemaDesigner()

schema = designer.design("""
A project management app where:
- Users belong to organizations
- Organizations have multiple projects
- Projects contain tasks with assignees and due dates
- Tasks can have comments and file attachments
- Users can watch tasks for notifications
""")

ddl = to_ddl(schema)
print(ddl)

# Save to file
with open("schema.sql", "w") as f:
    f.write(ddl)

# Print design notes
for note in schema.get("notes", []):
    print(f"💡 {note}")

def validate_schema(schema: dict) -> list[str]:
    """Return a list of warnings for the generated schema."""
    warnings = []
    table_names = {t["name"] for t in schema["tables"]}

    for table in schema["tables"]:
        col_names = [c["name"] for c in table["columns"]]

        # Check for timestamps
        if "created_at" not in col_names:
            warnings.append(
                f"⚠️  {table['name']}: missing created_at column"
            )

        # Check foreign key references exist
        for fk in table.get("foreign_keys", []):
            if fk["references_table"] not in table_names:
                warnings.append(
                    f"❌ {table['name']}.{fk['column']} references "
                    f"non-existent table {fk['references_table']}"
                )

        # Check for FLOAT on money columns
        for col in table["columns"]:
            if "FLOAT" in col["type"].upper() and \
               any(w in col["name"] for w in ["price", "cost", "amount", "total"]):
                warnings.append(
                    f"⚠️  {table['name']}.{col['name']}: use DECIMAL for money, not FLOAT"
                )

    return warnings

def to_sqlalchemy(schema: dict) -> str:
    """Generate SQLAlchemy 2.0 model code from schema JSON."""
    designer = SchemaDesigner()

    msg = designer.client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=4096,
        messages=[{
            "role": "user",
            "content": f"""Convert this database schema JSON to SQLAlchemy 2.0 models
using mapped_column and DeclarativeBase. Include relationships,
type annotations, and __repr__ methods.

Schema: {json.dumps(schema, indent=2)}

Output ONLY the Python code, no markdown."""
        }]
    )
    return msg.content[0].text

# Usage
models_code = to_sqlalchemy(schema)
with open("models.py", "w") as f:
    f.write(models_code)

def refine(self, current_schema: dict, changes: str) -> dict:
    """Modify an existing schema based on new requirements."""
    msg = self.client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=4096,
        system=self.system_prompt,
        messages=[
            {"role": "user", "content": "Here is my current schema:"},
            {"role": "assistant", "content": json.dumps(current_schema)},
            {"role": "user", "content": f"Modify the schema: {changes}. Return the complete updated schema JSON."}
        ]
    )
    return json.loads(msg.content[0].text)

# Example: add multi-tenancy
updated = designer.refine(schema, """
Add row-level multi-tenancy:
- Every table gets a tenant_id column
- Add a tenants table with name, slug, plan
- All queries should scope by tenant_id
""")