pretensor 0.1.0a0__py3-none-any.whl

pretensor/__init__.py

@@ -0,0 +1,50 @@
+"""Pretensor Graph — schema knowledge graph backed by Kuzu.
+
+Heavy imports (e.g. :class:`GraphBuilder`) are lazy so ``import pretensor``
+works in tests that only touch the registry.
+"""
+
+# pyright: reportUnsupportedDunderAll=false
+# Names in __all__ are provided via __getattr__; Pyright does not model that.
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = [
+    "GraphBuilder",
+    "GraphEdge",
+    "GraphNode",
+    "GraphRegistry",
+    "KuzuStore",
+    "RelationshipCandidate",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "GraphBuilder":
+        from pretensor.core.builder import GraphBuilder
+
+        return GraphBuilder
+    if name == "GraphRegistry":
+        from pretensor.core.registry import GraphRegistry
+
+        return GraphRegistry
+    if name == "KuzuStore":
+        from pretensor.core.store import KuzuStore
+
+        return KuzuStore
+    if name == "GraphNode":
+        from pretensor.graph_models.node import GraphNode
+
+        return GraphNode
+    if name == "GraphEdge":
+        from pretensor.graph_models.edge import GraphEdge
+
+        return GraphEdge
+    if name == "RelationshipCandidate":
+        from pretensor.graph_models.relationship import RelationshipCandidate
+
+        return RelationshipCandidate
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)

pretensor/cli/__init__.py

@@ -0,0 +1 @@
+"""Command-line interface for Pretensor Graph."""

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,122 @@
+"""``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,
+    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)."
+        )
+        console.print(
+            f"Run: pretensor index {dsn!r} --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,120 @@
+"""``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_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,
+        ),
+        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_path_option(
+            ctx,
+            param_name="state_dir",
+            cli_value=state_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]"
+        )