seedloom 0.1.0__py3-none-any.whl

seedloom/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

seedloom/cli.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+import click
+import psycopg2
+from rich.console import Console
+from rich.table import Table as RichTable
+
+from .config import Config
+from .generator import generate_rows
+from .graph import CyclicDependencyError, resolve_seed_order
+from .inserter import existing_column_values, insert_rows, table_row_count
+from .introspect import introspect
+from .models import Schema
+from .providers import ProviderError, SUPPORTED_PROVIDERS, get_provider
+
+console = Console()
+SCHEMA_CACHE = Path(".seedloom_schema.json")
+
+
+@click.group()
+def main() -> None:
+    """seedloom — AI-powered database seeding.
+
+    Introspects your Postgres schema and uses Claude to generate realistic,
+    referentially-valid seed data.
+    """
+
+
+@main.command()
+def init() -> None:
+    """Connect to the database, introspect the schema, and cache it locally."""
+    try:
+        config = Config.load(require_provider=False)
+    except EnvironmentError as e:
+        console.print(f"[red]{e}[/red]")
+        sys.exit(1)
+
+    console.print("[cyan]Connecting and introspecting schema...[/cyan]")
+    try:
+        schema = introspect(config.database_url)
+    except psycopg2.OperationalError as e:
+        console.print(f"[red]Could not connect to database: {e}[/red]")
+        sys.exit(1)
+
+    if not schema.tables:
+        console.print("[yellow]No tables found in the 'public' schema.[/yellow]")
+        sys.exit(0)
+
+    SCHEMA_CACHE.write_text(json.dumps(schema.to_dict(), indent=2))
+
+    t = RichTable(title="Discovered schema")
+    t.add_column("Table")
+    t.add_column("Columns")
+    t.add_column("Foreign Keys")
+    for table in schema.tables.values():
+        fks = ", ".join(f"{fk.column}->{fk.ref_table}.{fk.ref_column}" for fk in table.foreign_keys)
+        t.add_row(table.name, str(len(table.columns)), fks or "-")
+    console.print(t)
+    console.print(f"[green]Schema cached to {SCHEMA_CACHE}[/green]. Run 'seedloom run' next.")
+
+
+@main.command()
+@click.option("--rows", default=10, show_default=True, help="Rows to generate per table.")
+@click.option("--tables", default=None, help="Comma-separated subset of tables to seed (default: all).")
+@click.option("--dry-run", is_flag=True, help="Generate data and print it without inserting.")
+@click.option(
+    "--provider",
+    default=None,
+    help=f"Override provider from config. Supported: {', '.join(SUPPORTED_PROVIDERS)}.",
+)
+@click.option("--model", default=None, help="Override model from config.")
+@click.option("--base-url", default=None, help="Override base URL (openai_compatible or self-hosted endpoints).")
+@click.option("--host", default=None, help="Override Ollama host (default: http://localhost:11434).")
+def run(
+    rows: int,
+    tables: str | None,
+    dry_run: bool,
+    provider: str | None,
+    model: str | None,
+    base_url: str | None,
+    host: str | None,
+) -> None:
+    """Generate and insert seed data, respecting foreign key order."""
+    try:
+        config = Config.load(provider_override=provider or "")
+    except EnvironmentError as e:
+        console.print(f"[red]{e}[/red]")
+        sys.exit(1)
+
+    if not SCHEMA_CACHE.exists():
+        console.print("[red]No cached schema found. Run 'seedloom init' first.[/red]")
+        sys.exit(1)
+
+    schema = Schema.from_dict(json.loads(SCHEMA_CACHE.read_text()))
+
+    try:
+        order = resolve_seed_order(schema)
+    except CyclicDependencyError as e:
+        console.print(f"[red]{e}[/red]")
+        sys.exit(1)
+
+    if tables:
+        wanted = set(t.strip() for t in tables.split(","))
+        order = [t for t in order if t in wanted]
+
+    try:
+        active_provider = get_provider(
+            config.provider,
+            api_key=config.api_key,
+            model=model or config.model,
+            base_url=base_url or config.base_url,
+            host=host or config.host,
+        )
+    except ProviderError as e:
+        console.print(f"[red]{e}[/red]")
+        sys.exit(1)
+
+    console.print(f"[cyan]Using provider: {config.provider}[/cyan]")
+    conn = None if dry_run else psycopg2.connect(config.database_url)
+
+    fk_pools: dict[str, dict[str, list]] = {}  # table -> column -> values
+
+    referenced_columns: dict[str, set[str]] = {}
+    for t in schema.tables.values():
+        for fk in t.foreign_keys:
+            referenced_columns.setdefault(fk.ref_table, set()).add(fk.ref_column)
+
+    try:
+        for table_name in order:
+            table = schema.tables[table_name]
+            needed_columns = sorted(referenced_columns.get(table_name, set()))
+
+            to_generate = rows
+            if conn is not None:
+                existing_count = table_row_count(conn, table_name)
+                if existing_count > 0 and needed_columns:
+                    existing_values = existing_column_values(conn, table_name, needed_columns)
+                    for col, vals in existing_values.items():
+                        if vals:
+                            fk_pools.setdefault(table_name, {})[col] = vals
+
+                if existing_count >= rows:
+                    console.print(
+                        f"[yellow]Skipping '{table_name}' — already has {existing_count} row(s) "
+                        f"(>= {rows} requested).[/yellow]"
+                    )
+                    continue
+
+                to_generate = rows - existing_count
+                if existing_count > 0:
+                    console.print(
+                        f"[cyan]'{table_name}' has {existing_count} row(s); generating "
+                        f"{to_generate} more to reach {rows}...[/cyan]"
+                    )
+                else:
+                    console.print(f"[cyan]Generating {to_generate} rows for '{table_name}'...[/cyan]")
+            else:
+                console.print(f"[cyan]Generating {to_generate} rows for '{table_name}'...[/cyan]")
+
+            fk_value_pool: dict[str, list] = {}
+            for fk in table.foreign_keys:
+                parent_pool = fk_pools.get(fk.ref_table, {}).get(fk.ref_column, [])
+                if parent_pool:
+                    fk_value_pool[fk.column] = parent_pool
+
+            try:
+                generated = generate_rows(active_provider, table, to_generate, fk_value_pool)
+            except ProviderError as e:
+                console.print(f"[red]{e}[/red]")
+                sys.exit(1)
+
+            if dry_run:
+                console.print(generated)
+                continue
+
+            inserted_values = insert_rows(conn, table, generated, needed_columns)
+            for col, vals in inserted_values.items():
+                if vals:
+                    fk_pools.setdefault(table_name, {}).setdefault(col, [])
+                    fk_pools[table_name][col].extend(vals)
+
+            console.print(f"[green]Inserted {len(generated)} rows into '{table_name}'.[/green]")
+    finally:
+        if conn:
+            conn.close()
+
+    console.print("[bold green]Done.[/bold green]")
+
+
+if __name__ == "__main__":
+    main()

seedloom/config.py

@@ -0,0 +1,85 @@
+"""Configuration loading: env vars + optional .env file, no external deps."""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+from .providers import NO_KEY_REQUIRED, SUPPORTED_PROVIDERS
+
+_PROVIDER_KEY_ENV: dict[str, str] = {
+    "anthropic": "ANTHROPIC_API_KEY",
+    "openai": "OPENAI_API_KEY",
+    "gemini": "GEMINI_API_KEY",
+    "groq": "GROQ_API_KEY",
+    "together": "TOGETHER_API_KEY",
+    "fireworks": "FIREWORKS_API_KEY",
+    "openrouter": "OPENROUTER_API_KEY",
+    "deepseek": "DEEPSEEK_API_KEY",
+    "mistral": "MISTRAL_API_KEY",
+    "openai_compatible": "OPENAI_COMPATIBLE_API_KEY",
+}
+
+
+def _load_dotenv(path: Path = Path(".env")) -> None:
+    """Minimal .env loader — avoids pulling in python-dotenv as a dependency."""
+    if not path.exists():
+        return
+    for line in path.read_text().splitlines():
+        line = line.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+        key, _, value = line.partition("=")
+        key = key.strip()
+        value = value.strip().strip('"').strip("'")
+        os.environ.setdefault(key, value)
+
+
+@dataclass
+class Config:
+    database_url: str
+    provider: str = "anthropic"
+    api_key: str = ""
+    model: str = ""
+    base_url: str = ""
+    host: str = ""
+
+    @classmethod
+    def load(cls, provider_override: str = "", require_provider: bool = True) -> "Config":
+        _load_dotenv()
+        db_url = os.environ.get("DATABASE_URL", "")
+        provider = (provider_override or os.environ.get("SEEDLOOM_PROVIDER", "anthropic")).lower()
+        model = os.environ.get("SEEDLOOM_MODEL", "")
+        base_url = os.environ.get("SEEDLOOM_BASE_URL", "")
+        host = os.environ.get("SEEDLOOM_HOST", "")
+
+        missing = []
+        if not db_url:
+            missing.append("DATABASE_URL")
+
+        if provider not in SUPPORTED_PROVIDERS:
+            raise EnvironmentError(
+                f"Unknown provider '{provider}'. Supported: {', '.join(SUPPORTED_PROVIDERS)}."
+            )
+
+        api_key = ""
+        if require_provider and provider not in NO_KEY_REQUIRED:
+            key_env = _PROVIDER_KEY_ENV.get(provider, f"{provider.upper()}_API_KEY")
+            api_key = os.environ.get(key_env, "")
+            if not api_key:
+                missing.append(key_env)
+
+        if missing:
+            raise EnvironmentError(
+                f"Missing required environment variable(s): {', '.join(missing)}. "
+                "Set them in your shell or in a .env file in the current directory."
+            )
+
+        return cls(
+            database_url=db_url,
+            provider=provider,
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            host=host,
+        )

seedloom/generator.py

@@ -0,0 +1,235 @@
+"""Generate realistic seed rows for a table using a pluggable LLM provider.
+
+Key design choice: referential integrity is enforced *structurally*, not by
+hoping the model behaves. Foreign-key columns are generated as a JSON Schema
+`enum` of the actual parent-row key values already inserted — the model picks
+from real values, it can't invent a dangling reference.
+"""
+from __future__ import annotations
+
+import random
+import uuid
+from typing import Any
+
+from .models import Column, Table
+from .providers import Provider
+
+_MEDIA_KEYWORDS = (
+    "avatar",
+    "photo",
+    "image",
+    "picture",
+    "logo",
+    "banner",
+    "thumbnail",
+    "cover",
+    "icon",
+)
+
+
+def _is_media_url_column(col: Column) -> bool:
+    if col.data_type not in ("text", "character varying", "character"):
+        return False
+    name = col.name.lower()
+    return any(k in name for k in _MEDIA_KEYWORDS)
+
+
+def _random_media_url(col: Column) -> str:
+    name = col.name.lower()
+    seed = uuid.uuid4().hex[:12]
+    if "avatar" in name or "headshot" in name or "profile" in name:
+        return f"https://i.pravatar.cc/300?u={seed}"
+    if "logo" in name or "icon" in name:
+        return f"https://picsum.photos/seed/{seed}/200/200"
+    if "banner" in name or "cover" in name:
+        return f"https://picsum.photos/seed/{seed}/1200/400"
+    return f"https://picsum.photos/seed/{seed}/600/400"
+
+
+_VIDEO_KEYWORDS = (
+    "video",
+    "mp4",
+    "clip",
+    "trailer",
+    "movie",
+    "recording",
+    "footage",
+)
+
+_SAMPLE_VIDEO_URLS = (
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/BigBuckBunny.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ElephantsDream.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerBlazes.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerEscapes.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerJoyrides.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerMeltdowns.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/Sintel.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/SubaruOutbackOnStreetAndDirt.mp4",
+    "https://commondatastorage.googleapis.com/gtv-videos-bucket/sample/TearsOfSteel.mp4",
+)
+
+
+def _is_video_url_column(col: Column) -> bool:
+    if col.data_type not in ("text", "character varying", "character"):
+        return False
+    name = col.name.lower()
+    return any(k in name for k in _VIDEO_KEYWORDS)
+
+
+def _random_video_url() -> str:
+    return random.choice(_SAMPLE_VIDEO_URLS)
+
+_PG_TYPE_TO_JSON_SCHEMA: dict[str, dict[str, Any]] = {
+    "integer": {"type": "integer"},
+    "bigint": {"type": "integer"},
+    "smallint": {"type": "integer"},
+    "numeric": {"type": "number"},
+    "real": {"type": "number"},
+    "double precision": {"type": "number"},
+    "boolean": {"type": "boolean"},
+    "text": {"type": "string"},
+    "character varying": {"type": "string"},
+    "character": {"type": "string"},
+    "uuid": {"type": "string"},
+    "date": {"type": "string", "description": "ISO 8601 date, e.g. 2024-03-15"},
+    "timestamp without time zone": {"type": "string", "description": "ISO 8601 datetime"},
+    "timestamp with time zone": {"type": "string", "description": "ISO 8601 datetime with offset"},
+    "json": {"type": "object"},
+    "jsonb": {"type": "object"},
+}
+
+
+def _column_schema(col: Column) -> dict[str, Any]:
+    if col.enum_values:
+        return {"type": "string", "enum": col.enum_values}
+    base = _PG_TYPE_TO_JSON_SCHEMA.get(col.data_type, {"type": "string"})
+    schema = dict(base)
+    if col.char_max_length and schema.get("type") == "string":
+        schema["maxLength"] = col.char_max_length
+    return schema
+
+
+def build_row_schema(
+    table: Table, fk_value_pool: dict[str, list[Any]]
+) -> tuple[dict[str, Any], list[str]]:
+    """Returns (json_schema_for_one_row, list_of_generatable_column_names).
+
+    fk_value_pool maps column_name -> already-inserted parent key values,
+    for columns that are foreign keys. Columns with an empty pool (parent
+    table not seeded yet / no rows) are skipped — caller should seed in
+    dependency order so this shouldn't happen for non-nullable FKs.
+    """
+    properties: dict[str, Any] = {}
+    generatable: list[str] = []
+    fk_columns = {fk.column for fk in table.foreign_keys}
+
+    for col in table.columns:
+        if col.is_auto_generated:
+            continue
+        if col.name in fk_columns:
+            pool = [v for v in (fk_value_pool.get(col.name) or []) if v not in (None, "")]
+            if not pool:
+                continue
+            properties[col.name] = {"enum": pool}
+        else:
+            properties[col.name] = _column_schema(col)
+        generatable.append(col.name)
+
+    required = [
+        c for c in generatable
+        if not (table.column(c) and table.column(c).is_nullable)
+    ]
+    schema = {"type": "object", "properties": properties, "required": required}
+    return schema, generatable
+
+
+_NULL_LITERALS = {"null", "none", "n/a", "na", ""}
+
+
+def _sanitize_row(
+    table: Table, row: dict[str, Any], fk_value_pool: dict[str, list[Any]]
+) -> dict[str, Any]:
+    fk_columns = {fk.column for fk in table.foreign_keys}
+    cleaned: dict[str, Any] = {}
+    for key, value in row.items():
+        col = table.column(key)
+        if (
+            isinstance(value, str)
+            and value.strip().lower() in _NULL_LITERALS
+            and col
+            and col.is_nullable
+        ):
+            cleaned[key] = None
+            continue
+        if key in fk_columns:
+            pool = fk_value_pool.get(key) or []
+            if pool and value not in pool:
+                match = next((item for item in pool if str(item) == str(value)), None)
+                if match is not None:
+                    value = match
+                elif isinstance(value, int) and 1 <= value <= len(pool):
+                    value = pool[value - 1]
+                elif isinstance(value, int) and 0 <= value < len(pool):
+                    value = pool[value]
+                else:
+                    value = random.choice(pool)
+        elif (
+            isinstance(value, str)
+            and col
+            and col.char_max_length
+            and len(value) > col.char_max_length
+        ):
+            value = value[: col.char_max_length]
+        cleaned[key] = value
+    return cleaned
+
+
+def generate_rows(
+    provider: Provider,
+    table: Table,
+    count: int,
+    fk_value_pool: dict[str, list[Any]],
+    context_hint: str = "",
+) -> list[dict[str, Any]]:
+    row_schema, columns = build_row_schema(table, fk_value_pool)
+    if not columns:
+        return [{}]
+
+    tool_schema = {
+        "type": "object",
+        "properties": {
+            "rows": {
+                "type": "array",
+                "minItems": count,
+                "maxItems": count,
+                "items": row_schema,
+            }
+        },
+        "required": ["rows"],
+    }
+
+    system = (
+        "You generate realistic, internally-consistent fake data for database seeding. "
+        "Values should look like real-world data (plausible names, emails matching names, "
+        "believable dates/amounts), not placeholder text like 'test1'. "
+        "Never reuse the exact same value twice within the batch unless the column is clearly "
+        "meant to repeat (e.g. a status field)."
+    )
+    user_prompt = (
+        f"Generate {count} realistic rows for the table `{table.name}`.\n"
+        f"{context_hint}\n"
+        "Call the generate_rows tool with the data."
+    )
+
+    rows = provider.generate(system, user_prompt, tool_schema, tool_name="generate_rows")
+    if not rows:
+        raise RuntimeError(f"Provider did not return any rows for table {table.name}")
+    for row in rows:
+        for c in columns:
+            col = table.column(c)
+            if col and _is_media_url_column(col):
+                row[c] = _random_media_url(col)
+            elif col and _is_video_url_column(col):
+                row[c] = _random_video_url()
+    return [_sanitize_row(table, row, fk_value_pool) for row in rows]

seedloom/graph.py

@@ -0,0 +1,50 @@
+"""Resolve the order tables must be seeded in, based on foreign key dependencies.
+
+A table depends on every table its foreign keys point to (excluding self-references,
+which are seeded as NULL-first-then-update or just left nullable). Raises on a
+genuine cycle between two *different* tables, since that can't be seeded without
+deferred constraints (out of scope for v1 — surfaced as a clear error instead of
+a silent wrong answer).
+"""
+from __future__ import annotations
+
+from .models import Schema
+
+
+class CyclicDependencyError(Exception):
+    pass
+
+
+def resolve_seed_order(schema: Schema) -> list[str]:
+    deps: dict[str, set[str]] = {}
+    for table in schema.tables.values():
+        table_deps = set()
+        for fk in table.foreign_keys:
+            if fk.ref_table != table.name and fk.ref_table in schema.tables:
+                table_deps.add(fk.ref_table)
+        deps[table.name] = table_deps
+
+    ordered: list[str] = []
+    visited: set[str] = set()
+    in_progress: set[str] = set()
+
+    def visit(name: str, path: list[str]) -> None:
+        if name in visited:
+            return
+        if name in in_progress:
+            raise CyclicDependencyError(
+                f"Cyclic foreign key dependency detected: {' -> '.join(path + [name])}. "
+                "seedloom can't resolve insert order for mutually-dependent tables in v1 — "
+                "consider making one of the FKs nullable and seeding it in a second pass."
+            )
+        in_progress.add(name)
+        for dep in deps.get(name, set()):
+            visit(dep, path + [name])
+        in_progress.discard(name)
+        visited.add(name)
+        ordered.append(name)
+
+    for table_name in schema.tables:
+        visit(table_name, [])
+
+    return ordered

seedloom/inserter.py

@@ -0,0 +1,93 @@
+"""Insert generated rows into Postgres, returning values for whichever
+columns other tables' foreign keys point at, so downstream tables always
+have a real, valid pool to pick from — not just the primary key."""
+from __future__ import annotations
+
+from typing import Any
+
+import psycopg2
+import psycopg2.extras
+from rich.console import Console
+
+from .models import Table
+
+console = Console()
+
+
+def _adapt_value(value: Any) -> Any:
+    if isinstance(value, (dict, list)):
+        return psycopg2.extras.Json(value)
+    return value
+
+
+def table_row_count(conn, table_name: str) -> int:
+    with conn.cursor() as cur:
+        cur.execute(f'SELECT COUNT(*) FROM "{table_name}"')
+        return cur.fetchone()[0]
+
+
+def existing_column_values(conn, table_name: str, columns: list[str]) -> dict[str, list[Any]]:
+    """Fetch current values for the given columns (e.g. columns other tables'
+    foreign keys reference — which may or may not be the primary key)."""
+    if not columns:
+        return {}
+    col_list = ", ".join(f'"{c}"' for c in columns)
+    result: dict[str, list[Any]] = {c: [] for c in columns}
+    with conn.cursor() as cur:
+        cur.execute(f'SELECT {col_list} FROM "{table_name}"')
+        for record in cur.fetchall():
+            for i, c in enumerate(columns):
+                result[c].append(record[i])
+    return result
+
+
+def insert_rows(
+    conn, table: Table, rows: list[dict[str, Any]], needed_columns: list[str] | None = None
+) -> dict[str, list[Any]]:
+    """Insert rows for one table. Returns {column_name: [values]} for each
+    column in `needed_columns` (typically every column some other table's
+    foreign key references), covering both auto-generated and model-supplied
+    values.
+
+    Rows that collide with an existing unique/primary key are skipped (logged,
+    not treated as an error) rather than aborting the run.
+    """
+    if not rows or not rows[0]:
+        return {}
+
+    needed_columns = needed_columns or []
+    columns = sorted({c for row in rows for c in row.keys()})
+    returning_clause = ""
+    if needed_columns:
+        returning_clause = " RETURNING " + ", ".join(f'"{c}"' for c in needed_columns)
+
+    col_list = ", ".join(f'"{c}"' for c in columns)
+    placeholders = ", ".join(f"%({c})s" for c in columns)
+    query = (
+        f'INSERT INTO "{table.name}" ({col_list}) VALUES ({placeholders}) '
+        f'ON CONFLICT DO NOTHING{returning_clause}'
+    )
+
+    collected: dict[str, list[Any]] = {c: [] for c in needed_columns}
+    skipped = 0
+    with conn.cursor() as cur:
+        for row in rows:
+            adapted_row = {c: _adapt_value(row.get(c)) for c in columns}
+            cur.execute(query, adapted_row)
+            if returning_clause:
+                result = cur.fetchone()
+                if result is not None:
+                    for i, c in enumerate(needed_columns):
+                        collected[c].append(result[i])
+                else:
+                    skipped += 1
+            elif cur.rowcount == 0:
+                skipped += 1
+    conn.commit()
+
+    if skipped:
+        console.print(
+            f"[yellow]  {skipped} row(s) already existed in '{table.name}', skipped.[/yellow]"
+        )
+
+    return collected