pretensor 0.1.0__py3-none-any.whl

pretensor/cli/commands/analyze.py

@@ -0,0 +1,201 @@
+"""``pretensor analyze`` — scan a repo for SQL and link code to graph tables."""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+from pathlib import Path
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from pretensor.cli import constants as cli_constants
+from pretensor.cli.config_file import get_cli_config, resolve_aliased_path_option
+from pretensor.enrichment.analyze.pipeline import (
+    EmptyGraphError,
+    run_analyze_enrichment,
+)
+from pretensor.enrichment.analyze.summary import AnalyzeSummary
+from pretensor.mcp.service_registry import (
+    load_registry,
+    open_store_for_entry,
+    release_store,
+    resolve_registry_entry,
+)
+
+__all__ = ["register_analyze_command"]
+
+_EXIT_ERROR = 1
+_MAX_DISPLAY_ROWS = 20
+
+
+def register_analyze_command(app: typer.Typer, *, console: Console) -> None:
+    """Register ``pretensor analyze`` onto ``app``."""
+
+    @app.command("analyze")
+    def analyze_command(
+        path: Path = typer.Argument(
+            Path("."),
+            help="Repository root to scan (default: current directory).",
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        connection: str = typer.Option(
+            ...,
+            "--connection",
+            help="Required. Connection whose tables SQL refs resolve against.",
+        ),
+        service: str | None = typer.Option(
+            None,
+            "--service",
+            help="Consumer-service label (default: scanned directory basename).",
+        ),
+        include: list[str] = typer.Option(
+            [],
+            "--include",
+            help="Gitignore-style glob to include (repeatable).",
+        ),
+        exclude: list[str] = typer.Option(
+            [],
+            "--exclude",
+            help="Gitignore-style glob to exclude (repeatable).",
+        ),
+        max_file_bytes: int = typer.Option(
+            1_000_000,
+            "--max-file-bytes",
+            help="Skip files larger than this many bytes.",
+        ),
+        min_confidence: float = typer.Option(
+            0.6,
+            "--min-confidence",
+            help="Minimum confidence to emit a CONSUMES edge.",
+        ),
+        default_schema: str = typer.Option(
+            "public",
+            "--default-schema",
+            help="Schema assumed for unqualified table names in scanned SQL.",
+        ),
+        dry_run: bool = typer.Option(
+            False,
+            "--dry-run",
+            help="Parse and print the plan only; make no graph writes.",
+        ),
+        as_json: bool = typer.Option(
+            False,
+            "--json",
+            help="Emit the machine-readable summary as JSON (suppresses the table).",
+        ),
+        state_dir: Path = typer.Option(
+            cli_constants.DEFAULT_STATE_DIR,
+            "--state-dir",
+            help="Directory containing registry.json.",
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        graph_dir: Path | None = typer.Option(
+            None,
+            "--graph-dir",
+            hidden=True,
+            help="Deprecated alias for --state-dir.",
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        ctx: typer.Context = typer.Option(None, hidden=True),
+    ) -> None:
+        """Scan a repository for SQL string literals and link them to graph tables."""
+        cli_config = get_cli_config(ctx)
+        state_dir = resolve_aliased_path_option(
+            ctx,
+            primary_param="state_dir",
+            primary_value=state_dir,
+            alias_param="graph_dir",
+            alias_value=graph_dir,
+            config_value=cli_config.state_dir,
+        )
+
+        service_name = service or path.name
+
+        reg = load_registry(state_dir)
+        entry = resolve_registry_entry(reg, connection)
+        if entry is None:
+            _empty_graph_error(console, connection)
+            raise typer.Exit(_EXIT_ERROR)
+
+        store = open_store_for_entry(entry)
+        try:
+            summary = run_analyze_enrichment(
+                path,
+                store,
+                entry.connection_name,
+                service_name=service_name,
+                includes=include,
+                excludes=exclude,
+                max_file_bytes=max_file_bytes,
+                min_confidence=min_confidence,
+                default_schema=default_schema,
+                dry_run=dry_run,
+            )
+        except EmptyGraphError:
+            _empty_graph_error(console, entry.connection_name)
+            raise typer.Exit(_EXIT_ERROR) from None
+        finally:
+            release_store(store)
+
+        if as_json:
+            typer.echo(json.dumps(dataclasses.asdict(summary), indent=2))
+            return
+
+        _render_summary(console, summary, service_name=service_name, dry_run=dry_run)
+
+
+def _empty_graph_error(console: Console, connection: str) -> None:
+    """Print the hard-error hint; no auto-index is performed (by design)."""
+    console.print(
+        f"[red]The graph for connection '{connection}' is empty.[/red] Run:\n"
+        f"  pretensor index --connection {connection}\n"
+        "first, then re-run pretensor analyze."
+    )
+
+
+def _render_summary(
+    console: Console,
+    summary: AnalyzeSummary,
+    *,
+    service_name: str,
+    dry_run: bool,
+) -> None:
+    rows = sorted(summary.rows, key=lambda r: r.confidence, reverse=True)
+    shown = rows[:_MAX_DISPLAY_ROWS]
+
+    title = (
+        "pretensor analyze (dry run — no writes)" if dry_run else "pretensor analyze"
+    )
+    table = Table(title=title)
+    table.add_column("service")
+    table.add_column("table")
+    table.add_column("op")
+    table.add_column("file:line")
+    table.add_column("confidence", justify="right")
+    for r in shown:
+        line = (
+            f"{r.file_path}:{r.line_start}"
+            if r.line_start == r.line_end
+            else f"{r.file_path}:{r.line_start}-{r.line_end}"
+        )
+        table.add_row(r.service_name, r.table, r.op, line, f"{r.confidence:.2f}")
+    console.print(table)
+
+    if len(rows) > len(shown):
+        console.print(f"[dim]… showing top {len(shown)} of {len(rows)} edges[/dim]")
+
+    verb = "would write" if dry_run else "wrote"
+    console.print(
+        f"Scanned {summary.files_scanned} files, "
+        f"found {summary.sql_candidates_found} SQL candidates, "
+        f"{verb} {summary.consumers_written} consumers / {summary.edges_written} edges "
+        f"({summary.cross_connection_dropped} dropped: cross-connection)."
+    )

pretensor/cli/commands/connections/__init__.py

@@ -0,0 +1,7 @@
+"""Database connection registration CLI commands (``add``, ``remove``)."""
+
+from __future__ import annotations
+
+from pretensor.cli.commands.connections.register import register_connection_commands
+
+__all__ = ["register_connection_commands"]

pretensor/cli/commands/connections/add_remove.py

@@ -0,0 +1,126 @@
+"""``add`` and ``remove`` cross-DB registry commands."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+from pretensor.cli import constants as cli_constants
+from pretensor.cli.config_file import (
+    get_cli_config,
+    resolve_optional_str_option,
+    resolve_path_option,
+)
+from pretensor.cli.paths import default_connection_name, unified_graph_path
+from pretensor.core.dsn_crypto import DSNEncryptor
+from pretensor.core.registry import GraphRegistry
+from pretensor.introspection.models.dsn import (
+    connection_config_from_url,
+    redact_dsn,
+    registry_dialect_for,
+)
+
+
+def _encryptor_for(state_dir: Path) -> DSNEncryptor:
+    return DSNEncryptor(state_dir / "keystore")
+
+
+def register_add_remove_commands(app: typer.Typer) -> None:
+    @app.command("add")
+    def add_command(
+        dsn: str = typer.Argument(
+            ..., help="Database DSN URL (postgresql:// or snowflake://)."
+        ),
+        name: str | None = typer.Option(
+            None,
+            "--name",
+            "-n",
+            help="Logical connection name (default: DB name from DSN).",
+        ),
+        state_dir: Path = typer.Option(
+            cli_constants.DEFAULT_STATE_DIR,
+            "--state-dir",
+            help="State directory for registry and graphs.",
+            file_okay=False,
+            dir_okay=True,
+            writable=True,
+            resolve_path=True,
+        ),
+        ctx: typer.Context = typer.Option(None, hidden=True),
+    ) -> None:
+        """Register a DSN (encrypted); run `pretensor index` to build the graph."""
+        cli_config = get_cli_config(ctx)
+        state_dir = resolve_path_option(
+            ctx,
+            param_name="state_dir",
+            cli_value=state_dir,
+            config_value=cli_config.state_dir,
+        )
+        name = resolve_optional_str_option(
+            ctx,
+            param_name="name",
+            cli_value=name,
+            config_value=cli_config.connection_defaults.name,
+        )
+        console = Console()
+        connection_name = name or default_connection_name(dsn)
+        enc = _encryptor_for(state_dir)
+        reg_path = state_dir / cli_constants.REGISTRY_FILENAME
+        reg = GraphRegistry(reg_path).load()
+        graph_path = unified_graph_path(state_dir)
+        cfg = connection_config_from_url(dsn, connection_name)
+        reg.upsert(
+            connection_name=connection_name,
+            database=cfg.database or connection_name,
+            dsn=dsn,
+            graph_path=graph_path,
+            unified_graph_path=graph_path,
+            encrypt_dsn=True,
+            encryptor=enc,
+            indexed_at=None,
+            dialect=registry_dialect_for(cfg.type),
+        )
+        reg.save()
+        console.print(
+            f"[green]Registered[/green] connection {connection_name!r} (DSN encrypted)."
+        )
+        # Mask the password in the hint so the DSN does not land in terminal
+        # scrollback, screen-shares or CI logs.
+        console.print(
+            f"Run: pretensor index '{redact_dsn(dsn)}' "
+            f"--name {connection_name} --unified"
+        )
+
+    @app.command("remove")
+    def remove_command(
+        name: str = typer.Argument(
+            ..., help="Connection name to remove from registry."
+        ),
+        state_dir: Path = typer.Option(
+            cli_constants.DEFAULT_STATE_DIR,
+            "--state-dir",
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        ctx: typer.Context = typer.Option(None, hidden=True),
+    ) -> None:
+        """Remove a connection from the registry (does not delete the unified graph file)."""
+        cli_config = get_cli_config(ctx)
+        state_dir = resolve_path_option(
+            ctx,
+            param_name="state_dir",
+            cli_value=state_dir,
+            config_value=cli_config.state_dir,
+        )
+        console = Console()
+        reg_path = state_dir / cli_constants.REGISTRY_FILENAME
+        reg = GraphRegistry(reg_path).load()
+        if reg.get(name) is None:
+            console.print(f"[red]Unknown connection[/red] {name!r}")
+            raise typer.Exit(1)
+        reg.remove(name)
+        reg.save()
+        console.print(f"[green]Removed[/green] {name!r} from registry.")

pretensor/cli/commands/connections/register.py

@@ -0,0 +1,12 @@
+"""Wire connection-management Typer commands onto the root app."""
+
+from __future__ import annotations
+
+import typer
+
+from pretensor.cli.commands.connections.add_remove import register_add_remove_commands
+
+
+def register_connection_commands(app: typer.Typer) -> None:
+    """Attach ``add`` and ``remove`` to ``app``."""
+    register_add_remove_commands(app)

pretensor/cli/commands/export.py

@@ -0,0 +1,131 @@
+"""``pretensor export`` command."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+from pretensor.cli import constants as cli_constants
+from pretensor.cli.config_file import get_cli_config, resolve_aliased_path_option
+from pretensor.core.portable_export import export_graph_payload
+from pretensor.core.registry import DatabaseNotFoundError, GraphRegistry
+from pretensor.core.store import KuzuStore
+
+
+def register_export_command(app: typer.Typer, *, console: Console) -> None:
+    @app.command("export")
+    def export_command(
+        database: str = typer.Option(
+            ...,
+            "--database",
+            "-d",
+            help="Registry connection name to export (see `pretensor list`).",
+        ),
+        output: Path = typer.Option(
+            ...,
+            "--output",
+            "-o",
+            help="Target JSON path for portable graph export.",
+            file_okay=True,
+            dir_okay=False,
+            writable=True,
+            resolve_path=True,
+        ),
+        state_dir: Path = typer.Option(
+            cli_constants.DEFAULT_STATE_DIR,
+            "--state-dir",
+            help="Directory containing registry.json and graph files.",
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        graph_dir: Path | None = typer.Option(
+            None,
+            "--graph-dir",
+            hidden=True,
+            help="Deprecated alias for --state-dir.",
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        ctx: typer.Context = typer.Option(
+            None,
+            hidden=True,
+        ),
+    ) -> None:
+        """Dump one indexed graph to portable pretty-printed JSON."""
+        cli_config = get_cli_config(ctx)
+        state_dir = resolve_aliased_path_option(
+            ctx,
+            primary_param="state_dir",
+            primary_value=state_dir,
+            alias_param="graph_dir",
+            alias_value=graph_dir,
+            config_value=cli_config.state_dir,
+        )
+        reg_path = state_dir / cli_constants.REGISTRY_FILENAME
+        if not reg_path.exists():
+            console.print("[red]No registry found.[/red] Run `pretensor index` first.")
+            raise typer.Exit(1)
+        try:
+            reg = GraphRegistry(reg_path).load()
+            entry = reg.require(database)
+        except DatabaseNotFoundError:
+            console.print(
+                f"[red]Unknown connection {database!r}.[/red] "
+                "Use `pretensor list` to discover available names."
+            )
+            raise typer.Exit(1)
+        except Exception as exc:
+            console.print(
+                f"[red]Cannot read registry file {reg_path}:[/red] {exc}\n"
+                "The registry.json may be corrupt. Delete it and re-run `pretensor index`."
+            )
+            raise typer.Exit(1) from exc
+        graph_path = Path(entry.unified_graph_path or entry.graph_path)
+        if not graph_path.exists():
+            console.print(
+                f"[red]Graph file missing:[/red] {graph_path}\n"
+                "Re-run `pretensor index` before exporting."
+            )
+            raise typer.Exit(1)
+        try:
+            store = KuzuStore(graph_path)
+        except Exception as exc:
+            console.print(
+                f"[red]Cannot open graph file {graph_path}:[/red] {exc}\n"
+                "The file may be corrupt. Re-run `pretensor index`."
+            )
+            raise typer.Exit(1) from exc
+        try:
+            payload = export_graph_payload(
+                store,
+                connection_name=entry.connection_name,
+                database_name=entry.database,
+                graph_path=graph_path,
+            )
+        finally:
+            store.close()
+        try:
+            output.parent.mkdir(parents=True, exist_ok=True)
+            output.write_text(
+                json.dumps(payload, indent=2, ensure_ascii=False) + "\n",
+                encoding="utf-8",
+            )
+        except PermissionError as exc:
+            console.print(f"[red]Cannot write export file {output}:[/red] {exc}")
+            raise typer.Exit(1) from exc
+        except OSError as exc:
+            console.print(f"[red]Cannot write export file {output}:[/red] {exc}")
+            raise typer.Exit(1) from exc
+        stats = payload["stats"]
+        console.print(f"[green]Graph exported:[/green] {output}")
+        console.print(
+            "[dim]"
+            f"node_types={stats['node_types']} edge_types={stats['edge_types']} "
+            f"nodes={stats['node_count']} edges={stats['edge_count']}"
+            "[/dim]"
+        )