pretensor 0.1.0a0__py3-none-any.whl

pretensor/visibility/config.py

@@ -0,0 +1,126 @@
+"""YAML visibility rules: hidden schemas/tables/columns and named profiles."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field
+from ruamel.yaml import YAML
+
+__all__ = [
+    "VisibilityProfileRule",
+    "VisibilityConfig",
+    "load_visibility_config",
+    "default_visibility_path",
+    "merge_profile_into_base",
+]
+
+
+def default_visibility_path(graph_dir: Path) -> Path:
+    """Default location for visibility rules under the graph state directory."""
+    return graph_dir / "visibility.yml"
+
+
+class VisibilityProfileRule(BaseModel):
+    """Overlay merged onto base visibility when ``--profile`` is set."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    hidden_schemas: list[str] = Field(default_factory=list)
+    hidden_tables: list[str] = Field(default_factory=list)
+    hidden_columns: list[str] = Field(default_factory=list)
+    hidden_table_types: list[str] = Field(default_factory=list)
+    allowed_schemas: list[str] = Field(default_factory=list)
+    allowed_tables: list[str] = Field(default_factory=list)
+
+
+class VisibilityConfig(BaseModel):
+    """Parsed ``visibility.yml`` (base section)."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    hidden_schemas: list[str] = Field(default_factory=list)
+    hidden_tables: list[str] = Field(default_factory=list)
+    hidden_columns: list[str] = Field(default_factory=list)
+    hidden_table_types: list[str] = Field(default_factory=list)
+    allowed_schemas: list[str] = Field(default_factory=list)
+    allowed_tables: list[str] = Field(default_factory=list)
+    profiles: dict[str, VisibilityProfileRule] = Field(default_factory=dict)
+
+
+def load_visibility_config(path: Path) -> VisibilityConfig:
+    """Load and validate visibility YAML from disk.
+
+    Args:
+        path: File to read (typically ``.pretensor/visibility.yml``).
+
+    Returns:
+        Parsed config. Missing file yields an empty config.
+
+    Raises:
+        ValueError: If YAML is present but invalid.
+    """
+    if not path.exists():
+        return VisibilityConfig()
+    raw_text = path.read_text(encoding="utf-8")
+    if not raw_text.strip():
+        return VisibilityConfig()
+    yaml = YAML(typ="safe")
+    try:
+        data = yaml.load(raw_text)
+    except Exception as exc:
+        raise ValueError(f"Invalid visibility YAML: {path}: {exc}") from exc
+    if data is None:
+        return VisibilityConfig()
+    if not isinstance(data, dict):
+        raise ValueError(f"visibility.yml must be a mapping at root: {path}")
+    profiles_raw = data.get("profiles")
+    base_data = {k: v for k, v in data.items() if k != "profiles"}
+    cfg = VisibilityConfig.model_validate(base_data)
+    if profiles_raw is None:
+        return cfg
+    if not isinstance(profiles_raw, dict):
+        raise ValueError(f"`profiles` must be a mapping in {path}")
+    profiles: dict[str, VisibilityProfileRule] = {}
+    for name, body in profiles_raw.items():
+        if not isinstance(name, str) or not name.strip():
+            continue
+        if body is None:
+            profiles[name.strip()] = VisibilityProfileRule()
+        elif isinstance(body, dict):
+            profiles[name.strip()] = VisibilityProfileRule.model_validate(body)
+        else:
+            raise ValueError(f"Profile {name!r} must be a mapping in {path}")
+    return cfg.model_copy(update={"profiles": profiles})
+
+
+def merge_profile_into_base(
+    base: VisibilityConfig,
+    profile_name: str | None,
+) -> VisibilityConfig:
+    """Return a new config: base lists + optional profile overlay."""
+    if not profile_name or not profile_name.strip():
+        return base
+    key = profile_name.strip()
+    prof = base.profiles.get(key)
+    if prof is None:
+        raise ValueError(f"Unknown visibility profile: {key!r}")
+    merged_allowed = (
+        [*base.allowed_schemas, *prof.allowed_schemas]
+        if prof.allowed_schemas
+        else list(base.allowed_schemas)
+    )
+    merged_allowed_tables = (
+        [*base.allowed_tables, *prof.allowed_tables]
+        if prof.allowed_tables
+        else list(base.allowed_tables)
+    )
+    return VisibilityConfig(
+        hidden_schemas=[*base.hidden_schemas, *prof.hidden_schemas],
+        hidden_tables=[*base.hidden_tables, *prof.hidden_tables],
+        hidden_columns=[*base.hidden_columns, *prof.hidden_columns],
+        hidden_table_types=[*base.hidden_table_types, *prof.hidden_table_types],
+        allowed_schemas=merged_allowed,
+        allowed_tables=merged_allowed_tables,
+        profiles=base.profiles,
+    )

pretensor/visibility/filter.py

@@ -0,0 +1,143 @@
+"""Runtime visibility checks for tables and columns (fnmatch globs)."""
+
+from __future__ import annotations
+
+import fnmatch
+from dataclasses import dataclass
+
+from pretensor.visibility.config import VisibilityConfig
+
+__all__ = ["VisibilityFilter"]
+
+
+@dataclass(frozen=True, slots=True)
+class VisibilityFilter:
+    """Evaluates table/column visibility for one merged :class:`VisibilityConfig`."""
+
+    _config: VisibilityConfig
+
+    @classmethod
+    def from_config(cls, config: VisibilityConfig) -> VisibilityFilter:
+        return cls(_config=config)
+
+    def is_table_visible(
+        self,
+        connection_name: str,
+        schema_name: str,
+        table_name: str,
+        *,
+        table_type: str | None = None,
+    ) -> bool:
+        """True if the physical table may appear in MCP responses and traversal."""
+        if table_type and self._table_type_hidden(table_type):
+            return False
+        if self._schema_denied(schema_name):
+            return False
+        if self._allowed_schemas_active() and not self._schema_allowed(schema_name):
+            return False
+        if self._allowed_tables_active() and not self._table_allowed(
+            connection_name, schema_name, table_name
+        ):
+            return False
+        return not self._table_pattern_matches(connection_name, schema_name, table_name)
+
+    def visible_columns(
+        self,
+        connection_name: str,
+        schema_name: str,
+        table_name: str,
+        column_names: list[str],
+    ) -> list[str]:
+        """Return column names that are not hidden for this table."""
+        return [
+            c
+            for c in column_names
+            if not self._column_hidden(schema_name, table_name, c)
+        ]
+
+    def is_schema_table_node_id_visible(self, node_id: str) -> bool:
+        """Interpret ``connection::schema::table`` node ids (SchemaTable)."""
+        parts = node_id.split("::")
+        if len(parts) != 3:
+            return True
+        connection_name, schema_name, table_name = parts
+        return self.is_table_visible(connection_name, schema_name, table_name)
+
+    def _allowed_schemas_active(self) -> bool:
+        return bool(self._config.allowed_schemas)
+
+    def _schema_allowed(self, schema_name: str) -> bool:
+        for pat in self._config.allowed_schemas:
+            if fnmatch.fnmatchcase(schema_name, pat):
+                return True
+        return False
+
+    def _allowed_tables_active(self) -> bool:
+        return bool(self._config.allowed_tables)
+
+    def _table_allowed(
+        self, connection_name: str, schema_name: str, table_name: str
+    ) -> bool:
+        qualified = f"{schema_name}.{table_name}"
+        conn_qualified = f"{connection_name}::{schema_name}.{table_name}"
+        for pat in self._config.allowed_tables:
+            if not pat.strip():
+                continue
+            if "::" in pat:
+                if fnmatch.fnmatchcase(conn_qualified, pat):
+                    return True
+                continue
+            if "." in pat:
+                if fnmatch.fnmatchcase(qualified, pat):
+                    return True
+                continue
+            if fnmatch.fnmatchcase(table_name, pat):
+                return True
+        return False
+
+    def _schema_denied(self, schema_name: str) -> bool:
+        for pat in self._config.hidden_schemas:
+            if fnmatch.fnmatchcase(schema_name, pat):
+                return True
+        return False
+
+    def _table_pattern_matches(
+        self, connection_name: str, schema_name: str, table_name: str
+    ) -> bool:
+        qualified = f"{schema_name}.{table_name}"
+        conn_qualified = f"{connection_name}::{schema_name}.{table_name}"
+        for pat in self._config.hidden_tables:
+            if not pat.strip():
+                continue
+            if "::" in pat:
+                if fnmatch.fnmatchcase(conn_qualified, pat):
+                    return True
+                continue
+            if "." in pat:
+                if fnmatch.fnmatchcase(qualified, pat):
+                    return True
+                continue
+            if fnmatch.fnmatchcase(table_name, pat):
+                return True
+        return False
+
+    def _table_type_hidden(self, table_type: str) -> bool:
+        tt = table_type.strip().lower()
+        return tt in {t.strip().lower() for t in self._config.hidden_table_types}
+
+    def _column_hidden(self, schema_name: str, table_name: str, column_name: str) -> bool:
+        qualified_3 = f"{schema_name}.{table_name}.{column_name}"
+        qualified_2 = f"{table_name}.{column_name}"
+        for pat in self._config.hidden_columns:
+            if not pat.strip():
+                continue
+            if pat.count(".") >= 2:
+                if fnmatch.fnmatchcase(qualified_3, pat):
+                    return True
+            elif "." in pat:
+                if fnmatch.fnmatchcase(qualified_2, pat):
+                    return True
+            else:
+                if fnmatch.fnmatchcase(column_name, pat):
+                    return True
+        return False

pretensor/visibility/kuzu_helpers.py

@@ -0,0 +1,32 @@
+"""Kuzu helpers for visibility-scoped node sets."""
+
+from __future__ import annotations
+
+from pretensor.core.store import KuzuStore
+from pretensor.visibility.filter import VisibilityFilter
+
+__all__ = ["visible_schema_table_node_ids"]
+
+
+def visible_schema_table_node_ids(
+    store: KuzuStore,
+    connection_name: str,
+    visibility_filter: VisibilityFilter | None,
+) -> set[str] | None:
+    """Return visible ``SchemaTable`` node ids, or None when filtering is disabled."""
+    if visibility_filter is None:
+        return None
+    rows = store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {connection_name: $cn})
+        RETURN t.node_id, t.schema_name, t.table_name
+        """,
+        {"cn": connection_name},
+    )
+    out: set[str] = set()
+    for nid, sn, tn in rows:
+        if visibility_filter.is_table_visible(
+            connection_name, str(sn or ""), str(tn or "")
+        ):
+            out.add(str(nid))
+    return out

pretensor/visibility/runtime.py

@@ -0,0 +1,36 @@
+"""Load merged visibility rules for a graph state directory."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pretensor.visibility.config import (
+    default_visibility_path,
+    load_visibility_config,
+    merge_profile_into_base,
+)
+from pretensor.visibility.filter import VisibilityFilter
+
+__all__ = ["load_visibility_filter_for_graph_dir"]
+
+
+def load_visibility_filter_for_graph_dir(
+    graph_dir: Path,
+    *,
+    visibility_path: Path | None = None,
+    profile: str | None = None,
+) -> VisibilityFilter:
+    """Load ``visibility.yml`` and return a :class:`VisibilityFilter`.
+
+    Args:
+        graph_dir: State directory (``--graph-dir`` / ``.pretensor``).
+        visibility_path: Optional explicit file; defaults to ``graph_dir/visibility.yml``.
+        profile: Optional named profile merged onto the base section.
+
+    Returns:
+        Filter instance (allows all tables when no file or empty rules).
+    """
+    path = visibility_path if visibility_path is not None else default_visibility_path(graph_dir)
+    base = load_visibility_config(path)
+    merged = merge_profile_into_base(base, profile)
+    return VisibilityFilter.from_config(merged)

pretensor/visibility/sync_grants.py

@@ -0,0 +1,188 @@
+"""Merge DB SELECT grants into ``visibility.yml`` profile ``allowed_tables``."""
+
+from __future__ import annotations
+
+import copy
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Mapping
+
+from ruamel.yaml import YAML
+
+from pretensor.connectors.base import BaseConnector, TableGrant
+
+__all__ = [
+    "allowed_table_patterns_for_grants",
+    "merge_grant_profiles_into_visibility_doc",
+    "write_visibility_yaml_document",
+]
+
+
+def allowed_table_patterns_for_grants(
+    grants: list[TableGrant],
+    connection_name: str,
+) -> dict[str, list[str]]:
+    """Group grants by grantee; each value is sorted ``connection::schema.table`` patterns."""
+    by_grantee: dict[str, set[str]] = defaultdict(set)
+    for row in grants:
+        g = row.grantee.strip()
+        if not g:
+            continue
+        schema = row.schema_name.strip()
+        table = row.table_name.strip()
+        if not schema or not table:
+            continue
+        by_grantee[g].add(f"{connection_name}::{schema}.{table}")
+    return {k: sorted(v) for k, v in sorted(by_grantee.items(), key=lambda kv: kv[0])}
+
+
+def _is_grant_only_profile(body: Mapping[str, Any]) -> bool:
+    """True if the profile body only carries ``allowed_tables`` (grant-sync shape)."""
+    keys = frozenset(body.keys())
+    return keys == frozenset({"allowed_tables"})
+
+
+def merge_grant_profiles_into_visibility_doc(
+    *,
+    existing_root: Mapping[str, Any] | None,
+    grantee_to_tables: Mapping[str, list[str]],
+    roles_filter: frozenset[str] | None,
+) -> dict[str, Any]:
+    """Return a YAML-root mapping: base keys preserved, ``profiles`` merged for grant sync.
+
+    Profiles not targeted by this sync pass are copied verbatim. For each targeted
+    grantee, ``allowed_tables`` is set to the computed list; other keys on that profile
+    are preserved. Grantees with no tables omit ``allowed_tables`` unless the profile
+    has other keys (then ``allowed_tables`` is removed to avoid an empty allowlist).
+
+    Args:
+        existing_root: Parsed root mapping from the current file, or ``None``.
+        grantee_to_tables: Grantee name -> ``allowed_tables`` patterns (sorted).
+        roles_filter: When set, only these grantees receive grant-driven updates.
+
+    Returns:
+        New root dict suitable for YAML serialization.
+    """
+    root: dict[str, Any] = (
+        copy.deepcopy(dict(existing_root)) if existing_root else {}
+    )
+    profiles_any = root.get("profiles")
+    profiles: dict[str, Any] = (
+        copy.deepcopy(dict(profiles_any)) if isinstance(profiles_any, dict) else {}
+    )
+
+    grant_keys = frozenset(grantee_to_tables.keys())
+    if roles_filter is None:
+        stale = [
+            name
+            for name, body in list(profiles.items())
+            if name not in grant_keys
+            and isinstance(body, dict)
+            and _is_grant_only_profile(body)
+        ]
+        for name in stale:
+            profiles.pop(name, None)
+
+    if roles_filter is not None:
+        targets = sorted(roles_filter)
+    else:
+        targets = sorted(grant_keys)
+
+    for grantee in targets:
+        tables = list(grantee_to_tables.get(grantee, []))
+        existing_body = profiles.get(grantee)
+        body: dict[str, Any] = (
+            copy.deepcopy(dict(existing_body))
+            if isinstance(existing_body, dict)
+            else {}
+        )
+        other_keys = {k for k in body if k != "allowed_tables"}
+        if tables:
+            body["allowed_tables"] = tables
+        else:
+            body.pop("allowed_tables", None)
+            if not other_keys:
+                profiles.pop(grantee, None)
+                continue
+        profiles[grantee] = body
+
+    if profiles:
+        root["profiles"] = profiles
+    else:
+        root.pop("profiles", None)
+
+    return root
+
+
+def write_visibility_yaml_document(path: Path, document: Mapping[str, Any]) -> None:
+    """Atomically write a mapping to ``path`` using block-style YAML.
+
+    Writes to a sibling ``.tmp`` file first and renames on success, so an
+    interrupted run (Ctrl-C, crash, disk full) never leaves ``visibility.yml``
+    truncated — the file is either the previous content or the new content,
+    never a partial write. ``os.replace`` is atomic on POSIX and Windows.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    yaml = YAML()
+    yaml.default_flow_style = False
+    yaml.indent(mapping=2, sequence=4, offset=2)
+    tmp_path = path.with_name(path.name + ".tmp")
+    try:
+        with tmp_path.open("w", encoding="utf-8") as handle:
+            yaml.dump(dict(document), handle)
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(tmp_path, path)
+    except BaseException:
+        # Clean up partial temp file so a retry starts from a clean slate.
+        try:
+            tmp_path.unlink()
+        except FileNotFoundError:
+            pass
+        raise
+
+
+def load_visibility_yaml_root(path: Path) -> dict[str, Any] | None:
+    """Load the raw YAML root mapping, or ``None`` if missing/empty."""
+    if not path.exists():
+        return None
+    text = path.read_text(encoding="utf-8")
+    if not text.strip():
+        return None
+    yaml = YAML(typ="safe")
+    data = yaml.load(text)
+    if data is None:
+        return None
+    if not isinstance(data, dict):
+        msg = f"visibility.yml must be a mapping at root: {path}"
+        raise ValueError(msg)
+    return data
+
+
+def run_sync_grants(
+    connector: BaseConnector,
+    *,
+    output_path: Path,
+    connection_name: str,
+    roles_filter: frozenset[str] | None,
+) -> int:
+    """Read grants from ``connector``, merge into ``output_path``, return grantee count.
+
+    The connector must already be connected. The file on disk is replaced atomically
+    with merged content; base visibility keys and non-target profiles are preserved.
+
+    Returns the number of distinct grantees for which ``allowed_tables`` was written
+    (0 if the connector returned no grants). Callers that need the parsed
+    :class:`VisibilityConfig` should call :func:`load_visibility_config` afterward.
+    """
+    grants = connector.get_table_grants()
+    grantee_to_tables = allowed_table_patterns_for_grants(grants, connection_name)
+    existing_root = load_visibility_yaml_root(output_path)
+    merged = merge_grant_profiles_into_visibility_doc(
+        existing_root=existing_root,
+        grantee_to_tables=grantee_to_tables,
+        roles_filter=roles_filter,
+    )
+    write_visibility_yaml_document(output_path, merged)
+    return len(grantee_to_tables)