seedforge 0.1.0__py3-none-any.whl

seedforge/__init__.py

@@ -0,0 +1,3 @@
+"""SeedForge - AI-powered test data generator for databases."""
+
+__version__ = "0.1.0"

seedforge/ai.py

@@ -0,0 +1,150 @@
+"""AI-слой: Claude анализирует схему и генерирует контекстно-зависимые данные.
+
+Используется как premium-фича. Отправляет только метаданные схемы (имена таблиц/колонок),
+НЕ реальные данные пользователя.
+"""
+
+import json
+import os
+
+
+def get_ai_generators(tables: dict, api_key: str | None = None) -> dict:
+    """Попросить Claude проанализировать схему и предложить генераторы.
+
+    Возвращает dict: {"table.column": {"generator": "...", "args": {...}}}
+    """
+    api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        return {}
+
+    # Собираем метаданные схемы (только имена, типы, связи — НЕ данные)
+    schema_desc = _build_schema_description(tables)
+
+    try:
+        import anthropic
+        client = anthropic.Anthropic(api_key=api_key)
+
+        response = client.messages.create(
+            model="claude-haiku-4-5-20251001",
+            max_tokens=4096,
+            messages=[{
+                "role": "user",
+                "content": f"""Analyze this database schema and suggest realistic test data generators for each column.
+
+Schema:
+{schema_desc}
+
+For each column, suggest what kind of realistic data should be generated.
+Focus on columns where the name alone isn't enough to determine the right generator
+(e.g., "name" in an "organizations" table should be a company name, not a person name).
+
+Return a JSON object where keys are "table.column" and values have:
+- "generator": one of the Faker methods (e.g., "company", "name", "sentence") or a custom description
+- "example": an example value
+- "values": (optional) a list of realistic values to choose from (for status/type/role fields)
+
+Only include columns where your suggestion differs from the obvious default.
+Return ONLY valid JSON, no markdown, no explanation."""
+            }],
+        )
+
+        text = response.content[0].text.strip()
+        # Убираем markdown если есть
+        if text.startswith("```"):
+            text = text.split("\n", 1)[1]
+            text = text.rsplit("```", 1)[0]
+
+        return json.loads(text)
+
+    except ImportError:
+        return {}
+    except Exception:
+        return {}
+
+
+def generate_with_ai(
+    table_name: str,
+    columns: list[dict],
+    row_count: int,
+    api_key: str | None = None,
+    context: str = "",
+) -> list[dict] | None:
+    """Генерировать данные для таблицы через Claude.
+
+    Используется для небольших таблиц (до 50 строк) где нужна высокая реалистичность.
+    Для больших объёмов используем rule-based генерацию.
+    """
+    api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        return None
+
+    # Ограничиваем до 50 строк для AI-генерации (стоимость)
+    row_count = min(row_count, 50)
+
+    col_desc = ", ".join(
+        f"{c['name']} ({c['type']}{'?' if c.get('nullable') else ''})"
+        for c in columns
+    )
+
+    try:
+        import anthropic
+        client = anthropic.Anthropic(api_key=api_key)
+
+        prompt = f"""Generate {row_count} realistic rows for the "{table_name}" table.
+Columns: {col_desc}
+{f"Context: {context}" if context else ""}
+
+Requirements:
+- Data must be realistic and internally consistent
+- If column is a name in an organization table, use company names
+- If column is a status, use realistic statuses
+- Dates should be recent (within last 2 years)
+- Respect nullable columns (occasionally set to null)
+
+Return a JSON array of objects. Each object has column names as keys.
+Return ONLY valid JSON, no markdown, no explanation."""
+
+        response = client.messages.create(
+            model="claude-haiku-4-5-20251001",
+            max_tokens=8192,
+            messages=[{"role": "user", "content": prompt}],
+        )
+
+        text = response.content[0].text.strip()
+        if text.startswith("```"):
+            text = text.split("\n", 1)[1]
+            text = text.rsplit("```", 1)[0]
+
+        return json.loads(text)
+
+    except ImportError:
+        return None
+    except Exception:
+        return None
+
+
+def _build_schema_description(tables: dict) -> str:
+    """Построить текстовое описание схемы для AI."""
+    lines = []
+    for table_name, table in tables.items():
+        cols = []
+        for col in table.columns:
+            parts = [f"{col.name} {col.data_type}"]
+            if col.is_primary:
+                parts.append("PK")
+            if col.fk_table:
+                parts.append(f"FK→{col.fk_table}.{col.fk_column}")
+            if not col.nullable:
+                parts.append("NOT NULL")
+            if col.is_unique:
+                parts.append("UNIQUE")
+            if col.enum_values:
+                parts.append(f"ENUM({', '.join(col.enum_values[:5])})")
+            cols.append(" ".join(parts))
+
+        lines.append(f"TABLE {table_name}:")
+        for c in cols:
+            lines.append(f"  - {c}")
+        lines.append("")
+
+    return "\n".join(lines)

seedforge/cli.py

@@ -0,0 +1,311 @@
+"""CLI-интерфейс SeedForge."""
+
+import typer
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+from rich import print as rprint
+from typing import Optional
+from pathlib import Path
+
+from seedforge.config import Config, DEFAULT_CONFIG_FILE
+from seedforge.introspector import create_introspector
+from seedforge.graph import DependencyGraph
+from seedforge.generators import DataGenerator
+from seedforge.inserter import BatchInserter
+
+app = typer.Typer(
+    name="seedforge",
+    help="AI-powered test data generator. Reads your DB schema, generates realistic FK-valid data.",
+    no_args_is_help=True,
+)
+console = Console()
+
+
+@app.command()
+def connect(
+    db_url: str = typer.Argument(..., help="PostgreSQL connection string (postgresql://user:pass@host:port/dbname)"),
+    save: bool = typer.Option(True, help="Save connection to .seedforge.yaml"),
+):
+    """Connect to a PostgreSQL database and save the connection."""
+    console.print(f"\n[bold blue]Connecting to database...[/bold blue]")
+
+    try:
+        introspector = create_introspector(db_url)
+        info = introspector.get_db_info()
+        introspector.close()
+
+        console.print(f"[bold green]Connected![/bold green] {info['database']} @ {info['host']}")
+        console.print(f"  Tables: {info['table_count']}")
+        console.print(f"  PostgreSQL: {info['version']}")
+
+        if save:
+            config = Config.load()
+            config.db_url = db_url
+            config.save()
+            console.print(f"\n[dim]Saved to {DEFAULT_CONFIG_FILE}[/dim]")
+
+    except Exception as e:
+        console.print(f"[bold red]Connection failed:[/bold red] {e}")
+        raise typer.Exit(1)
+
+
+@app.command()
+def inspect(
+    db_url: Optional[str] = typer.Argument(None, help="PostgreSQL connection string"),
+    schema: str = typer.Option("public", help="Database schema to inspect"),
+):
+    """Inspect database schema: tables, columns, foreign keys."""
+    db_url = _resolve_db_url(db_url)
+
+    introspector = Introspector(db_url)
+    tables = introspector.get_tables(schema)
+    introspector.close()
+
+    if not tables:
+        console.print("[yellow]No tables found.[/yellow]")
+        raise typer.Exit(0)
+
+    # Граф зависимостей + порядок вставки
+    graph = DependencyGraph(tables)
+    order = graph.topological_sort()
+
+    console.print(f"\n[bold]Found {len(tables)} tables[/bold] (insertion order):\n")
+
+    for i, table_name in enumerate(order, 1):
+        table = tables[table_name]
+        t = Table(title=f"{i}. {table_name}", title_style="bold cyan", show_lines=False)
+        t.add_column("Column", style="white")
+        t.add_column("Type", style="green")
+        t.add_column("Nullable", style="yellow", width=8)
+        t.add_column("FK → ", style="magenta")
+
+        for col in table.columns:
+            fk_str = f"{col.fk_table}.{col.fk_column}" if col.fk_table else ""
+            nullable = "YES" if col.nullable else "NO"
+            t.add_row(col.name, col.data_type, nullable, fk_str)
+
+        console.print(t)
+        console.print()
+
+    # Статистика
+    total_cols = sum(len(t.columns) for t in tables.values())
+    total_fks = sum(1 for t in tables.values() for c in t.columns if c.fk_table)
+    console.print(f"[dim]Total: {len(tables)} tables, {total_cols} columns, {total_fks} foreign keys[/dim]\n")
+
+
+@app.command()
+def generate(
+    db_url: Optional[str] = typer.Argument(None, help="PostgreSQL connection string"),
+    rows: int = typer.Option(100, "--rows", "-r", help="Rows per table"),
+    schema: str = typer.Option("public", help="Database schema"),
+    seed: Optional[int] = typer.Option(None, "--seed", "-s", help="Random seed for deterministic generation"),
+    export: Optional[str] = typer.Option(None, "--export", "-e", help="Export to file (sql/json) instead of inserting"),
+    tables: Optional[str] = typer.Option(None, "--tables", "-t", help="Comma-separated list of tables to fill"),
+    clean: bool = typer.Option(False, "--clean", help="TRUNCATE tables before inserting (CASCADE)"),
+    dry_run: bool = typer.Option(False, "--dry-run", help="Show what would be generated without inserting"),
+):
+    """Generate realistic test data and insert into the database."""
+    db_url = _resolve_db_url(db_url)
+
+    with console.status("[bold blue]Reading schema...[/bold blue]"):
+        introspector = create_introspector(db_url)
+        all_tables = introspector.get_tables(schema)
+
+    if not all_tables:
+        console.print("[yellow]No tables found.[/yellow]")
+        introspector.close()
+        raise typer.Exit(0)
+
+    # Фильтр таблиц
+    if tables:
+        table_list = [t.strip() for t in tables.split(",")]
+        filtered = {k: v for k, v in all_tables.items() if k in table_list}
+        if not filtered:
+            console.print(f"[red]Tables not found: {tables}[/red]")
+            introspector.close()
+            raise typer.Exit(1)
+        # Добавляем родительские таблицы (для FK)
+        for tbl in list(filtered.values()):
+            for col in tbl.columns:
+                if col.fk_table and col.fk_table in all_tables and col.fk_table not in filtered:
+                    filtered[col.fk_table] = all_tables[col.fk_table]
+        all_tables = filtered
+
+    # Порядок вставки
+    graph = DependencyGraph(all_tables)
+    order = graph.topological_sort()
+
+    console.print(f"\n[bold]Generating {rows} rows for {len(order)} tables[/bold]")
+    if seed is not None:
+        console.print(f"[dim]Seed: {seed}[/dim]")
+    console.print()
+
+    # Генерация данных
+    generator = DataGenerator(seed=seed)
+    generated_data = {}
+
+    for table_name in order:
+        table = all_tables[table_name]
+        with console.status(f"[blue]Generating {table_name}...[/blue]"):
+            data = generator.generate_table(table, rows, generated_data)
+            generated_data[table_name] = data
+        console.print(f"  [green]✓[/green] {table_name}: {len(data)} rows")
+
+    # Экспорт или вставка
+    if export:
+        _export_data(export, generated_data, all_tables, order)
+    elif dry_run:
+        console.print(f"\n[yellow]Dry run — no data inserted.[/yellow]")
+        _show_preview(generated_data, order)
+    else:
+        engine = introspector.get_db_info().get("engine", "PostgreSQL")
+        inserter = BatchInserter(introspector.connection, engine=engine)
+        if clean:
+            with console.status("[yellow]Cleaning tables...[/yellow]"):
+                inserter.truncate_tables(order)
+            console.print("[yellow]Tables truncated.[/yellow]\n")
+
+        with console.status("[bold blue]Inserting data...[/bold blue]"):
+            inserter.insert_all(generated_data, all_tables, order)
+        console.print(f"\n[bold green]Done![/bold green] Inserted {sum(len(d) for d in generated_data.values())} rows into {len(order)} tables.\n")
+
+    introspector.close()
+
+
+@app.command()
+def ai_generate(
+    db_url: Optional[str] = typer.Argument(None, help="Database connection string"),
+    rows: int = typer.Option(20, "--rows", "-r", help="Rows per table (max 50 for AI)"),
+    schema: str = typer.Option("public", help="Database schema"),
+    api_key: Optional[str] = typer.Option(None, "--api-key", envvar="ANTHROPIC_API_KEY", help="Anthropic API key"),
+    export: Optional[str] = typer.Option(None, "--export", "-e", help="Export to file (sql/json)"),
+):
+    """Generate data using Claude AI for maximum realism (requires ANTHROPIC_API_KEY)."""
+    from seedforge.ai import generate_with_ai
+
+    if not api_key:
+        console.print("[red]Set ANTHROPIC_API_KEY or use --api-key[/red]")
+        raise typer.Exit(1)
+
+    db_url = _resolve_db_url(db_url)
+    rows = min(rows, 50)
+
+    with console.status("[bold blue]Reading schema...[/bold blue]"):
+        introspector = create_introspector(db_url)
+        all_tables = introspector.get_tables(schema)
+
+    if not all_tables:
+        console.print("[yellow]No tables found.[/yellow]")
+        introspector.close()
+        raise typer.Exit(0)
+
+    from seedforge.graph import DependencyGraph
+    graph = DependencyGraph(all_tables)
+    order = graph.topological_sort()
+
+    console.print(f"\n[bold]AI generating {rows} rows for {len(order)} tables[/bold]\n")
+
+    generated_data = {}
+    for table_name in order:
+        table = all_tables[table_name]
+        columns = [
+            {"name": c.name, "type": c.data_type, "nullable": c.nullable}
+            for c in table.columns
+            if not (c.is_primary and c.is_serial)
+        ]
+
+        with console.status(f"[blue]AI generating {table_name}...[/blue]"):
+            data = generate_with_ai(table_name, columns, rows, api_key=api_key)
+
+        if data:
+            generated_data[table_name] = data
+            console.print(f"  [green]✓[/green] {table_name}: {len(data)} rows (AI)")
+        else:
+            console.print(f"  [yellow]⚠[/yellow] {table_name}: AI failed, skipping")
+
+    if export:
+        _export_data(export, generated_data, all_tables, order)
+    else:
+        console.print(f"\n[bold green]Done![/bold green] Generated {sum(len(d) for d in generated_data.values())} rows.")
+        _show_preview(generated_data, order)
+
+    introspector.close()
+
+
+@app.command()
+def version():
+    """Show SeedForge version."""
+    from seedforge import __version__
+    console.print(f"SeedForge v{__version__}")
+
+
+def _resolve_db_url(db_url: Optional[str]) -> str:
+    """Получить DB URL из аргумента или конфига."""
+    if db_url:
+        return db_url
+    config = Config.load()
+    if config.db_url:
+        return config.db_url
+    console.print("[red]No database URL. Run 'seedforge connect <url>' first or pass it as argument.[/red]")
+    raise typer.Exit(1)
+
+
+def _export_data(format: str, data: dict, tables: dict, order: list):
+    """Экспорт данных в файл."""
+    if format == "sql":
+        from seedforge.inserter import BatchInserter
+        sql = BatchInserter.generate_sql(data, tables, order)
+        output_file = "seedforge_export.sql"
+        Path(output_file).write_text(sql)
+        console.print(f"\n[green]Exported to {output_file}[/green]")
+    elif format == "json":
+        import json
+        output_file = "seedforge_export.json"
+        # Конвертируем в JSON-сериализуемый формат
+        serializable = {}
+        for table_name in order:
+            serializable[table_name] = [
+                {k: _json_safe(v) for k, v in row.items()}
+                for row in data[table_name]
+            ]
+        Path(output_file).write_text(json.dumps(serializable, indent=2, ensure_ascii=False))
+        console.print(f"\n[green]Exported to {output_file}[/green]")
+    else:
+        console.print(f"[red]Unknown format: {format}. Use 'sql' or 'json'.[/red]")
+
+
+def _json_safe(value):
+    """Конвертировать значение в JSON-безопасный тип."""
+    import datetime
+    import decimal
+    import uuid
+    if isinstance(value, (datetime.date, datetime.datetime)):
+        return value.isoformat()
+    if isinstance(value, decimal.Decimal):
+        return float(value)
+    if isinstance(value, uuid.UUID):
+        return str(value)
+    return value
+
+
+def _show_preview(data: dict, order: list):
+    """Показать превью данных."""
+    for table_name in order:
+        rows = data[table_name]
+        if not rows:
+            continue
+        t = Table(title=table_name, title_style="bold cyan")
+        cols = list(rows[0].keys())
+        for col in cols:
+            t.add_column(col)
+        for row in rows[:5]:
+            t.add_row(*[str(row.get(c, ""))[:40] for c in cols])
+        if len(rows) > 5:
+            t.add_row(*[f"... ({len(rows)} total)" if i == 0 else "..." for i, c in enumerate(cols)])
+        console.print(t)
+        console.print()
+
+
+if __name__ == "__main__":
+    app()

seedforge/config.py

@@ -0,0 +1,53 @@
+"""Конфигурация SeedForge (.seedforge.yaml)."""
+
+import yaml
+from pathlib import Path
+from dataclasses import dataclass, field
+
+DEFAULT_CONFIG_FILE = ".seedforge.yaml"
+
+
+@dataclass
+class Config:
+    db_url: str = ""
+    default_rows: int = 100
+    default_schema: str = "public"
+    seed: int | None = None
+    exclude_tables: list[str] = field(default_factory=list)
+    # Кастомные генераторы для колонок
+    # Формат: {"table.column": {"type": "faker_method", "args": {...}}}
+    custom_generators: dict = field(default_factory=dict)
+
+    @classmethod
+    def load(cls, path: str = DEFAULT_CONFIG_FILE) -> "Config":
+        """Загрузить конфиг из файла."""
+        config_path = Path(path)
+        if not config_path.exists():
+            return cls()
+        with open(config_path) as f:
+            data = yaml.safe_load(f) or {}
+        return cls(
+            db_url=data.get("db_url", ""),
+            default_rows=data.get("default_rows", 100),
+            default_schema=data.get("default_schema", "public"),
+            seed=data.get("seed"),
+            exclude_tables=data.get("exclude_tables", []),
+            custom_generators=data.get("custom_generators", {}),
+        )
+
+    def save(self, path: str = DEFAULT_CONFIG_FILE):
+        """Сохранить конфиг в файл."""
+        data = {
+            "db_url": self.db_url,
+            "default_rows": self.default_rows,
+            "default_schema": self.default_schema,
+        }
+        if self.seed is not None:
+            data["seed"] = self.seed
+        if self.exclude_tables:
+            data["exclude_tables"] = self.exclude_tables
+        if self.custom_generators:
+            data["custom_generators"] = self.custom_generators
+
+        with open(path, "w") as f:
+            yaml.dump(data, f, default_flow_style=False, allow_unicode=True)