kiri-cli 1.0.0b3__py3-none-any.whl

kiri_core/lifecycle/source_introspection.py

@@ -0,0 +1,488 @@
+"""Deterministic source-schema introspection helpers (ADR 0090)."""
+
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Protocol
+
+import yaml
+from pydantic import BaseModel, ConfigDict
+
+from kiri_core.schema.ingestion import (
+    DatabaseSpec,
+    IngestionFormat,
+    IngestionMode,
+    RestApiIncremental,
+    RestApiPagination,
+    RestApiSpec,
+)
+from kiri_core.schema.source import Source
+from kiri_core.types import Classification
+
+
+class InferredColumn(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    name: str
+    data_type: str = "string"
+    nullable: bool = True
+    description: str | None = None
+    pii: bool = False
+
+
+class RestApiIntrospectionResult(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    record_path: str | None
+    columns: tuple[InferredColumn, ...]
+    pagination: RestApiPagination
+    incremental: RestApiIncremental | None = None
+
+
+class SourceIntrospectionDraft(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    yaml_text: str
+    source_path: Path | None = None
+    note: str
+
+
+class AIColumnAnnotation(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    name: str
+    data_type: str | None = None
+    description: str | None = None
+    pii: bool | None = None
+
+
+class AISchemaAnnotations(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    columns: tuple[AIColumnAnnotation, ...] = ()
+    watermark_column: str | None = None
+    watermark_reason: str | None = None
+
+
+class AIGatewayLike(Protocol):
+    def generate(
+        self,
+        *,
+        prompt_name: str,
+        variables: dict[str, Any],
+        caller: str,
+        classification: str,
+    ) -> Any: ...
+
+
+_TIMESTAMP_NAME_RE = re.compile(
+    r"(updated|modified|created|time|timestamp|date|datum|systemdatum)",
+    re.I,
+)
+_ID_NAME_RE = re.compile(r"(^id$|_id$|id$|nummer$|nr$)", re.I)
+_PII_NAME_RE = re.compile(r"(email|e[-_]?post|phone|telefon|personnummer|ssn|name|namn)", re.I)
+
+
+def infer_rest_api_schema(payload: Any) -> RestApiIntrospectionResult:
+    """Infer record path, columns, pagination, and watermark from one JSON payload."""
+    record_path, rows = _best_record_array(payload)
+    if not rows:
+        raise ValueError("sample JSON contains no array of object records to introspect")
+    columns = _infer_columns(rows)
+    pagination = _infer_pagination(payload)
+    incremental = _infer_incremental(columns)
+    return RestApiIntrospectionResult(
+        record_path=record_path,
+        columns=tuple(columns),
+        pagination=pagination,
+        incremental=incremental,
+    )
+
+
+def apply_ai_annotations(
+    deterministic: RestApiIntrospectionResult,
+    annotations: AISchemaAnnotations,
+) -> RestApiIntrospectionResult:
+    """Apply AI annotations only to columns that exist in the deterministic sample."""
+    by_name = {col.name: col for col in deterministic.columns}
+    updated = dict(by_name)
+    for ann in annotations.columns:
+        current = updated.get(ann.name)
+        if current is None:
+            continue
+        data_type = ann.data_type if ann.data_type in _TYPE_VOCAB else current.data_type
+        updated[ann.name] = current.model_copy(
+            update={
+                "data_type": data_type,
+                "description": ann.description or current.description,
+                "pii": current.pii if ann.pii is None else bool(ann.pii),
+            }
+        )
+    incremental = deterministic.incremental
+    if annotations.watermark_column in by_name:
+        incremental = RestApiIncremental(
+            incremental_column=str(annotations.watermark_column),
+            request_param="since",
+        )
+    return deterministic.model_copy(
+        update={
+            "columns": tuple(updated[name] for name in by_name),
+            "incremental": incremental,
+        }
+    )
+
+
+def infer_with_ai(
+    deterministic: RestApiIntrospectionResult,
+    sample_rows: list[dict[str, Any]],
+    *,
+    gateway: AIGatewayLike,
+    source_name: str,
+    table_name: str,
+    classification: str,
+) -> RestApiIntrospectionResult:
+    """Ask the audited AI gateway for semantic annotations over sampled columns."""
+    response = gateway.generate(
+        prompt_name="infer_source_schema",
+        variables={
+            "source_name": source_name,
+            "table_name": table_name,
+            "columns": json.dumps([c.model_dump() for c in deterministic.columns], sort_keys=True),
+            "sample_rows": json.dumps(sample_rows[:20], default=str, sort_keys=True),
+        },
+        caller="cli.source.introspect.ai",
+        classification=classification,
+    )
+    raw = _strip_fences(str(response.content))
+    data = json.loads(raw)
+    annotations = AISchemaAnnotations.model_validate(data)
+    return apply_ai_annotations(deterministic, annotations)
+
+
+def render_rest_api_source_draft(
+    *,
+    project_root: Path,
+    source_name: str,
+    owner: str,
+    classification: Classification,
+    domain: str,
+    table_name: str,
+    rest_api: RestApiSpec,
+    inference: RestApiIntrospectionResult,
+    ai_used: bool = False,
+) -> SourceIntrospectionDraft:
+    """Render an unreviewed source YAML draft from REST introspection."""
+    rest_payload = rest_api.model_copy(
+        update={
+            "record_path": inference.record_path,
+            "pagination": inference.pagination,
+            "incremental": inference.incremental,
+        }
+    ).model_dump(mode="python", exclude_none=True)
+    payload = _source_payload(
+        source_name=source_name,
+        owner=owner,
+        classification=classification,
+        domain=domain,
+        table_name=table_name,
+        columns=inference.columns,
+        ingestion={
+            "mode": IngestionMode.REST_API.value,
+            "format": IngestionFormat.JSON.value,
+            "rest_api": rest_payload,
+        },
+        note=(
+            "UNREVIEWED source-introspection draft. Review inferred columns/types and "
+            "remove kiri.introspection_unreviewed to accept."
+            + (" AI annotations were applied through AIGateway." if ai_used else "")
+        ),
+    )
+    Source.model_validate(payload["sources"][0])
+    return SourceIntrospectionDraft(
+        yaml_text=str(yaml.safe_dump(payload, sort_keys=False, default_flow_style=False)),
+        source_path=project_root / "sources" / f"{source_name}.yml",
+        note=payload["sources"][0]["kiri"]["introspection_note"],
+    )
+
+
+def render_database_source_draft(
+    *,
+    project_root: Path,
+    source_name: str,
+    owner: str,
+    classification: Classification,
+    domain: str,
+    table_name: str,
+    database: DatabaseSpec,
+    columns: tuple[InferredColumn, ...],
+) -> SourceIntrospectionDraft:
+    payload = _source_payload(
+        source_name=source_name,
+        owner=owner,
+        classification=classification,
+        domain=domain,
+        table_name=table_name,
+        columns=columns,
+        ingestion={
+            "mode": IngestionMode.DATABASE.value,
+            "format": IngestionFormat.JSON.value,
+            "database": database.model_dump(mode="python", exclude_none=True),
+        },
+        note=(
+            "UNREVIEWED database-catalog introspection draft. Review inferred "
+            "columns/types and remove kiri.introspection_unreviewed to accept."
+        ),
+    )
+    Source.model_validate(payload["sources"][0])
+    return SourceIntrospectionDraft(
+        yaml_text=str(yaml.safe_dump(payload, sort_keys=False, default_flow_style=False)),
+        source_path=project_root / "sources" / f"{source_name}.yml",
+        note=payload["sources"][0]["kiri"]["introspection_note"],
+    )
+
+
+def _source_payload(
+    *,
+    source_name: str,
+    owner: str,
+    classification: Classification,
+    domain: str,
+    table_name: str,
+    columns: tuple[InferredColumn, ...],
+    ingestion: dict[str, Any],
+    note: str,
+) -> dict[str, Any]:
+    return {
+        "sources": [
+            {
+                "name": source_name,
+                "owner": owner,
+                "classification": classification.value,
+                "domain": domain,
+                "kiri": {
+                    "introspection_unreviewed": True,
+                    "introspection_note": note,
+                },
+                "tables": [
+                    {
+                        "name": table_name,
+                        "columns": [_column_source_payload(col) for col in columns],
+                        "ingestion": ingestion,
+                    }
+                ],
+            }
+        ]
+    }
+
+
+def _column_source_payload(col: InferredColumn) -> dict[str, Any]:
+    payload: dict[str, Any] = {
+        "name": col.name,
+        "data_type": col.data_type,
+    }
+    if col.description:
+        payload["description"] = col.description
+    if col.pii:
+        payload["pii"] = True
+    return payload
+
+
+def _best_record_array(payload: Any) -> tuple[str | None, list[dict[str, Any]]]:
+    candidates: list[tuple[str | None, list[dict[str, Any]], int]] = []
+
+    def walk(value: Any, path: str | None) -> None:
+        if isinstance(value, list):
+            if value and all(isinstance(item, dict) for item in value):
+                rows = [item for item in value if isinstance(item, dict)]
+                if _has_homogeneous_keys(rows) and _has_record_scalar(rows):
+                    common_keys = set(rows[0])
+                    for row in rows[1:]:
+                        common_keys &= set(row)
+                    candidates.append((path, rows, len(rows) * 10 + len(common_keys)))
+            return
+        if isinstance(value, dict):
+            for key, child in value.items():
+                child_path = key if path in (None, "") else f"{path}.{key}"
+                walk(child, child_path)
+
+    walk(payload, None)
+    if not candidates:
+        return None, []
+    path, rows, _score = max(
+        candidates,
+        key=lambda c: (c[2], 0 if c[0] is None else c[0].count(".")),
+    )
+    return path, rows
+
+
+def _has_homogeneous_keys(rows: list[dict[str, Any]]) -> bool:
+    """Accept record arrays only when every object has the same key set.
+
+    ADR 0090 drafts must be deterministic and must not union unrelated object
+    shapes. Exact key-set equality keeps array detection explainable and fail
+    closed for mixed envelopes such as ``[{id, name}, {error, message}]``.
+    """
+    if not rows:
+        return False
+    modal = set(rows[0])
+    return all(set(row) == modal for row in rows[1:])
+
+
+def _has_record_scalar(rows: list[dict[str, Any]]) -> bool:
+    return any(
+        not isinstance(value, (dict, list))
+        for row in rows
+        for value in row.values()
+        if value is not None
+    )
+
+
+def _infer_pagination(payload: Any) -> RestApiPagination:
+    """Infer pagination only from explicit payload evidence.
+
+    Non-paginated payloads default to ``kind='none'`` so unreviewed drafts do
+    not fabricate page parameters.
+    """
+    keys = _payload_keys(payload)
+    cursor_markers = {"cursor", "next_cursor", "nextCursor", "next"}
+    page_markers = {"page", "total_pages", "totalPages", "offset", "limit", "per_page"}
+    if keys & cursor_markers:
+        return RestApiPagination(kind="cursor", cursor_param="cursor")
+    if keys & page_markers:
+        return RestApiPagination(kind="page", page_param="page", limit_param="limit")
+    return RestApiPagination(kind="none")
+
+
+def _payload_keys(payload: Any) -> set[str]:
+    keys: set[str] = set()
+
+    def walk(value: Any) -> None:
+        if isinstance(value, dict):
+            keys.update(str(key) for key in value)
+            for child in value.values():
+                walk(child)
+        elif isinstance(value, list):
+            for item in value[:5]:
+                walk(item)
+
+    walk(payload)
+    return keys
+
+
+_TYPE_VOCAB = {"boolean", "int64", "float64", "timestamp", "date", "string"}
+
+
+def _infer_columns(rows: list[dict[str, Any]]) -> list[InferredColumn]:
+    names = sorted({name for row in rows for name in row})
+    columns: list[InferredColumn] = []
+    for name in names:
+        values = [row.get(name) for row in rows]
+        non_null = [v for v in values if v is not None]
+        columns.append(
+            InferredColumn(
+                name=name,
+                data_type=_infer_type(name, non_null),
+                nullable=len(non_null) != len(values),
+                pii=bool(_PII_NAME_RE.search(name)),
+            )
+        )
+    return columns
+
+
+def _infer_type(name: str, values: list[Any]) -> str:
+    if not values:
+        return "string"
+    if all(isinstance(v, bool) for v in values):
+        return "boolean"
+    if all(isinstance(v, int) and not isinstance(v, bool) for v in values):
+        return "int64"
+    if all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in values):
+        return "float64"
+    if all(isinstance(v, str) and _looks_like_datetime(v) for v in values):
+        return "timestamp" if _TIMESTAMP_NAME_RE.search(name) else "date"
+    return "string"
+
+
+def _looks_like_datetime(value: str) -> bool:
+    s = value.strip()
+    if not s:
+        return False
+    for candidate in (s, s.replace("Z", "+00:00")):
+        try:
+            datetime.fromisoformat(candidate)
+            return True
+        except ValueError:
+            continue
+    return False
+
+
+def _infer_incremental(columns: list[InferredColumn]) -> RestApiIncremental | None:
+    by_type = [c for c in columns if c.data_type in {"timestamp", "date"}]
+    named_time = [c for c in by_type if _TIMESTAMP_NAME_RE.search(c.name)]
+    if named_time:
+        return RestApiIncremental(incremental_column=named_time[0].name, request_param="since")
+    named_id = [c for c in columns if c.data_type == "int64" and _ID_NAME_RE.search(c.name)]
+    if named_id:
+        return RestApiIncremental(incremental_column=named_id[0].name, request_param="since")
+    return None
+
+
+def columns_from_database_catalog(rows: list[tuple[str, str, bool]]) -> tuple[InferredColumn, ...]:
+    """Convert database catalog rows ``(name, db_type, nullable)`` to source columns."""
+    return tuple(
+        InferredColumn(name=name, data_type=_map_database_type(db_type), nullable=nullable)
+        for name, db_type, nullable in rows
+    )
+
+
+def _map_database_type(db_type: str) -> str:
+    t = db_type.lower()
+    if any(part in t for part in ("int", "serial", "number")):
+        return "int64"
+    if any(part in t for part in ("double", "float", "real", "decimal", "numeric")):
+        return "float64"
+    if "bool" in t:
+        return "boolean"
+    if any(part in t for part in ("timestamp", "datetime")):
+        return "timestamp"
+    if "date" in t:
+        return "date"
+    return "string"
+
+
+def _strip_fences(text: str) -> str:
+    s = text.strip()
+    if not s.startswith("```"):
+        return s
+    lines = s.splitlines()
+    if lines and lines[0].startswith("```"):
+        lines = lines[1:]
+    if lines and lines[-1].startswith("```"):
+        lines = lines[:-1]
+    return "\n".join(lines).strip()
+
+
+def sample_rows_at_path(payload: Any, record_path: str | None) -> list[dict[str, Any]]:
+    current = resolve_dotted_path(payload, record_path)
+    if isinstance(current, list):
+        return [item for item in current if isinstance(item, dict)]
+    return []
+
+
+def resolve_dotted_path(payload: Any, record_path: str | None) -> Any:
+    """Resolve an ADR 0090 dotted JSON object path using literal key names."""
+    current = payload
+    if record_path:
+        for segment in record_path.split("."):
+            if not isinstance(current, dict):
+                return None
+            current = current.get(segment)
+    return current
+
+
+def column_type_counts(columns: tuple[InferredColumn, ...]) -> Counter[str]:
+    return Counter(col.data_type for col in columns)