dataface 0.1.2__py3-none-any.whl

dataface/agent_api/schema.py

@@ -0,0 +1,658 @@
+"""Typed schema verb — hierarchical data drill-down: source → schema → table → column.
+
+Every response is a single ``SchemaResponse`` envelope wrapped in
+``sources``, regardless of drill level. Leaves are bare scalars;
+response-level provenance lives in the ``_meta`` footer.
+
+Wildcard targeting at every level: ``*`` for everything, ``"a,b"`` for a
+finite list, ``"stg_*"`` for fnmatch globs. The resolver expands
+wildcards once at entry; this verb is just the wire-shape adapter.
+"""
+
+from __future__ import annotations
+
+import re
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from dataface.core.compile.errors import DatafaceError
+from dataface.core.errors import (
+    DF_EXECUTE_COLUMN_NOT_FOUND,
+    DF_EXECUTE_SCHEMA_NOT_FOUND,
+    DF_EXECUTE_TABLE_NOT_FOUND,
+)
+from dataface.core.errors.structured import StructuredError
+from dataface.core.execute.adapters import AdapterRegistry
+from dataface.core.inspect.cache_factory import (
+    build_resolver,
+    project_root_for_registry,
+)
+from dataface.core.inspect.partition_types import TablePartitions
+from dataface.core.inspect.resolver import LayeredSchemaResolver, is_exact_target
+from dataface.core.schema.guidance import attach_empty_table_list_hints
+
+if TYPE_CHECKING:
+    # SuperSchemaSource lives in the private dataface-super-schema package.
+    # Accepted by schema() as an injection point for callers with the private
+    # package installed. At runtime the type is Any; mypy sees the class via
+    # the guard.
+    from dataface_super_schema.inspect.sources.super_schema import SuperSchemaSource
+
+
+class SchemaArgs(BaseModel):
+    """Browse the data hierarchy: source → schema → table → column.
+
+    The deepest non-None arg determines the return tier. No args: lists
+    configured data sources (level 1). source=X: lists schema/namespace
+    names in that source (level 2). source=X, schema=Y: lists tables with
+    lean summaries (level 3 — name, kind, row_count, description, tags
+    when known; full column profiles require drilling). source=X,
+    schema=Y, table=T: deep-profiles that table (level 4). source=X,
+    schema=Y, table=T, column=C: column profile (level 5). Each level
+    accepts ``*``, comma-lists (``"a,b,c"``), and fnmatch globs
+    (``"stg_*"``) for multi-target queries. Always wrapped in ``sources``;
+    response-level provenance lives in the ``_meta`` footer. If
+    it returns zero tables, do not assume the project has no data — some
+    projects query files directly (for example via read_csv(...)) instead
+    of persistent database tables.
+
+    Power-user note: ``--json | jq`` is the escape hatch for ad-hoc
+    cross-cutting queries the curated verbs don't anticipate — e.g. find
+    every column ending in ``_at``, or list tables above a row-count
+    threshold. The JSON shape is a stable named-dict tree; pipe it into
+    ``jq`` and use ``.sources.<src>.schemas.<schema>.tables`` as the root.
+    For keyword and predicate search across descriptions, tags, and roles
+    use ``dft schema -s <keyword>`` instead.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    source: str | None = Field(
+        None,
+        description="Data source name. Omit to list all configured sources.",
+    )
+    schema_name: str | None = Field(
+        None,
+        alias="schema",
+        description=(
+            "Schema/dataset/namespace name. Requires source. Omit to list "
+            "schemas in the source. Accepts wildcards: '*', comma-lists, "
+            "fnmatch globs ('stg_*')."
+        ),
+    )
+    table: str | None = Field(
+        None,
+        description=(
+            "Table name. Requires source and schema. Accepts wildcards: "
+            "'*', comma-lists, fnmatch globs ('stg_*')."
+        ),
+    )
+    column: str | None = Field(
+        None,
+        description=(
+            "Column name. Requires source, schema, and table. Accepts "
+            "wildcards: '*', comma-lists, fnmatch globs ('*_id')."
+        ),
+    )
+    force_refresh: bool = Field(
+        False,
+        description=(
+            "Re-profile a specific table even if cached data exists; valid only "
+            "with source, schema, and table, and without column."
+        ),
+    )
+    table_search: str | None = Field(
+        None,
+        description=(
+            "Regexp filter (re.search, re.IGNORECASE) applied at level 3 "
+            "(source + schema → tables). A table is included when the pattern matches "
+            "its name or any cached column name — the match is unanchored, so 'id' "
+            "matches 'user_id'. Un-profiled tables filter on name only; no extra DB "
+            "queries are issued. Example: 'revenue|arr|mrr'."
+        ),
+    )
+
+
+class ResponseMeta(BaseModel):
+    """Response-level provenance footer.
+
+    Present on successful resolver responses. Validation-error envelopes may
+    omit it. ``sources_consulted`` reports which layers contributed:
+    ``["super_schema"]`` on cache hit, ``["dbt_adapter", "dbt_manifest"]``
+    on cache miss with manifest, ``["dbt_adapter"]`` without manifest, ``[]``
+    for level 1 (registry-only).
+
+    ``retrieved_at`` and ``cache_built_at`` are ISO-8601 strings, not
+    ``datetime`` objects — the MCP transport calls plain ``json.dumps``
+    on the model_dump output and would choke on a datetime instance.
+    The resolver stamps them via ``datetime.isoformat()`` upstream.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    retrieved_at: str
+    sources_consulted: list[Literal["super_schema", "dbt_adapter", "dbt_manifest"]] = []
+    cache_built_at: str | None = None
+
+
+class TableProfile(BaseModel):
+    """Typed level-4 table profile fields added by the schema verb.
+
+    The named fields (``partitions``, ``last_modified``) enumerate the
+    agent-api typed contract; additional fields emitted by the
+    ``LayeredSchemaResolver`` (rows, description, columns, …) flow through
+    via ``extra='allow'`` because the resolver owns the level-4 wire shape
+    and this model is the bridge — not the spec.
+    """
+
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    partitions: TablePartitions | None = None
+    last_modified: datetime | None = None
+
+
+class SchemaResponse(BaseModel):
+    """Single envelope for every drill level.
+
+    ``sources`` is the named-dict tree (``dict[str, ...]`` at every level).
+    ``_meta`` carries response-level provenance. ``success`` and ``errors``
+    preserve the MCP envelope guarantee. ``structured_errors`` carries
+    fully-typed error objects for MCP consumers that can act on code/doc_url.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    success: bool = True
+    sources: dict[str, dict[str, Any]] = Field(
+        default_factory=dict,
+        description="Schema data keyed by source name (tables, columns, etc.).",
+    )
+    supported_query_types: list[str] | None = None
+    meta: ResponseMeta | None = Field(
+        default=None, alias="_meta", description="Pagination and timing metadata."
+    )
+    errors: list[str] = []
+    structured_errors: list[StructuredError] = []
+    # Empty-tables UX: when level 3 returns nothing the verb attaches
+    # discoverability hints. Optional; absent on success-with-results.
+    message: str | None = None
+    warnings: list[str] | None = None
+    hints: list[str] | None = None
+
+
+_LINEAGE_DEPTH_MIN = 1
+_LINEAGE_DEPTH_MAX = 10
+
+
+def schema(
+    source: str | None = None,
+    schema: str | None = None,
+    table: str | None = None,
+    column: str | None = None,
+    *,
+    cache_path: Path | None = None,
+    cache_source: SuperSchemaSource | None = None,
+    force_refresh: bool = False,
+    lineage_depth: int = 1,
+    table_search: str | None = None,
+    surface: Literal["cli", "mcp"] = "mcp",
+    adapter_registry: AdapterRegistry,
+) -> SchemaResponse:
+    """Browse the data hierarchy: source → schema → table → column.
+
+    ``table_search`` is a regexp (``re.IGNORECASE``) applied at level 3 (schema
+    → tables). A table passes if the pattern matches its name OR any cached
+    column name. For un-profiled tables (no cached column list), name-only
+    matching is applied — no extra DB queries are issued.
+
+    ``cache_path``: Explicit path to ``super_schema.json``. When ``None``
+    (default), the path is auto-derived from the registry's project root.
+    Pass an explicit path only when the cache lives outside the project root.
+
+    ``cache_source``: Inject a pre-constructed ``SuperSchemaSource`` directly.
+    Takes priority over ``cache_path`` when both are provided. Use when
+    the caller has already constructed the cache source.
+
+    ``surface``: Controls vocabulary in empty-cache hints. ``"cli"`` emits
+    ``dft inspect`` / ``dft search`` / ``dft query`` verbs; ``"mcp"`` (default)
+    emits MCP tool names (``search_dashboards``, ``execute_query``).
+    """
+    if not (_LINEAGE_DEPTH_MIN <= lineage_depth <= _LINEAGE_DEPTH_MAX):
+        return SchemaResponse(
+            success=False,
+            errors=[
+                f"lineage_depth must be between {_LINEAGE_DEPTH_MIN} and "
+                f"{_LINEAGE_DEPTH_MAX}, got {lineage_depth}"
+            ],
+        )
+
+    invalid = _validate_schema_request(
+        source=source,
+        schema=schema,
+        table=table,
+        column=column,
+        force_refresh=force_refresh,
+        table_search=table_search,
+    )
+    if invalid is not None:
+        return invalid
+
+    compiled_search: re.Pattern[str] | None = None
+    if table_search is not None:
+        try:
+            compiled_search = re.compile(table_search, re.IGNORECASE)
+        except re.error as exc:
+            return SchemaResponse(
+                success=False,
+                errors=[f"table_search is not a valid regexp: {exc}"],
+            )
+
+    if source is None:
+        response = _list_sources_response(adapter_registry)
+        response.supported_query_types = sorted(adapter_registry.supported_types)
+        return response
+
+    try:
+        resolver = build_resolver(
+            adapter_registry, cache_path, cache_source=cache_source
+        )
+        raw = _dispatch_schema_request(
+            resolver=resolver,
+            source=source,
+            schema=schema,
+            table=table,
+            column=column,
+            force_refresh=force_refresh,
+            lineage_depth=lineage_depth,
+        )
+        response = SchemaResponse.model_validate(raw)
+        _validate_table_profiles(
+            response,
+            source=source,
+            schema=schema,
+            table=table,
+            column=column,
+        )
+        not_found = _validate_exact_targets_found(
+            response,
+            resolver=resolver,
+            source=source,
+            schema=schema,
+            table=table,
+            column=column,
+        )
+        if not_found is not None:
+            return not_found
+    except DatafaceError as exc:
+        return SchemaResponse(
+            success=False,
+            errors=[str(exc)],
+            structured_errors=[exc.to_structured()],
+        )
+    except Exception as exc:  # noqa: BLE001 — verb returns errors via response envelope
+        return SchemaResponse(success=False, errors=[str(exc)])
+
+    if schema is not None and table is None:
+        assert source is not None
+        if compiled_search is not None:
+            _filter_tables_by_search(
+                response, source=source, rx=compiled_search, resolver=resolver
+            )
+            # Skip the "no tables found" cache-guidance hints — the empty result is
+            # from the filter, not a missing cache. Attach a filter-specific message.
+            # compiled_search was built from table_search, so table_search is str here.
+            assert isinstance(table_search, str)
+            _attach_filtered_empty_message(
+                response, source=source, pattern=table_search
+            )
+        else:
+            project_root = project_root_for_registry(adapter_registry)
+            attach_empty_table_list_hints(
+                response,
+                adapter_registry,
+                source,
+                uses_direct_file_queries=lambda: _project_uses_direct_file_queries(
+                    project_root
+                ),
+                surface=surface,
+            )
+
+    return response
+
+
+def _validate_exact_targets_found(
+    response: SchemaResponse,
+    resolver: LayeredSchemaResolver,
+    source: str,
+    schema: str | None,
+    table: str | None,
+    column: str | None,
+) -> SchemaResponse | None:
+    """Check whether every exact drill-down target exists in the source universe.
+
+    Only validates tiers where the spec is exact AND all parent tiers are also
+    exact. Returns a ``SchemaResponse(success=False, ...)`` envelope on the
+    first missing tier, or ``None`` when all exact targets were found (or a
+    wildcard appears at some tier).
+
+    Universe queries (``list_schemas``, ``list_tables``) are deferred to the
+    error path — only called when the response tree is missing the target.
+    This avoids extra adapter round-trips on successful requests.
+    """
+    if not is_exact_target(schema):
+        return None
+
+    # Check schema tier using the response tree (fast path).
+    # The resolver omits a schema entry when none of its tables matched the
+    # table spec, so absence can mean either "schema not found" or "schema
+    # exists but no tables matched". We resolve the ambiguity lazily by
+    # querying the universe only when the schema is absent from the response.
+    source_data = response.sources.get(source, {})
+    schema_in_response = schema in source_data.get("schemas", {})
+
+    if not schema_in_response:
+        # Query the real schema universe to distinguish "schema not found"
+        # from "schema found but no tables matched".
+        all_schemas = sorted(
+            resolver.list_schemas(source=source)
+            .get("sources", {})
+            .get(source, {})
+            .get("schemas", {})
+            .keys()
+        )
+        if schema not in all_schemas:
+            exc = DatafaceError.from_code(
+                DF_EXECUTE_SCHEMA_NOT_FOUND,
+                schema=schema,
+                source=source,
+                available=", ".join(all_schemas) if all_schemas else "(none)",
+            )
+            return SchemaResponse(
+                success=False,
+                errors=[str(exc)],
+                structured_errors=[exc.to_structured()],
+            )
+        # Schema exists but was omitted because no tables matched.
+        # Fall through to the table-tier check.
+
+    if not is_exact_target(table):
+        return None
+
+    # is_exact_target(schema) returned True above, so schema is not None.
+    assert schema is not None
+
+    # Check table tier using the response tree.
+    schema_data = source_data.get("schemas", {}).get(schema, {})
+    table_in_response = table in schema_data.get("tables", {})
+
+    if not table_in_response:
+        all_tables = sorted(
+            resolver.list_tables(source=source, schema=schema)
+            .get("sources", {})
+            .get(source, {})
+            .get("schemas", {})
+            .get(schema, {})
+            .get("tables", {})
+            .keys()
+        )
+        if table not in all_tables:
+            exc = DatafaceError.from_code(
+                DF_EXECUTE_TABLE_NOT_FOUND,
+                table=table,
+                schema=schema,
+                source=source,
+                available=", ".join(all_tables) if all_tables else "(none)",
+            )
+            return SchemaResponse(
+                success=False,
+                errors=[str(exc)],
+                structured_errors=[exc.to_structured()],
+            )
+        # Table exists but is absent from the response because the column
+        # filter already dropped it. Fall through to the column-tier check.
+
+    if not is_exact_target(column):
+        return None
+
+    # Check column tier.
+    # When the table is in the response, the column list is reliable (exact
+    # table profile was dispatched). When the table was dropped by the column
+    # filter (table_in_response is False), the column is definitively absent —
+    # we confirmed the table exists above, so the column is not found.
+    if table_in_response:
+        table_data = schema_data.get("tables", {}).get(table, {})
+        columns_present = set((table_data.get("columns") or {}).keys())
+    else:
+        # Table was dropped by the column filter — fetch the column universe
+        # from the live resolver to build the "available" hint.
+        # is_exact_target(table) returned True above, so table is not None.
+        assert table is not None
+        full_profile = resolver.profile_table(source=source, schema=schema, table=table)
+        raw_cols = (
+            full_profile.get("sources", {})
+            .get(source, {})
+            .get("schemas", {})
+            .get(schema, {})
+            .get("tables", {})
+            .get(table, {})
+            .get("columns", {})
+        )
+        columns_present = set(raw_cols.keys())
+
+    if column not in columns_present:
+        exc = DatafaceError.from_code(
+            DF_EXECUTE_COLUMN_NOT_FOUND,
+            col_name=column,
+            table=table,
+            schema=schema,
+            source=source,
+            available=(
+                ", ".join(sorted(columns_present)) if columns_present else "(none)"
+            ),
+        )
+        return SchemaResponse(
+            success=False,
+            errors=[str(exc)],
+            structured_errors=[exc.to_structured()],
+        )
+
+    return None
+
+
+def _validate_table_profiles(
+    response: SchemaResponse,
+    source: str,
+    schema: str | None,
+    table: str | None,
+    column: str | None,
+) -> None:
+    """Validate level-4 profile-table dicts against the agent API shape."""
+    if schema is None or table is None or column is not None:
+        return
+    schemas = response.sources.get(source, {}).get("schemas", {})
+    for schema_entry in schemas.values():
+        tables = schema_entry.get("tables", {})
+        for raw_profile in tables.values():
+            if isinstance(raw_profile, dict):
+                TableProfile.model_validate(raw_profile)
+
+
+def _list_sources_response(adapter_registry: AdapterRegistry) -> SchemaResponse:
+    sources: dict[str, dict[str, Any]] = {}
+    for entry in adapter_registry.list_sql_sources():
+        source_entry = dict(entry)
+        name = str(source_entry.pop("name"))
+        sources[name] = source_entry
+    return SchemaResponse(
+        sources=sources,
+        _meta=ResponseMeta(
+            retrieved_at=datetime.now(timezone.utc).isoformat(),
+            sources_consulted=[],
+        ),
+    )
+
+
+def _filter_tables_by_search(
+    response: SchemaResponse,
+    source: str,
+    rx: re.Pattern[str],
+    resolver: LayeredSchemaResolver,
+) -> None:
+    """Remove tables from a level-3 response that don't match ``rx``.
+
+    A table passes when the compiled pattern matches its name, or when the
+    cache has a profile for the table and any cached column name matches.
+    Un-profiled tables (no cache entry) are filtered on name only —
+    no extra DB queries are issued.
+
+    Reads the cache file at most once per schema (not once per table).
+
+    Mutates ``response`` in place.
+    """
+    schemas = response.sources.get(source, {}).get("schemas", {})
+    for schema_name, schema_entry in schemas.items():
+        tables = schema_entry.get("tables")
+        if not tables:
+            continue
+        # Single read for the whole schema — avoids N file-parses in the loop.
+        # When no cache source is configured (OSS dbt-only mode), there are no
+        # profiled column names; filtering falls back to name-only matching.
+        col_names_by_table: dict[str, set[str]] = (
+            resolver.cache.column_names_for_schema(schema_name)
+            if resolver.cache is not None
+            else {}
+        )
+        to_drop = [
+            name
+            for name in tables
+            if not _table_matches(rx, name, col_names_by_table.get(name))
+        ]
+        for name in to_drop:
+            del tables[name]
+
+
+def _attach_filtered_empty_message(
+    response: SchemaResponse,
+    source: str,
+    pattern: str,
+) -> None:
+    """Set a filter-specific message when table_search empties the table list.
+
+    Prevents the "no cache / run dft inspect" guidance from being shown when
+    the result is empty solely because the regexp matched nothing.
+    """
+    schemas = response.sources.get(source, {}).get("schemas", {})
+    total_tables = sum(len(se.get("tables") or {}) for se in schemas.values())
+    if total_tables == 0:
+        response.message = f"No tables matched table_search={pattern!r}."
+
+
+def _table_matches(
+    rx: re.Pattern[str],
+    table_name: str,
+    col_names: set[str] | None,
+) -> bool:
+    """Return True if ``rx`` matches the table name or any column name.
+
+    ``col_names`` is ``None`` for un-profiled tables — name-only match applies.
+    No DB queries are issued.
+    """
+    if rx.search(table_name):
+        return True
+    return col_names is not None and any(rx.search(c) for c in col_names)
+
+
+def _validate_schema_request(
+    source: str | None,
+    schema: str | None,
+    table: str | None,
+    column: str | None,
+    force_refresh: bool,
+    table_search: str | None,
+) -> SchemaResponse | None:
+    if table_search is not None and (
+        source is None or schema is None or table is not None or column is not None
+    ):
+        return SchemaResponse(
+            success=False,
+            errors=[
+                "table_search is only valid at the schema level "
+                "(source and schema provided, without table or column)"
+            ],
+        )
+    if force_refresh and (
+        source is None or schema is None or table is None or column is not None
+    ):
+        return SchemaResponse(
+            success=False,
+            errors=[
+                "force_refresh is only valid when querying a specific table "
+                "(source, schema, and table, without column)"
+            ],
+        )
+    if column is not None and (table is None or source is None or schema is None):
+        return SchemaResponse(
+            success=False,
+            errors=["Must provide source, schema, and table when querying a column"],
+        )
+    if table is not None and (source is None or schema is None):
+        return SchemaResponse(
+            success=False,
+            errors=["Must provide source and schema when querying a table"],
+        )
+    if schema is not None and source is None:
+        return SchemaResponse(
+            success=False,
+            errors=["Must provide source when querying a schema"],
+        )
+    return None
+
+
+def _dispatch_schema_request(
+    resolver: LayeredSchemaResolver,
+    source: str,
+    schema: str | None,
+    table: str | None,
+    column: str | None,
+    force_refresh: bool,
+    lineage_depth: int,
+) -> dict[str, Any]:
+    if column is not None:
+        # mypy: guaranteed by _validate_schema_request before dispatch.
+        assert schema is not None and table is not None
+        return resolver.profile_column(
+            source=source, schema=schema, table=table, column=column
+        )
+    if table is not None:
+        # mypy: guaranteed by _validate_schema_request before dispatch.
+        assert schema is not None
+        return resolver.profile_table(
+            source=source,
+            schema=schema,
+            table=table,
+            fresh=force_refresh,
+            lineage_depth=lineage_depth,
+        )
+    if schema is not None:
+        return resolver.list_tables(source=source, schema=schema)
+    return resolver.list_schemas(source=source)
+
+
+def _project_uses_direct_file_queries(project_root: Path) -> bool:
+    from dataface.agent_api.dashboards import list_dashboards
+
+    for dashboard in list_dashboards(directory=project_root, recursive=True).dashboards:
+        try:
+            content = dashboard.absolute_path.read_text()
+        except OSError:
+            continue
+        if (
+            "read_csv(" in content
+            or "read_parquet(" in content
+            or "read_json(" in content
+        ):
+            return True
+    return False