pretensor 0.1.0__py3-none-any.whl

pretensor/cli/config_file.py

@@ -0,0 +1,442 @@
+"""CLI config-file loading and option precedence helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import typer
+from click.core import ParameterSource
+from ruamel.yaml import YAML
+
+from pretensor.config import GraphConfig
+from pretensor.errors import PretensorError
+
+__all__ = [
+    "DEFAULT_CONFIG_PATH",
+    "CliConfigError",
+    "ConnectionDefaults",
+    "LlmDefaults",
+    "PretensorCliConfig",
+    "SourceConfig",
+    "VisibilityDefaults",
+    "get_cli_config",
+    "load_cli_config",
+    "resolve_aliased_path_option",
+    "resolve_optional_path_option",
+    "resolve_optional_str_option",
+    "resolve_path_option",
+]
+
+DEFAULT_CONFIG_PATH = Path(".pretensor") / "config.yaml"
+
+
+class CliConfigError(PretensorError, ValueError):
+    """Raised when the CLI config file is unreadable or invalid."""
+
+
+@dataclass(frozen=True, slots=True)
+class LlmDefaults:
+    """LLM defaults accepted by config for compatibility with external workflows."""
+
+    model: str | None = None
+    budget_usd: float | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class VisibilityDefaults:
+    """Visibility defaults for ``index`` and ``serve``."""
+
+    path: Path | None = None
+    profile: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ConnectionDefaults:
+    """Default connection settings for DSN-driven commands."""
+
+    name: str | None = None
+    dialect: str | None = None
+
+
+_SOURCE_ALLOWED_KEYS = frozenset(
+    {
+        "dialect",
+        "host",
+        "port",
+        "user",
+        "password",
+        "database",
+        # Snowflake
+        "account",
+        "schema",
+        "warehouse",
+        "role",
+        "private_key_path",
+        "private_key_passphrase",
+        # BigQuery
+        "project",
+        "dataset",
+        "location",
+    }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class SourceConfig:
+    """Declarative database source definition from ``sources:`` config block."""
+
+    dialect: str
+    host: str | None = None
+    port: int | None = None
+    user: str | None = None
+    password: str | None = None
+    database: str | None = None
+    # Snowflake
+    account: str | None = None
+    schema: str | None = None
+    warehouse: str | None = None
+    role: str | None = None
+    private_key_path: str | None = None
+    private_key_passphrase: str | None = None
+    # BigQuery
+    project: str | None = None
+    dataset: str | None = None
+    location: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class PretensorCliConfig:
+    """Normalized runtime config loaded from YAML (or empty defaults)."""
+
+    source_path: Path | None = None
+    state_dir: Path | None = None
+    graph: GraphConfig = GraphConfig()
+    llm: LlmDefaults = LlmDefaults()
+    visibility: VisibilityDefaults = VisibilityDefaults()
+    connection_defaults: ConnectionDefaults = ConnectionDefaults()
+    sources: dict[str, SourceConfig] = field(default_factory=dict)
+
+    @property
+    def loaded(self) -> bool:
+        """Whether this config came from a file on disk."""
+        return self.source_path is not None
+
+
+def _as_mapping(raw: Any, *, field: str) -> dict[str, Any]:
+    if raw is None:
+        return {}
+    if not isinstance(raw, dict):
+        raise CliConfigError(f"`{field}` must be a mapping in config YAML")
+    return raw
+
+
+def _resolved_path(raw: Any, *, field: str, base_dir: Path) -> Path:
+    if not isinstance(raw, str) or not raw.strip():
+        raise CliConfigError(f"`{field}` must be a non-empty string path")
+    value = Path(raw.strip()).expanduser()
+    if not value.is_absolute():
+        value = (base_dir / value).resolve()
+    else:
+        value = value.resolve()
+    return value
+
+
+def _resolved_optional_path(raw: Any, *, field: str, base_dir: Path) -> Path | None:
+    if raw is None:
+        return None
+    return _resolved_path(raw, field=field, base_dir=base_dir)
+
+
+def _normalized_optional_str(raw: Any, *, field: str) -> str | None:
+    if raw is None:
+        return None
+    if not isinstance(raw, str):
+        raise CliConfigError(f"`{field}` must be a string when provided")
+    value = raw.strip()
+    return value or None
+
+
+def _normalized_optional_float(raw: Any, *, field: str) -> float | None:
+    if raw is None:
+        return None
+    if isinstance(raw, bool):
+        raise CliConfigError(f"`{field}` must be a number when provided")
+    if isinstance(raw, (int, float)):
+        return float(raw)
+    raise CliConfigError(f"`{field}` must be a number when provided")
+
+
+def _graph_config_from_mapping(raw: Any) -> GraphConfig:
+    data = _as_mapping(raw, field="graph")
+    allowed = {
+        "stale_index_warning_days",
+        "clustering_resolution_override",
+        "join_path_max_depth",
+        "min_cluster_size_merge",
+        "collapse_shadow_aliases",
+        "lineage_in_max_for_alias",
+    }
+    unknown = sorted(set(data) - allowed)
+    if unknown:
+        joined = ", ".join(unknown)
+        raise CliConfigError(f"Unknown `graph` config key(s): {joined}")
+    try:
+        return GraphConfig(**data)
+    except TypeError as exc:
+        raise CliConfigError(f"Invalid `graph` config: {exc}") from exc
+
+
+def _source_config_from_mapping(
+    name: str,
+    raw: Any,
+) -> SourceConfig:
+    """Validate and build a :class:`SourceConfig` from a raw YAML mapping."""
+    if not isinstance(raw, dict):
+        raise CliConfigError(f"Source `{name}` must be a mapping")
+    data: dict[str, Any] = dict(raw)
+    unknown = sorted(set(data) - _SOURCE_ALLOWED_KEYS)
+    if unknown:
+        joined = ", ".join(unknown)
+        raise CliConfigError(f"Unknown key(s) in source `{name}`: {joined}")
+    dialect = data.get("dialect")
+    if not isinstance(dialect, str) or not dialect.strip():
+        raise CliConfigError(f"Source `{name}` requires a `dialect` string")
+    port = data.get("port")
+    if port is not None:
+        if isinstance(port, bool) or not isinstance(port, int):
+            raise CliConfigError(f"Source `{name}`: `port` must be an integer")
+    str_fields = {
+        "host",
+        "user",
+        "password",
+        "database",
+        "account",
+        "schema",
+        "warehouse",
+        "role",
+        "private_key_path",
+        "private_key_passphrase",
+        "project",
+        "dataset",
+        "location",
+    }
+    kwargs: dict[str, Any] = {"dialect": dialect.strip().lower(), "port": port}
+    for key in str_fields:
+        val = data.get(key)
+        if val is None:
+            kwargs[key] = None
+        elif not isinstance(val, str):
+            raise CliConfigError(f"Source `{name}`: `{key}` must be a string")
+        else:
+            kwargs[key] = val.strip() or None
+    return SourceConfig(**kwargs)
+
+
+def _parse_sources(
+    raw: Any,
+    secrets_raw: dict[str, Any] | None,
+) -> dict[str, SourceConfig]:
+    """Parse the ``sources:`` block and merge optional secrets overlay."""
+    mapping = _as_mapping(raw, field="sources")
+    if not mapping:
+        return {}
+    sources: dict[str, SourceConfig] = {}
+    for name, block in mapping.items():
+        if not isinstance(name, str) or not name.strip():
+            raise CliConfigError("Source names must be non-empty strings")
+        merged = dict(block) if isinstance(block, dict) else block
+        if secrets_raw and name in secrets_raw:
+            secret_block = secrets_raw[name]
+            if not isinstance(secret_block, dict):
+                raise CliConfigError(
+                    f"Secrets entry for source `{name}` must be a mapping"
+                )
+            if isinstance(merged, dict):
+                merged = {**merged, **secret_block}
+        sources[name.strip()] = _source_config_from_mapping(name.strip(), merged)
+    return sources
+
+
+def _load_secrets_file(base_dir: Path) -> dict[str, Any] | None:
+    """Load the optional ``sources.secrets.yaml`` sibling to the config file."""
+    secrets_path = base_dir / "sources.secrets.yaml"
+    if not secrets_path.exists():
+        return None
+    return _parse_yaml_file(secrets_path)
+
+
+def _parse_yaml_file(path: Path) -> dict[str, Any]:
+    yaml = YAML(typ="safe")
+    try:
+        loaded = yaml.load(path.read_text(encoding="utf-8"))
+    except OSError as exc:
+        raise CliConfigError(f"Cannot read config file {path}: {exc}") from exc
+    except Exception as exc:
+        raise CliConfigError(f"Invalid YAML in config file {path}: {exc}") from exc
+    if loaded is None:
+        return {}
+    if not isinstance(loaded, dict):
+        raise CliConfigError("Config file root must be a mapping")
+    return dict(loaded)
+
+
+def load_cli_config(config_path: Path | None) -> PretensorCliConfig:
+    """Load and normalize CLI defaults from YAML."""
+    target = (config_path or DEFAULT_CONFIG_PATH).expanduser()
+    explicit = config_path is not None
+    if not target.exists():
+        if explicit:
+            raise CliConfigError(f"Config file not found: {target}")
+        return PretensorCliConfig()
+
+    source = target.resolve()
+    base_dir = source.parent
+    raw = _parse_yaml_file(source)
+
+    state_dir = _resolved_optional_path(
+        raw.get("state_dir"),
+        field="state_dir",
+        base_dir=base_dir,
+    )
+    graph = _graph_config_from_mapping(raw.get("graph"))
+
+    llm_raw = _as_mapping(raw.get("llm"), field="llm")
+    llm = LlmDefaults(
+        model=_normalized_optional_str(llm_raw.get("model"), field="llm.model"),
+        budget_usd=_normalized_optional_float(
+            llm_raw.get("budget_usd"),
+            field="llm.budget_usd",
+        ),
+    )
+
+    visibility_raw = _as_mapping(raw.get("visibility"), field="visibility")
+    visibility = VisibilityDefaults(
+        path=_resolved_optional_path(
+            visibility_raw.get("path"),
+            field="visibility.path",
+            base_dir=base_dir,
+        ),
+        profile=_normalized_optional_str(
+            visibility_raw.get("profile"),
+            field="visibility.profile",
+        ),
+    )
+
+    conn_raw = _as_mapping(raw.get("connection_defaults"), field="connection_defaults")
+    connection_defaults = ConnectionDefaults(
+        name=_normalized_optional_str(
+            conn_raw.get("name"),
+            field="connection_defaults.name",
+        ),
+        dialect=_normalized_optional_str(
+            conn_raw.get("dialect"),
+            field="connection_defaults.dialect",
+        ),
+    )
+
+    secrets_raw = _load_secrets_file(base_dir)
+    sources = _parse_sources(raw.get("sources"), secrets_raw)
+
+    return PretensorCliConfig(
+        source_path=source,
+        state_dir=state_dir,
+        graph=graph,
+        llm=llm,
+        visibility=visibility,
+        connection_defaults=connection_defaults,
+        sources=sources,
+    )
+
+
+def get_cli_config(ctx: typer.Context | None) -> PretensorCliConfig:
+    """Return config stored in Typer context, or an empty config."""
+    if ctx is None:
+        return PretensorCliConfig()
+    obj = ctx.obj
+    if isinstance(obj, PretensorCliConfig):
+        return obj
+    if isinstance(obj, dict):
+        cfg = obj.get("config")
+        if isinstance(cfg, PretensorCliConfig):
+            return cfg
+    return PretensorCliConfig()
+
+
+def _is_user_override(ctx: typer.Context | None, param_name: str) -> bool:
+    if ctx is None:
+        return False
+    try:
+        source = ctx.get_parameter_source(param_name)
+    except Exception:
+        return False
+    if source is None:
+        return False
+    return source is not ParameterSource.DEFAULT
+
+
+def resolve_path_option(
+    ctx: typer.Context | None,
+    *,
+    param_name: str,
+    cli_value: Path,
+    config_value: Path | None,
+) -> Path:
+    """Return effective path option with ``CLI > config > default`` precedence."""
+    if _is_user_override(ctx, param_name):
+        return cli_value
+    return config_value or cli_value
+
+
+def resolve_aliased_path_option(
+    ctx: typer.Context | None,
+    *,
+    primary_param: str,
+    primary_value: Path,
+    alias_param: str,
+    alias_value: Path | None,
+    config_value: Path | None,
+) -> Path:
+    """Resolve a path option that has a hidden, deprecated alias flag.
+
+    Precedence: explicit ``primary_param`` > explicit ``alias_param`` > config
+    > ``primary_value`` default. If both flags are passed explicitly, the
+    primary (canonical) one wins.
+    """
+    if _is_user_override(ctx, primary_param):
+        return primary_value
+    if _is_user_override(ctx, alias_param) and alias_value is not None:
+        return alias_value
+    return config_value if config_value is not None else primary_value
+
+
+def resolve_optional_path_option(
+    ctx: typer.Context | None,
+    *,
+    param_name: str,
+    cli_value: Path | None,
+    config_value: Path | None,
+) -> Path | None:
+    """Return effective optional path with ``CLI > config > default`` precedence."""
+    if _is_user_override(ctx, param_name):
+        return cli_value
+    return config_value if config_value is not None else cli_value
+
+
+def resolve_optional_str_option(
+    ctx: typer.Context | None,
+    *,
+    param_name: str,
+    cli_value: str | None,
+    config_value: str | None,
+) -> str | None:
+    """Return effective optional string with ``CLI > config > default`` precedence."""
+    if _is_user_override(ctx, param_name):
+        return cli_value
+    return config_value if config_value is not None else cli_value
+
+
+def resolve_graph_config(ctx: typer.Context | None) -> GraphConfig:
+    """Return graph config from context, falling back to OSS defaults."""
+    return get_cli_config(ctx).graph

pretensor/cli/constants.py

@@ -0,0 +1,10 @@
+"""Shared CLI literals (state dir, registry filenames, graph layout)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+DEFAULT_STATE_DIR = Path(".pretensor")
+GRAPHS_SUBDIR = "graphs"
+REGISTRY_FILENAME = "registry.json"
+UNIFIED_GRAPH_BASENAME = "unified.kuzu"

pretensor/cli/dbt_enrichment.py

@@ -0,0 +1,96 @@
+"""Shared dbt manifest / sources path handling for CLI commands."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+from pretensor.core.store import KuzuStore
+from pretensor.enrichment.dbt.manifest import DbtManifest, DbtManifestError
+from pretensor.enrichment.dbt.pipeline import run_dbt_enrichment_from_manifest
+
+
+def resolve_dbt_sources_path(
+    manifest_path: Path, explicit: Path | None
+) -> Path | None:
+    """Return path to ``sources.json`` for freshness signals.
+
+    If ``explicit`` is set, it must exist as a file. If omitted, use
+    ``<manifest_dir>/sources.json`` when that file exists.
+
+    Args:
+        manifest_path: Resolved manifest path (must exist when called for defaulting).
+        explicit: User-provided ``--dbt-sources`` or ``None``.
+
+    Returns:
+        Path to read, or ``None`` when no sources file should be used.
+
+    Raises:
+        FileNotFoundError: When ``explicit`` is set but is not a file.
+    """
+    if explicit is not None:
+        if not explicit.is_file():
+            msg = f"dbt sources file not found or not a file: {explicit}"
+            raise FileNotFoundError(msg)
+        return explicit
+    default = manifest_path.parent / "sources.json"
+    return default if default.is_file() else None
+
+
+def validate_dbt_manifest_path(path: Path) -> None:
+    """Ensure ``--dbt-manifest`` points at an existing file before heavy work."""
+    if not path.is_file():
+        msg = f"dbt manifest not found or not a file: {path}"
+        raise FileNotFoundError(msg)
+
+
+def preload_dbt_manifest(
+    manifest_path: Path, sources_explicit: Path | None, *, console: Console
+) -> tuple[DbtManifest, Path | None]:
+    """Validate and parse the manifest (and ``sources.json``) up front.
+
+    Running this before any graph mutations means a malformed ``manifest.json``
+    aborts the command *before* the DB is re-introspected and patched, instead
+    of leaving a half-updated graph that then fails during enrichment.
+
+    Raises:
+        typer.Exit: With exit code 1 if either file is missing or the manifest
+            cannot be parsed.
+    """
+    try:
+        validate_dbt_manifest_path(manifest_path)
+        sources = resolve_dbt_sources_path(manifest_path, sources_explicit)
+        manifest = DbtManifest.load(manifest_path)
+    except DbtManifestError as exc:
+        console.print(f"[red]{exc}[/red]")
+        raise typer.Exit(1) from exc
+    except FileNotFoundError as exc:
+        console.print(f"[red]{exc}[/red]")
+        raise typer.Exit(1) from exc
+    return manifest, sources
+
+
+def apply_dbt_enrichment_cli(
+    *,
+    manifest: DbtManifest,
+    sources_path: Path | None,
+    store: KuzuStore,
+    connection_name: str,
+    console: Console,
+) -> None:
+    """Run dbt enrichment on an already-parsed manifest and print a summary line."""
+    summary = run_dbt_enrichment_from_manifest(
+        manifest, sources_path, store, connection_name
+    )
+
+    console.print(
+        "[green]dbt enrichment:[/green] "
+        f"lineage_edges={summary.lineage_edges} "
+        f"tables_enriched={summary.tables_enriched} "
+        f"tags_set={summary.tags_set} "
+        f"exposures_marked={summary.exposures_marked} "
+        f"freshness_rows={summary.freshness_rows_applied} "
+        f"tests_counted={summary.tests_counted}"
+    )

pretensor/cli/main.py

@@ -0,0 +1,109 @@
+"""Typer entrypoint for ``pretensor``."""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+from pretensor.benchmark import register_benchmark_command
+from pretensor.cli.commands.analyze import register_analyze_command
+from pretensor.cli.commands.connections import register_connection_commands
+from pretensor.cli.commands.export import register_export_command
+from pretensor.cli.commands.index import register_index_command
+from pretensor.cli.commands.list import register_list_command
+from pretensor.cli.commands.quickstart import register_quickstart_command
+from pretensor.cli.commands.reindex import register_reindex_command
+from pretensor.cli.commands.semantic import register_semantic_commands
+from pretensor.cli.commands.serve import register_serve_command
+from pretensor.cli.commands.sync_grants import register_sync_grants_command
+from pretensor.cli.commands.validate import register_validate_command
+from pretensor.cli.config_file import CliConfigError, load_cli_config
+from pretensor.cli.plugin import discover_cli_plugins
+from pretensor.observability import LogFormat, LogLevel, configure_logging
+
+app = typer.Typer(
+    no_args_is_help=True,
+    add_completion=False,
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+
+
+def _version_callback(value: bool) -> None:
+    if not value:
+        return
+    try:
+        ver = _pkg_version("pretensor")
+    except PackageNotFoundError:
+        ver = "unknown"
+    typer.echo(f"pretensor {ver}")
+    raise typer.Exit()
+
+
+@app.callback()
+def configure_root_logging(
+    ctx: typer.Context,
+    version: bool = typer.Option(
+        False,
+        "--version",
+        callback=_version_callback,
+        is_eager=True,
+        help="Show the pretensor version and exit.",
+    ),
+    log_level: LogLevel = typer.Option(
+        "warning",
+        "--log-level",
+        help="Logging level: debug, info, warning, error.",
+    ),
+    log_format: LogFormat = typer.Option(
+        "text",
+        "--log-format",
+        help="Log output format: text or json.",
+    ),
+    log_file: Path | None = typer.Option(
+        None,
+        "--log-file",
+        help="Optional file path for duplicate log output.",
+        dir_okay=False,
+        writable=True,
+        resolve_path=True,
+    ),
+    config: Path | None = typer.Option(
+        None,
+        "--config",
+        help="Optional path to pretensor CLI YAML config.",
+        dir_okay=False,
+        file_okay=True,
+        resolve_path=True,
+    ),
+) -> None:
+    """Configure process-wide logging for all CLI subcommands."""
+    configure_logging(level=log_level, log_format=log_format, log_file=log_file)
+    try:
+        cli_config = load_cli_config(config)
+    except CliConfigError as exc:
+        Console().print(f"[red]{exc}[/red]")
+        raise typer.Exit(1) from exc
+    ctx.obj = {"config": cli_config}
+
+
+register_serve_command(app)
+register_sync_grants_command(app)
+console = Console()
+
+register_connection_commands(app)
+
+register_benchmark_command(app)
+register_index_command(app, console=console)
+register_analyze_command(app, console=console)
+register_list_command(app, console=console)
+register_quickstart_command(app, console=console)
+register_reindex_command(app, console=console)
+register_export_command(app, console=console)
+register_semantic_commands(app, console=console)
+register_validate_command(app, console=console)
+
+discover_cli_plugins(app)

pretensor/cli/paths.py

@@ -0,0 +1,43 @@
+"""Path helpers for CLI state directory layout."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pretensor.cli import constants as c
+from pretensor.introspection.models.config import DatabaseType
+from pretensor.introspection.models.dsn import (
+    connection_config_from_url,
+    infer_database_type_from_dsn,
+)
+
+
+def default_connection_name(dsn: str) -> str:
+    """Use the database name from the URL when present, else ``default``."""
+    raw = dsn.strip()
+    try:
+        if infer_database_type_from_dsn(raw) == DatabaseType.BIGQUERY:
+            cfg = connection_config_from_url(raw, "_tmp")
+            logical_db = (cfg.database or "").strip()
+            if "/" in logical_db:
+                return logical_db.split("/", 1)[1]
+            return logical_db or "default"
+    except ValueError:
+        pass
+    cfg = connection_config_from_url(raw, "_tmp")
+    if cfg.database:
+        return cfg.database.replace("/", "_")
+    return "default"
+
+
+def graph_file_for_connection(state_dir: Path, connection_name: str) -> Path:
+    safe = connection_name.replace("/", "_")
+    return state_dir / c.GRAPHS_SUBDIR / f"{safe}.kuzu"
+
+
+def unified_graph_path(state_dir: Path) -> Path:
+    return state_dir / c.GRAPHS_SUBDIR / c.UNIFIED_GRAPH_BASENAME
+
+
+def keystore_path(state_dir: Path) -> Path:
+    return state_dir / "keystore"