interloper-toolkit 0.43.0__py3-none-any.whl

interloper_toolkit/lineage.py

@@ -0,0 +1,241 @@
+"""Lineage tools — dependency analysis, impact assessment, and DAG traversal."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from uuid import UUID
+
+from interloper_toolkit.context import ToolkitContext
+from interloper_toolkit.models import (
+    AssetRef,
+    CrossSourceDependencies,
+    CrossSourceEdge,
+    DependencyEdge,
+    DownstreamResult,
+    ImpactAnalysis,
+    LineageItem,
+    LineageResult,
+    ToolError,
+    UpstreamResult,
+)
+
+
+def get_upstream(ctx: ToolkitContext, asset_id: str) -> UpstreamResult | ToolError:
+    """Get the direct upstream dependencies of an asset.
+
+    Args:
+        asset_id: UUID of the asset to inspect.
+
+    Returns the list of upstream assets that this asset depends on,
+    including the parameter name used for each dependency.
+    """
+    try:
+        deps = ctx.store.list_relations(ctx.org_id, type="dependency")
+        target = UUID(asset_id)
+
+        upstream = []
+        for dep in deps:
+            if dep.src_id == target:
+                asset = ctx.store.get_component(dep.dst_id, kind="asset")
+                upstream.append(DependencyEdge(
+                    asset_id=str(dep.dst_id),
+                    param_name=dep.slot,
+                    asset_key=asset.key,
+                    source_id=str(asset.parent_id),
+                ))
+
+        return UpstreamResult(asset_id=asset_id, upstream=upstream)
+    except Exception as e:
+        return ToolError(error=str(e))
+
+
+def get_downstream(ctx: ToolkitContext, asset_id: str) -> DownstreamResult | ToolError:
+    """Get the direct downstream dependents of an asset.
+
+    Args:
+        asset_id: UUID of the asset to inspect.
+
+    Returns the list of assets that directly depend on this asset.
+    """
+    try:
+        deps = ctx.store.list_relations(ctx.org_id, type="dependency")
+        target = UUID(asset_id)
+
+        downstream = []
+        for dep in deps:
+            if dep.dst_id == target:
+                asset = ctx.store.get_component(dep.src_id, kind="asset")
+                downstream.append(DependencyEdge(
+                    asset_id=str(dep.src_id),
+                    param_name=dep.slot,
+                    asset_key=asset.key,
+                    source_id=str(asset.parent_id),
+                ))
+
+        return DownstreamResult(asset_id=asset_id, downstream=downstream)
+    except Exception as e:
+        return ToolError(error=str(e))
+
+
+def get_full_lineage(ctx: ToolkitContext, asset_id: str, direction: str = "upstream") -> LineageResult | ToolError:
+    """Recursively traverse the full lineage of an asset.
+
+    Args:
+        asset_id: UUID of the asset to start from.
+        direction: Either 'upstream' (ancestors) or 'downstream' (dependents).
+
+    Returns an ordered list of assets in the lineage chain with depth levels.
+    """
+    try:
+        adj, asset_info = _build_adjacency(ctx, direction)
+        target = UUID(asset_id)
+
+        visited: set[UUID] = set()
+        result: list[LineageItem] = []
+        queue: list[tuple[UUID, int]] = [(target, 0)]
+
+        while queue:
+            current, depth = queue.pop(0)
+            if current in visited:
+                continue
+            visited.add(current)
+            if current != target:
+                info = asset_info.get(current, {})
+                result.append(LineageItem(asset_id=str(current), depth=depth, **info))
+            for neighbor in adj.get(current, []):
+                if neighbor not in visited:
+                    queue.append((neighbor, depth + 1))
+
+        return LineageResult(
+            asset_id=asset_id,
+            direction=direction,
+            lineage_count=len(result),
+            lineage=result,
+        )
+    except Exception as e:
+        return ToolError(error=str(e))
+
+
+def impact_analysis(ctx: ToolkitContext, asset_id: str) -> ImpactAnalysis | ToolError:
+    """Analyze the downstream impact if an asset fails or is disabled.
+
+    Args:
+        asset_id: UUID of the asset to analyze.
+
+    Returns all downstream assets grouped by source, with total affected count.
+    """
+    try:
+        adj, asset_info = _build_adjacency(ctx, "downstream")
+        target = UUID(asset_id)
+
+        visited: set[UUID] = set()
+        affected: list[LineageItem] = []
+        queue: list[tuple[UUID, int]] = [(target, 0)]
+
+        while queue:
+            current, depth = queue.pop(0)
+            if current in visited:
+                continue
+            visited.add(current)
+            if current != target:
+                info = asset_info.get(current, {})
+                affected.append(LineageItem(asset_id=str(current), depth=depth, **info))
+            for neighbor in adj.get(current, []):
+                if neighbor not in visited:
+                    queue.append((neighbor, depth + 1))
+
+        # Group by source
+        by_source: dict[str, list[LineageItem]] = defaultdict(list)
+        for item in affected:
+            by_source[item.source_key or "unknown"].append(item)
+
+        return ImpactAnalysis(
+            asset_id=asset_id,
+            total_affected=len(affected),
+            by_source=dict(by_source),
+        )
+    except Exception as e:
+        return ToolError(error=str(e))
+
+
+def cross_source_dependencies(ctx: ToolkitContext) -> CrossSourceDependencies | ToolError:
+    """List all dependency edges that cross source boundaries.
+
+    Returns edges where the upstream and downstream assets belong to
+    different sources.
+    """
+    try:
+        deps = ctx.store.list_relations(ctx.org_id, type="dependency")
+        assets = ctx.store.list_components(ctx.org_id, kinds=["asset"])
+
+        asset_source: dict[UUID, UUID | None] = {}
+        asset_info: dict[UUID, AssetRef] = {}
+        for a in assets:
+            if not a.id:
+                continue
+            asset_source[a.id] = a.parent_id
+            asset_info[a.id] = AssetRef(asset_key=a.key, source_id=str(a.parent_id) if a.parent_id else "")
+
+        cross_deps = []
+        for dep in deps:
+            src_down = asset_source.get(dep.src_id)
+            src_up = asset_source.get(dep.dst_id)
+            if src_down and src_up and src_down != src_up:
+                cross_deps.append(CrossSourceEdge(
+                    downstream_asset_id=str(dep.src_id),
+                    downstream=asset_info.get(dep.src_id, AssetRef()),
+                    upstream_asset_id=str(dep.dst_id),
+                    upstream=asset_info.get(dep.dst_id, AssetRef()),
+                    param_name=dep.slot,
+                ))
+
+        return CrossSourceDependencies(cross_source_count=len(cross_deps), dependencies=cross_deps)
+    except Exception as e:
+        return ToolError(error=str(e))
+
+
+# -- Helpers -------------------------------------------------------------------
+
+
+def _build_adjacency(
+    ctx: ToolkitContext,
+    direction: str,
+) -> tuple[dict[UUID, list[UUID]], dict[UUID, dict[str, str]]]:
+    """Build an adjacency map and asset info lookup from all dependencies.
+
+    Args:
+        ctx: The toolkit context.
+        direction: 'upstream' builds child→parents; 'downstream' builds parent→children.
+
+    Returns:
+        ``(adjacency_map, asset_info_map)`` — info values feed
+        :class:`LineageItem` kwargs (asset_key, source_id, source_key).
+    """
+    deps = ctx.store.list_relations(ctx.org_id, type="dependency")
+    assets = ctx.store.list_components(ctx.org_id, kinds=["asset"])
+
+    asset_info: dict[UUID, dict[str, str]] = {}
+    source_keys: dict[UUID, str] = {}
+    for a in assets:
+        if not a.id:
+            continue
+        asset_info[a.id] = {"asset_key": a.key, "source_id": str(a.parent_id) if a.parent_id else ""}
+        if a.parent and a.parent_id:
+            source_keys[a.parent_id] = a.parent.key
+
+    # Add source_key to asset_info
+    for uid, info in asset_info.items():
+        sid_str = info["source_id"]
+        if sid_str:
+            info["source_key"] = source_keys.get(UUID(sid_str), "unknown")
+        else:
+            info["source_key"] = ""
+
+    adj: dict[UUID, list[UUID]] = defaultdict(list)
+    for dep in deps:
+        if direction == "upstream":
+            adj[dep.src_id].append(dep.dst_id)
+        else:
+            adj[dep.dst_id].append(dep.src_id)
+
+    return adj, asset_info

interloper_toolkit/models.py

@@ -0,0 +1,386 @@
+"""Typed result models for the toolkit surface.
+
+Every toolkit function returns ``<SuccessModel> | ToolError`` — a union
+discriminated by the literal ``status`` field, preserving the
+``{"status": "success" | "error"}`` envelope both AI surfaces already
+speak. The models are the tool contract: MCP derives per-tool output
+schemas from them, and the fields of row-projecting models act as an
+allowlist of what leaves the platform.
+
+Row-shaped payloads deliberately embed the interloper-db models (``Run``,
+``Backfill``, ``Event``, ``Component`` — already pydantic via SQLModel)
+rather than duplicating their shape: the full-row contract predates this
+module, and embedding makes the coupling visible in the signature. The one
+exception is :class:`ComponentSummary`, where projecting is the point —
+sensitive kinds must never expose config or credential material, so the
+model's fields fail closed instead of relying on conditional key insertion.
+
+Catalog definition payloads stay ``dict[str, Any]``: their content is
+definition-specific JSON schema material with no fixed shape to type.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Literal
+
+from interloper_db.models import AssetExecution, Backfill, Component, Event, Run
+from pydantic import BaseModel
+
+
+class ToolError(BaseModel):
+    """A failed tool call, as a structured result rather than an exception."""
+
+    status: Literal["error"] = "error"
+    error: str
+    valid_kinds: list[str] | None = None
+
+
+# -- Catalog --------------------------------------------------------------------
+
+
+class DefinitionCounts(BaseModel):
+    """Per-kind definition counts (the kind-less ``list_definitions`` call)."""
+
+    status: Literal["success"] = "success"
+    definition_counts: dict[str, int]
+    message: str
+
+
+class DefinitionEntry(BaseModel):
+    """One catalog definition, summarized for listing.
+
+    The trailing optional fields are kind-specific: ``asset_count`` /
+    ``in_collection`` / ``collection_count`` are populated for sources,
+    ``provider`` / ``required_fields`` / ``oauth`` / ``oauth_available``
+    for connections.
+    """
+
+    key: str
+    name: str
+    description: str | None = None
+    icon: str | None = None
+    tags: list[str] = []
+    asset_count: int | None = None
+    in_collection: bool | None = None
+    collection_count: int | None = None
+    provider: str | None = None
+    required_fields: list[str] | None = None
+    oauth: bool | None = None
+    oauth_available: bool | None = None
+
+
+class DefinitionList(BaseModel):
+    """Catalog definitions of one kind."""
+
+    status: Literal["success"] = "success"
+    kind: str
+    count: int
+    definitions: list[DefinitionEntry]
+
+
+class DefinitionDetail(BaseModel):
+    """One definition's full catalog payload (definition-specific shape)."""
+
+    status: Literal["success"] = "success"
+    definition: dict[str, Any]
+
+
+class AssetSchemaResult(BaseModel):
+    """The JSON schema and metadata of one asset within a source."""
+
+    status: Literal["success"] = "success"
+    source_key: str
+    asset_key: str
+    qualified_key: str
+    asset_schema: dict[str, Any] | None = None
+    partitioning: Any = None
+    tags: list[str] = []
+    requires: dict[str, Any] = {}
+    optional_requires: dict[str, Any] = {}
+
+
+class FieldMatch(BaseModel):
+    """One schema field matching a search query."""
+
+    source_key: str
+    asset_key: str
+    qualified_key: str
+    field_name: str
+    field_type: str
+    description: str
+
+
+class FieldSearchResult(BaseModel):
+    """Fields across all asset schemas matching a query."""
+
+    status: Literal["success"] = "success"
+    query: str
+    match_count: int
+    matches: list[FieldMatch]
+
+
+class SharedField(BaseModel):
+    """A field present in both compared schemas, with type agreement."""
+
+    field: str
+    type_a: str
+    type_b: str
+    type_match: bool
+
+
+class SchemaComparison(BaseModel):
+    """Side-by-side comparison of two asset schemas."""
+
+    status: Literal["success"] = "success"
+    asset_a: str
+    asset_b: str
+    shared_count: int
+    only_a_count: int
+    only_b_count: int
+    shared_fields: list[SharedField]
+    only_in_a: list[str]
+    only_in_b: list[str]
+
+
+# -- Collection -----------------------------------------------------------------
+
+
+class ComponentCounts(BaseModel):
+    """Per-kind component counts (the kind-less ``list_components`` call)."""
+
+    status: Literal["success"] = "success"
+    component_counts: dict[str, int]
+    message: str
+
+
+class ComponentSummary(BaseModel):
+    """A component instance, projected for listing.
+
+    Deliberately not the db row: these fields are an allowlist. ``config``
+    is only populated for non-sensitive kinds; credential-bearing kinds
+    surface identity and metadata alone.
+    """
+
+    id: str
+    key: str
+    name: str | None = None
+    type_name: str
+    created_at: datetime | None = None
+    config: dict[str, Any] | None = None
+    asset_count: int | None = None
+
+
+class ComponentList(BaseModel):
+    """The org's component instances of one kind."""
+
+    status: Literal["success"] = "success"
+    kind: str
+    count: int
+    components: list[ComponentSummary]
+
+
+# -- Lineage --------------------------------------------------------------------
+
+
+class DependencyEdge(BaseModel):
+    """A direct dependency edge from the perspective of one asset."""
+
+    asset_id: str
+    param_name: str
+    asset_key: str
+    source_id: str
+
+
+class UpstreamResult(BaseModel):
+    """Direct upstream dependencies of an asset."""
+
+    status: Literal["success"] = "success"
+    asset_id: str
+    upstream: list[DependencyEdge]
+
+
+class DownstreamResult(BaseModel):
+    """Direct downstream dependents of an asset."""
+
+    status: Literal["success"] = "success"
+    asset_id: str
+    downstream: list[DependencyEdge]
+
+
+class LineageItem(BaseModel):
+    """One asset in a lineage traversal, with its BFS depth."""
+
+    asset_id: str
+    depth: int
+    asset_key: str | None = None
+    source_id: str | None = None
+    source_key: str | None = None
+
+
+class LineageResult(BaseModel):
+    """The full recursive lineage of an asset in one direction."""
+
+    status: Literal["success"] = "success"
+    asset_id: str
+    direction: str
+    lineage_count: int
+    lineage: list[LineageItem]
+
+
+class ImpactAnalysis(BaseModel):
+    """All downstream assets affected by a failure, grouped by source."""
+
+    status: Literal["success"] = "success"
+    asset_id: str
+    total_affected: int
+    by_source: dict[str, list[LineageItem]]
+
+
+class AssetRef(BaseModel):
+    """A minimal asset reference on a cross-source edge."""
+
+    asset_key: str | None = None
+    source_id: str | None = None
+
+
+class CrossSourceEdge(BaseModel):
+    """A dependency edge whose endpoints belong to different sources."""
+
+    downstream_asset_id: str
+    downstream: AssetRef
+    upstream_asset_id: str
+    upstream: AssetRef
+    param_name: str
+
+
+class CrossSourceDependencies(BaseModel):
+    """All dependency edges crossing source boundaries."""
+
+    status: Literal["success"] = "success"
+    cross_source_count: int
+    dependencies: list[CrossSourceEdge]
+
+
+# -- Scheduling -----------------------------------------------------------------
+
+
+class JobList(BaseModel):
+    """The org's scheduled jobs (full component rows)."""
+
+    status: Literal["success"] = "success"
+    count: int
+    jobs: list[Component]
+
+
+class JobHealthStats(BaseModel):
+    """Success/failure statistics over a job's recent runs."""
+
+    total_recent_runs: int
+    success_count: int
+    failed_count: int
+    success_rate: float | None = None
+    avg_duration_seconds: float | None = None
+
+
+class JobHealth(BaseModel):
+    """A job's metadata plus health computed from its last runs."""
+
+    status: Literal["success"] = "success"
+    job: Component
+    health: JobHealthStats
+
+
+class RunList(BaseModel):
+    """Recent runs matching the filters."""
+
+    status: Literal["success"] = "success"
+    count: int
+    runs: list[Run]
+
+
+class RunDetail(BaseModel):
+    """One run with its event timeline and per-asset execution summary."""
+
+    status: Literal["success"] = "success"
+    run: Run
+    events: list[Event]
+    asset_executions: list[AssetExecution]
+
+
+class RunErrorEvent(BaseModel):
+    """One error event of a failed run."""
+
+    asset_key: str | None = None
+    error: str
+    timestamp: datetime
+
+
+class RunFailure(BaseModel):
+    """A failed run with its error events."""
+
+    run: Run
+    errors: list[RunErrorEvent]
+
+
+class FailureList(BaseModel):
+    """Recent failed runs with their errors."""
+
+    status: Literal["success"] = "success"
+    count: int
+    failures: list[RunFailure]
+
+
+class BackfillList(BaseModel):
+    """Backfills, optionally only the active ones."""
+
+    status: Literal["success"] = "success"
+    count: int
+    backfills: list[Backfill]
+
+
+# -- Analytics ------------------------------------------------------------------
+
+
+class RunHistorySummary(BaseModel):
+    """Aggregate run statistics over a look-back period."""
+
+    status: Literal["success"] = "success"
+    period_days: int
+    component_id: str | None = None
+    total_runs: int
+    by_status: dict[str, int]
+    success_rate: float | None = None
+    avg_duration_seconds: float | None = None
+
+
+class PartitionCoverage(BaseModel):
+    """Which partition dates in a range have successful runs."""
+
+    status: Literal["success"] = "success"
+    component_id: str
+    start_date: str
+    end_date: str
+    total_days: int
+    covered_days: int
+    missing_days: int
+    coverage_percent: float
+    missing_dates: list[str]
+
+
+class JobFreshness(BaseModel):
+    """One job's data freshness."""
+
+    job: Component
+    last_success_at: datetime | None = None
+    hours_since_success: float | None = None
+    stale: bool
+
+
+class FreshnessReport(BaseModel):
+    """Freshness across all enabled jobs."""
+
+    status: Literal["success"] = "success"
+    total_jobs: int
+    stale_count: int
+    jobs: list[JobFreshness]