diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 282a8f0..6a0bdca 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -23,7 +23,7 @@ jobs: uses: astral-sh/setup-uv@v6 - name: Install dependencies - run: uv pip install --system -e ".[dev]" + run: uv pip install --system -e ".[dev,mcp]" - name: Run Ruff run: ruff check . diff --git a/README.md b/README.md index f623a2f..b769207 100644 --- a/README.md +++ b/README.md @@ -26,8 +26,12 @@ FlowSense is designed to answer questions such as: - Median-based historical baselines - MAD-based robust Z-score drift detection - Severity classification -- Dependency-aware propagation analysis +- Task handoff delay analysis +- Task impact classification (`OWN_DRIFT`, `INHERITED_DELAY`, and `COMBINED`) +- Multi-hop and branching propagation analysis +- Primary root-cause selection - CLI-based DAG analysis +- MCP server integration ## Example @@ -40,13 +44,19 @@ Example output: ```text FlowSense Analysis — flowsense_demo -┏━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┓ -┃ Task ┃ Baseline ┃ Current ┃ Deviation ┃ Z-Score ┃ Severity ┃ -┡━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━┩ -│ extract │ 1.56s │ 1.61s │ +3.4% │ 0.17 │ NORMAL │ -│ transform │ 3.34s │ 9.61s │ +187.6% │ 7.61 │ CRITICAL │ -│ load │ 1.40s │ 2.11s │ +50.2% │ 3.17 │ MEDIUM │ -└───────────┴──────────┴─────────┴───────────┴─────────┴──────────┘ +┏━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━┓ +┃ Task ┃ Baseline ┃ Current ┃ Deviation ┃ Z-Score ┃ Severity ┃ Impact ┃ +┡━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━┩ +│ extract │ 1.56s │ 1.61s │ +3.4% │ 0.17 │ NORMAL │ NORMAL │ +│ transform │ 3.34s │ 9.61s │ +187.6% │ 7.61 │ CRITICAL │ OWN_DRIFT │ +│ load │ 1.40s │ 2.11s │ +50.2% │ 3.17 │ MEDIUM │ COMBINED │ +└───────────┴──────────┴─────────┴───────────┴─────────┴──────────┴───────────┘ + +Overall Severity: CRITICAL +Primary Origin: transform +Reason: OWN_DRIFT +Severity: CRITICAL +Propagation Score: 0.33 Propagation Analysis @@ -64,17 +74,16 @@ Apache Airflow Collector │ ▼ -Task Run History +Task Run and Handoff History │ ▼ - Drift Engine +Drift and Impact Analysis │ ▼ -Propagation Analysis +Propagation and Root-Cause Analysis │ ├── CLI - │ - └── MCP Server (planned) + └── MCP Server ``` ## Installation @@ -93,7 +102,7 @@ Create a virtual environment and install the project: ```bash uv venv --python 3.12 source .venv/bin/activate -uv pip install -e ".[dev]" +uv pip install -e ".[dev,mcp]" ``` ## Configuration @@ -126,6 +135,18 @@ Then run: flowsense analyze ``` +## MCP Server + +Start the FlowSense MCP server over stdio: + +```bash +flowsense-mcp +``` + +The server exposes the `analyze_airflow_dag` tool, which returns task drift, +handoff drift, impact classification, propagation paths, and primary root-cause +information for a DAG. + ## Development Run unit tests: @@ -167,7 +188,11 @@ src/flowsense/ ├── engine/ │ ├── drift.py │ ├── history.py -│ └── propagation.py +│ ├── impact.py +│ ├── propagation.py +│ ├── root_cause.py +│ └── timing.py +├── mcp/ └── models/ ``` @@ -194,13 +219,9 @@ Planned areas include: - DAG-level analysis models - configurable historical baseline windows - improved propagation scoring -- temporal delay analysis -- upstream/downstream impact separation -- root-cause analysis - change-point detection - trend detection - richer CLI reporting -- MCP server integration - broader Airflow compatibility testing ## License diff --git a/pyproject.toml b/pyproject.toml index d50a021..f966611 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,11 +21,12 @@ dev = [ ] mcp = [ - "mcp", + "mcp[cli]>=2.0", ] [project.scripts] flowsense = "flowsense.cli.main:app" +flowsense-mcp = "flowsense.mcp.server:main" [build-system] requires = ["setuptools>=75"] diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index 429e0c5..33d5604 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -4,10 +4,7 @@ from rich.console import Console from rich.table import Table -from flowsense.collector.airflow_client import AirflowClient -from flowsense.engine.drift import calculate_drift -from flowsense.engine.history import build_duration_history -from flowsense.engine.propagation import analyze_propagation +from flowsense.engine.analyzer import analyze_dag app = typer.Typer( name="flowsense", @@ -27,39 +24,12 @@ def main() -> None: def analyze( dag_id: str = typer.Argument( ..., - help="Airflow DAG ID to analyze.", + help="Airflow DAG id to analyze.", ), ) -> None: - client = AirflowClient() + analysis = analyze_dag(dag_id) - task_runs = client.collect_task_runs(dag_id) - - if not task_runs: - console.print(f"[red]No successful task runs found for DAG: {dag_id}[/red]") - raise typer.Exit(code=1) - - history = build_duration_history(task_runs) - - drift_results = {} - - for task_id, durations in history.items(): - result = calculate_drift( - task_id, - durations, - ) - - drift_results[task_id] = result - - dependencies = client.get_dag_dependencies(dag_id) - - propagation_results = analyze_propagation( - drift_results, - dependencies, - ) - - console.print() - console.print(f"[bold]FlowSense Analysis[/bold] — {dag_id}") - console.print() + console.print(f"\n[bold]FlowSense Analysis — {analysis.dag_id}[/bold]\n") table = Table() @@ -69,34 +39,38 @@ def analyze( table.add_column("Deviation") table.add_column("Z-Score") table.add_column("Severity") + table.add_column("Impact") + + for task_id, result in analysis.drift_results.items(): + impact = analysis.task_impacts.get(task_id) + impact_label = impact.classification if impact else "-" - for result in drift_results.values(): table.add_row( - result.task_id, + task_id, f"{result.baseline:.2f}s", f"{result.current:.2f}s", f"{result.deviation_percent:+.1f}%", f"{result.robust_z_score:.2f}", result.severity, + impact_label, ) console.print(table) - console.print() - console.print("[bold]Propagation Analysis[/bold]") + console.print(f"\nOverall Severity: [bold]{analysis.overall_severity}[/bold]") - if not propagation_results: - console.print("No propagation detected.") - return - - for propagation in propagation_results: - console.print() - console.print(f"Origin: [bold]{propagation.origin_task}[/bold]") - - console.print("Path: " + " -> ".join(propagation.path)) - - console.print(f"Propagation Score: {propagation.propagation_score:.2f}") + if analysis.primary_origin: + console.print(f"Primary Origin: [bold]{analysis.primary_origin.task_id}[/bold]") + console.print(f"Reason: [bold]{analysis.primary_origin.classification}[/bold]") + console.print(f"Severity: [bold]{analysis.primary_origin.severity}[/bold]") + console.print( + f"Propagation Score: {analysis.primary_origin.propagation_score:.2f}" + ) + if analysis.propagation_results: + console.print("\n[bold]Propagation Analysis[/bold]\n") -if __name__ == "__main__": - app() + for result in analysis.propagation_results: + console.print(f"Origin: {result.origin_task}") + console.print(f"Path: {' -> '.join(result.path)}") + console.print(f"Propagation Score: {result.propagation_score:.2f}") diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py new file mode 100644 index 0000000..aa45c3c --- /dev/null +++ b/src/flowsense/engine/analyzer.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from flowsense.collector.airflow_client import AirflowClient +from flowsense.engine.drift import calculate_drift +from flowsense.engine.history import build_duration_history +from flowsense.engine.impact import classify_task_impact +from flowsense.engine.propagation import analyze_propagation +from flowsense.engine.root_cause import select_primary_origin +from flowsense.engine.timing import ( + build_handoff_history, + calculate_handoff_drift, +) +from flowsense.models import DAGAnalysis + + +def analyze_dag(dag_id: str) -> DAGAnalysis: + client = AirflowClient() + + task_runs = client.collect_task_runs(dag_id) + duration_history = build_duration_history(task_runs) + + drift_results = {} + + for task_id, durations in duration_history.items(): + try: + drift_results[task_id] = calculate_drift( + task_id=task_id, + durations=durations, + ) + except ValueError: + continue + + dependencies = client.get_dag_dependencies(dag_id) + + handoff_history = build_handoff_history( + task_runs=task_runs, + dependencies=dependencies, + ) + + handoff_drift_results = {} + + for edge, delays in handoff_history.items(): + _upstream_task, downstream_task = edge + + try: + handoff_drift_results[edge] = calculate_handoff_drift( + upstream_task=_upstream_task, + downstream_task=downstream_task, + handoff_delays=delays, + ) + except ValueError: + continue + + task_impacts = {} + + for task_id, task_drift in drift_results.items(): + upstream_handoff_drifts = [ + drift + for ( + _upstream_task, + downstream_task, + ), drift in handoff_drift_results.items() + if downstream_task == task_id + ] + + task_impacts[task_id] = classify_task_impact( + task_id=task_id, + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + + propagation_results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + severity_order = { + "NORMAL": 0, + "MEDIUM": 1, + "HIGH": 2, + "CRITICAL": 3, + } + + overall_severity = "NORMAL" + + all_drift_results = [ + *drift_results.values(), + *handoff_drift_results.values(), + ] + + if all_drift_results: + overall_severity = max( + all_drift_results, + key=lambda result: severity_order[result.severity], + ).severity + + primary_origin = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies=dependencies, + propagation_results=propagation_results, + ) + + return DAGAnalysis( + dag_id=dag_id, + runs_analyzed=len({run.dag_run_id for run in task_runs}), + overall_severity=overall_severity, + primary_origin=primary_origin, + drift_results=drift_results, + handoff_drift_results=handoff_drift_results, + task_impacts=task_impacts, + propagation_results=propagation_results, + dependencies=dependencies, + ) diff --git a/src/flowsense/engine/impact.py b/src/flowsense/engine/impact.py new file mode 100644 index 0000000..fd69ca0 --- /dev/null +++ b/src/flowsense/engine/impact.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult + + +@dataclass +class TaskImpact: + task_id: str + classification: str + task_severity: str + upstream_handoff_severity: str | None + + +def classify_task_impact( + task_id: str, + task_drift: DriftResult, + upstream_handoff_drifts: list[DriftResult], +) -> TaskImpact: + anomalous_handoffs = [ + drift + for drift in upstream_handoff_drifts + if drift.severity + in { + "MEDIUM", + "HIGH", + "CRITICAL", + } + ] + + task_is_anomalous = task_drift.severity in { + "MEDIUM", + "HIGH", + "CRITICAL", + } + + if task_is_anomalous and anomalous_handoffs: + classification = "COMBINED" + elif task_is_anomalous: + classification = "OWN_DRIFT" + elif anomalous_handoffs: + classification = "INHERITED_DELAY" + else: + classification = "NORMAL" + + upstream_handoff_severity = None + + if anomalous_handoffs: + severity_order = { + "NORMAL": 0, + "MEDIUM": 1, + "HIGH": 2, + "CRITICAL": 3, + } + + upstream_handoff_severity = max( + anomalous_handoffs, + key=lambda result: severity_order[result.severity], + ).severity + + return TaskImpact( + task_id=task_id, + classification=classification, + task_severity=task_drift.severity, + upstream_handoff_severity=upstream_handoff_severity, + ) diff --git a/src/flowsense/engine/propagation.py b/src/flowsense/engine/propagation.py index c5382af..507cdbc 100644 --- a/src/flowsense/engine/propagation.py +++ b/src/flowsense/engine/propagation.py @@ -20,49 +20,150 @@ class PropagationResult: propagation_score: float +def _build_reverse_dependencies( + dependencies: dict[str, list[str]], +) -> dict[str, list[str]]: + reverse_dependencies: dict[str, list[str]] = { + task_id: [] for task_id in dependencies + } + + for upstream_task, downstream_tasks in dependencies.items(): + for downstream_task in downstream_tasks: + reverse_dependencies.setdefault(downstream_task, []).append(upstream_task) + + return reverse_dependencies + + +def _has_anomalous_upstream( + task_id: str, + drift_results: dict[str, DriftResult], + reverse_dependencies: dict[str, list[str]], + visited: set[str] | None = None, +) -> bool: + if visited is None: + visited = set() + + if task_id in visited: + return False + + visited = {*visited, task_id} + + for upstream_task in reverse_dependencies.get(task_id, []): + upstream_drift = drift_results.get(upstream_task) + + if upstream_drift is None or upstream_drift.severity == "NORMAL": + continue + + if upstream_drift.severity in { + "HIGH", + "CRITICAL", + }: + return True + + if _has_anomalous_upstream( + task_id=upstream_task, + drift_results=drift_results, + reverse_dependencies=reverse_dependencies, + visited=visited, + ): + return True + + return False + + +def _find_propagation_paths( + origin_task: str, + current_task: str, + drift_results: dict[str, DriftResult], + dependencies: dict[str, list[str]], + path: list[str], + visited: set[str], +) -> list[list[str]]: + paths: list[list[str]] = [] + + for downstream_task in dependencies.get(current_task, []): + if downstream_task in visited: + continue + + downstream_drift = drift_results.get(downstream_task) + + if downstream_drift is None: + continue + + if downstream_drift.severity == "NORMAL": + continue + + next_path = [*path, downstream_task] + next_visited = {*visited, downstream_task} + + child_paths = _find_propagation_paths( + origin_task=origin_task, + current_task=downstream_task, + drift_results=drift_results, + dependencies=dependencies, + path=next_path, + visited=next_visited, + ) + + if child_paths: + paths.extend(child_paths) + else: + paths.append(next_path) + + return paths + + def analyze_propagation( drift_results: dict[str, DriftResult], dependencies: dict[str, list[str]], ) -> list[PropagationResult]: results: list[PropagationResult] = [] + reverse_dependencies = _build_reverse_dependencies(dependencies) + for task_id, drift in drift_results.items(): if drift.severity not in {"HIGH", "CRITICAL"}: continue - downstream = dependencies.get(task_id, []) + if _has_anomalous_upstream( + task_id=task_id, + drift_results=drift_results, + reverse_dependencies=reverse_dependencies, + ): + continue - affected: list[str] = [] + paths = _find_propagation_paths( + origin_task=task_id, + current_task=task_id, + drift_results=drift_results, + dependencies=dependencies, + path=[task_id], + visited={task_id}, + ) - for downstream_task in downstream: - downstream_drift = drift_results.get(downstream_task) + for path in paths: + affected_tasks = path[1:] - if downstream_drift is None: + if not affected_tasks: continue - if downstream_drift.severity != "NORMAL": - affected.append(downstream_task) - - if not affected: - continue - - origin_score = SEVERITY_SCORE[drift.severity] + origin_score = SEVERITY_SCORE[drift.severity] - downstream_scores = [ - SEVERITY_SCORE[drift_results[task].severity] for task in affected - ] + downstream_scores = [ + SEVERITY_SCORE[drift_results[task].severity] for task in affected_tasks + ] - propagation_score = sum(downstream_scores) / ( - len(downstream_scores) * origin_score - ) + propagation_score = sum(downstream_scores) / ( + len(downstream_scores) * origin_score + ) - results.append( - PropagationResult( - origin_task=task_id, - affected_tasks=affected, - path=[task_id, *affected], - propagation_score=propagation_score, + results.append( + PropagationResult( + origin_task=task_id, + affected_tasks=affected_tasks, + path=path, + propagation_score=propagation_score, + ) ) - ) return results diff --git a/src/flowsense/engine/root_cause.py b/src/flowsense/engine/root_cause.py new file mode 100644 index 0000000..4e2e03f --- /dev/null +++ b/src/flowsense/engine/root_cause.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact +from flowsense.engine.propagation import PropagationResult + + +@dataclass +class RootCauseResult: + task_id: str + classification: str + severity: str + propagation_score: float + + +SEVERITY_SCORE = { + "NORMAL": 0, + "MEDIUM": 1, + "HIGH": 2, + "CRITICAL": 3, +} + + +def _build_reverse_dependencies( + dependencies: dict[str, list[str]], +) -> dict[str, list[str]]: + reverse_dependencies: dict[str, list[str]] = {} + + for upstream_task, downstream_tasks in dependencies.items(): + reverse_dependencies.setdefault(upstream_task, []) + + for downstream_task in downstream_tasks: + reverse_dependencies.setdefault( + downstream_task, + [], + ).append(upstream_task) + + return reverse_dependencies + + +def _has_candidate_upstream( + task_id: str, + candidate_tasks: set[str], + reverse_dependencies: dict[str, list[str]], + visited: set[str] | None = None, +) -> bool: + if visited is None: + visited = set() + + if task_id in visited: + return False + + visited.add(task_id) + + for upstream_task in reverse_dependencies.get(task_id, []): + if upstream_task in candidate_tasks: + return True + + if _has_candidate_upstream( + task_id=upstream_task, + candidate_tasks=candidate_tasks, + reverse_dependencies=reverse_dependencies, + visited=visited, + ): + return True + + return False + + +def select_primary_origin( + drift_results: dict[str, DriftResult], + task_impacts: dict[str, TaskImpact], + dependencies: dict[str, list[str]], + propagation_results: list[PropagationResult], +) -> RootCauseResult | None: + candidate_tasks = { + task_id + for task_id, impact in task_impacts.items() + if impact.classification + in { + "OWN_DRIFT", + "COMBINED", + } + } + + if not candidate_tasks: + return None + + reverse_dependencies = _build_reverse_dependencies( + dependencies, + ) + + root_candidates = { + task_id + for task_id in candidate_tasks + if not _has_candidate_upstream( + task_id=task_id, + candidate_tasks=candidate_tasks, + reverse_dependencies=reverse_dependencies, + ) + } + + if not root_candidates: + root_candidates = candidate_tasks + + propagation_scores: dict[str, float] = {} + + for result in propagation_results: + current_score = propagation_scores.get( + result.origin_task, + 0.0, + ) + + propagation_scores[result.origin_task] = max( + current_score, + result.propagation_score, + ) + + def ranking(task_id: str) -> tuple[float, int]: + propagation_score = propagation_scores.get( + task_id, + 0.0, + ) + + drift = drift_results.get(task_id) + + severity_score = SEVERITY_SCORE[drift.severity] if drift else 0 + + return ( + propagation_score, + severity_score, + ) + + primary_task = max( + root_candidates, + key=ranking, + ) + + impact = task_impacts[primary_task] + drift = drift_results[primary_task] + + return RootCauseResult( + task_id=primary_task, + classification=impact.classification, + severity=drift.severity, + propagation_score=propagation_scores.get( + primary_task, + 0.0, + ), + ) diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py new file mode 100644 index 0000000..7c43cac --- /dev/null +++ b/src/flowsense/engine/timing.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult, calculate_drift +from flowsense.models import TaskRun + + +@dataclass +class HandoffTiming: + upstream_task: str + downstream_task: str + dag_run_id: str + handoff_delay: float + + +def calculate_handoff_delay( + upstream_run: TaskRun, + downstream_run: TaskRun, +) -> HandoffTiming: + if upstream_run.end_date is None: + raise ValueError("Upstream task end_date is required.") + + if downstream_run.start_date is None: + raise ValueError("Downstream task start_date is required.") + + if upstream_run.dag_run_id != downstream_run.dag_run_id: + raise ValueError("Task runs must belong to the same DAG run.") + + handoff_delay = (downstream_run.start_date - upstream_run.end_date).total_seconds() + + return HandoffTiming( + upstream_task=upstream_run.task_id, + downstream_task=downstream_run.task_id, + dag_run_id=upstream_run.dag_run_id, + handoff_delay=handoff_delay, + ) + + +def build_handoff_history( + task_runs: list[TaskRun], + dependencies: dict[str, list[str]], +) -> dict[tuple[str, str], list[float]]: + runs_by_id: dict[str, dict[str, TaskRun]] = {} + + for task_run in task_runs: + runs_by_id.setdefault(task_run.dag_run_id, {})[task_run.task_id] = task_run + + history: dict[tuple[str, str], list[float]] = {} + + for tasks_in_run in runs_by_id.values(): + for upstream_task, downstream_tasks in dependencies.items(): + upstream_run = tasks_in_run.get(upstream_task) + + if upstream_run is None: + continue + + for downstream_task in downstream_tasks: + downstream_run = tasks_in_run.get(downstream_task) + + if downstream_run is None: + continue + + try: + timing = calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + except ValueError: + continue + + edge = ( + upstream_task, + downstream_task, + ) + + history.setdefault(edge, []).append(timing.handoff_delay) + + return history + + +def calculate_handoff_drift( + upstream_task: str, + downstream_task: str, + handoff_delays: list[float], +) -> DriftResult: + edge_id = f"{upstream_task}->{downstream_task}" + + return calculate_drift( + task_id=edge_id, + durations=handoff_delays, + ) diff --git a/src/flowsense/mcp/__init__.py b/src/flowsense/mcp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py new file mode 100644 index 0000000..af8d483 --- /dev/null +++ b/src/flowsense/mcp/server.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from mcp.server import MCPServer + +from flowsense.engine.analyzer import analyze_dag +from flowsense.models import DAGAnalysis + +mcp = MCPServer("FlowSense Engine") + + +def serialize_analysis(analysis: DAGAnalysis) -> dict: + return { + "dag_id": analysis.dag_id, + "runs_analyzed": analysis.runs_analyzed, + "overall_severity": analysis.overall_severity, + "primary_origin": ( + { + "task_id": analysis.primary_origin.task_id, + "classification": analysis.primary_origin.classification, + "severity": analysis.primary_origin.severity, + "propagation_score": analysis.primary_origin.propagation_score, + } + if analysis.primary_origin + else None + ), + "drift_results": { + task_id: { + "baseline": result.baseline, + "current": result.current, + "mad": result.mad, + "robust_z_score": result.robust_z_score, + "deviation_percent": result.deviation_percent, + "severity": result.severity, + } + for task_id, result in analysis.drift_results.items() + }, + "handoff_drift_results": { + f"{upstream}->{downstream}": { + "baseline": result.baseline, + "current": result.current, + "mad": result.mad, + "robust_z_score": result.robust_z_score, + "deviation_percent": result.deviation_percent, + "severity": result.severity, + } + for ( + upstream, + downstream, + ), result in analysis.handoff_drift_results.items() + }, + "task_impacts": { + task_id: { + "classification": impact.classification, + "task_severity": impact.task_severity, + "upstream_handoff_severity": impact.upstream_handoff_severity, + } + for task_id, impact in analysis.task_impacts.items() + }, + "propagation_results": [ + { + "origin_task": result.origin_task, + "affected_tasks": result.affected_tasks, + "path": result.path, + "propagation_score": result.propagation_score, + } + for result in analysis.propagation_results + ], + "dependencies": analysis.dependencies, + } + + +@mcp.tool() +def analyze_airflow_dag(dag_id: str) -> dict: + """Analyze an Apache Airflow DAG for temporal drift and propagation.""" + analysis = analyze_dag(dag_id) + return serialize_analysis(analysis) + + +def main() -> None: + mcp.run() + + +if __name__ == "__main__": + main() diff --git a/src/flowsense/models/__init__.py b/src/flowsense/models/__init__.py index 7de82b2..133820d 100644 --- a/src/flowsense/models/__init__.py +++ b/src/flowsense/models/__init__.py @@ -1,3 +1,7 @@ +from flowsense.models.dag_analysis import DAGAnalysis from flowsense.models.task_run import TaskRun -__all__ = ["TaskRun"] +__all__ = [ + "DAGAnalysis", + "TaskRun", +] diff --git a/src/flowsense/models/dag_analysis.py b/src/flowsense/models/dag_analysis.py new file mode 100644 index 0000000..e963ffb --- /dev/null +++ b/src/flowsense/models/dag_analysis.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact +from flowsense.engine.propagation import PropagationResult +from flowsense.engine.root_cause import RootCauseResult + + +@dataclass +class DAGAnalysis: + dag_id: str + runs_analyzed: int + overall_severity: str + primary_origin: RootCauseResult | None + drift_results: dict[str, DriftResult] + handoff_drift_results: dict[tuple[str, str], DriftResult] + task_impacts: dict[str, TaskImpact] + propagation_results: list[PropagationResult] + dependencies: dict[str, list[str]] diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py new file mode 100644 index 0000000..f76d9b8 --- /dev/null +++ b/tests/test_analyzer.py @@ -0,0 +1,242 @@ +from datetime import UTC, datetime, timedelta +from unittest.mock import MagicMock, patch + +from flowsense.engine.analyzer import analyze_dag +from flowsense.models import TaskRun + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_identifies_primary_origin( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + client.collect_task_runs.return_value = [ + TaskRun( + dag_id="demo", + dag_run_id=f"run_{index}", + task_id="transform", + state="success", + duration=duration, + ) + for index, duration in enumerate( + [3.0, 3.1, 2.9, 3.0, 9.5], + start=1, + ) + ] + [ + TaskRun( + dag_id="demo", + dag_run_id=f"run_{index}", + task_id="load", + state="success", + duration=duration, + ) + for index, duration in enumerate( + [1.0, 1.1, 0.9, 1.0, 2.0], + start=1, + ) + ] + + client.get_dag_dependencies.return_value = { + "transform": ["load"], + "load": [], + } + + analysis = analyze_dag("demo") + + assert analysis.overall_severity == "CRITICAL" + assert analysis.primary_origin is not None + assert analysis.primary_origin.task_id == "transform" + assert len(analysis.propagation_results) == 1 + assert analysis.propagation_results[0].origin_task == "transform" + assert analysis.propagation_results[0].affected_tasks == ["load"] + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_identifies_isolated_primary_origin( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + client.collect_task_runs.return_value = [ + TaskRun( + dag_id="demo", + dag_run_id=f"run_{index}", + task_id="transform", + state="success", + duration=duration, + ) + for index, duration in enumerate( + [3.0, 3.1, 2.9, 3.0, 9.5], + start=1, + ) + ] + + client.get_dag_dependencies.return_value = { + "transform": [], + } + + analysis = analyze_dag("demo") + + assert analysis.propagation_results == [] + assert analysis.primary_origin is not None + assert analysis.primary_origin.task_id == "transform" + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_calculates_handoff_drift( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + base_time = datetime( + 2026, + 8, + 20, + 10, + 0, + tzinfo=UTC, + ) + + handoff_delays = [ + 2.0, + 2.1, + 1.9, + 2.0, + 8.0, + ] + + task_runs = [] + + for index, delay in enumerate( + handoff_delays, + start=1, + ): + dag_run_id = f"run_{index}" + + extract_start = base_time + timedelta(minutes=index) + extract_end = extract_start + timedelta(seconds=1) + + transform_start = extract_end + timedelta(seconds=delay) + transform_end = transform_start + timedelta(seconds=3) + + task_runs.extend( + [ + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="extract", + state="success", + start_date=extract_start, + end_date=extract_end, + duration=1.0, + ), + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="transform", + state="success", + start_date=transform_start, + end_date=transform_end, + duration=3.0, + ), + ] + ) + + client.collect_task_runs.return_value = task_runs + + client.get_dag_dependencies.return_value = { + "extract": ["transform"], + "transform": [], + } + + analysis = analyze_dag("demo") + + edge = ("extract", "transform") + + assert edge in analysis.handoff_drift_results + + handoff_drift = analysis.handoff_drift_results[edge] + + assert handoff_drift.current == 8.0 + assert handoff_drift.baseline == 2.0 + assert handoff_drift.severity == "CRITICAL" + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_uses_handoff_for_overall_severity( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + base_time = datetime( + 2026, + 8, + 20, + 10, + 0, + tzinfo=UTC, + ) + + handoff_delays = [ + 2.0, + 2.1, + 1.9, + 2.0, + 8.0, + ] + + task_runs = [] + + for index, delay in enumerate( + handoff_delays, + start=1, + ): + dag_run_id = f"run_{index}" + + extract_start = base_time + timedelta(minutes=index) + extract_end = extract_start + timedelta(seconds=1) + + transform_start = extract_end + timedelta(seconds=delay) + transform_end = transform_start + timedelta(seconds=3) + + task_runs.extend( + [ + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="extract", + state="success", + start_date=extract_start, + end_date=extract_end, + duration=1.0, + ), + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="transform", + state="success", + start_date=transform_start, + end_date=transform_end, + duration=3.0, + ), + ] + ) + + client.collect_task_runs.return_value = task_runs + + client.get_dag_dependencies.return_value = { + "extract": ["transform"], + "transform": [], + } + + analysis = analyze_dag("demo") + + assert analysis.drift_results["extract"].severity == "NORMAL" + assert analysis.drift_results["transform"].severity == "NORMAL" + + assert ( + analysis.handoff_drift_results[("extract", "transform")].severity == "CRITICAL" + ) + + assert analysis.overall_severity == "CRITICAL" diff --git a/tests/test_impact.py b/tests/test_impact.py new file mode 100644 index 0000000..e3b1b36 --- /dev/null +++ b/tests/test_impact.py @@ -0,0 +1,149 @@ +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import classify_task_impact + + +def test_classifies_inherited_delay() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=3.1, + mad=0.2, + robust_z_score=0.3, + deviation_percent=3.3, + severity="NORMAL", + ) + + upstream_handoff_drifts = [ + DriftResult( + task_id="extract->transform", + baseline=2.0, + current=8.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ) + ] + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + + assert result.task_id == "transform" + assert result.classification == "INHERITED_DELAY" + assert result.task_severity == "NORMAL" + assert result.upstream_handoff_severity == "CRITICAL" + + +def test_classifies_own_drift() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=9.0, + mad=0.3, + robust_z_score=6.5, + deviation_percent=200.0, + severity="CRITICAL", + ) + + upstream_handoff_drifts = [ + DriftResult( + task_id="extract->transform", + baseline=2.0, + current=2.1, + mad=0.2, + robust_z_score=0.3, + deviation_percent=5.0, + severity="NORMAL", + ) + ] + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + + assert result.classification == "OWN_DRIFT" + assert result.task_severity == "CRITICAL" + assert result.upstream_handoff_severity is None + + +def test_classifies_normal_task() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=3.1, + mad=0.2, + robust_z_score=0.2, + deviation_percent=3.3, + severity="NORMAL", + ) + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=[], + ) + + assert result.classification == "NORMAL" + assert result.task_severity == "NORMAL" + assert result.upstream_handoff_severity is None + + +def test_classifies_combined_impact() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=8.0, + mad=0.3, + robust_z_score=5.5, + deviation_percent=166.7, + severity="CRITICAL", + ) + + upstream_handoff_drifts = [ + DriftResult( + task_id="extract->transform", + baseline=2.0, + current=7.0, + mad=0.2, + robust_z_score=5.0, + deviation_percent=250.0, + severity="CRITICAL", + ) + ] + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + + assert result.classification == "COMBINED" + assert result.task_severity == "CRITICAL" + assert result.upstream_handoff_severity == "CRITICAL" + + +def test_classifies_medium_task_as_own_drift() -> None: + task_drift = DriftResult( + task_id="load", + baseline=1.4, + current=2.1, + mad=0.2, + robust_z_score=3.17, + deviation_percent=50.0, + severity="MEDIUM", + ) + + result = classify_task_impact( + task_id="load", + task_drift=task_drift, + upstream_handoff_drifts=[], + ) + + assert result.classification == "OWN_DRIFT" + assert result.task_severity == "MEDIUM" + assert result.upstream_handoff_severity is None diff --git a/tests/test_mcp_integration.py b/tests/test_mcp_integration.py new file mode 100644 index 0000000..f045671 --- /dev/null +++ b/tests/test_mcp_integration.py @@ -0,0 +1,54 @@ +import os +import sys + +import pytest +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +@pytest.mark.integration +@pytest.mark.anyio +async def test_mcp_analyze_airflow_dag_end_to_end() -> None: + dag_id = os.getenv( + "FLOWSENSE_TEST_DAG_ID", + "flowsense_demo", + ) + + env = os.environ.copy() + + existing_pythonpath = env.get("PYTHONPATH") + + if existing_pythonpath: + env["PYTHONPATH"] = f"src:{existing_pythonpath}" + else: + env["PYTHONPATH"] = "src" + + server_params = StdioServerParameters( + command=sys.executable, + args=[ + "-m", + "flowsense.mcp.server", + ], + env=env, + ) + + async with ( + stdio_client(server_params) as (read, write), + ClientSession(read, write) as session, + ): + await session.initialize() + + result = await session.call_tool( + "analyze_airflow_dag", + { + "dag_id": dag_id, + }, + ) + + assert result.is_error is False + assert result.content diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py new file mode 100644 index 0000000..bbee8f1 --- /dev/null +++ b/tests/test_mcp_server.py @@ -0,0 +1,136 @@ +import os +import sys + +import pytest +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact +from flowsense.engine.root_cause import RootCauseResult +from flowsense.mcp.server import serialize_analysis +from flowsense.models import DAGAnalysis + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +@pytest.mark.anyio +async def test_mcp_server_exposes_analyze_tool() -> None: + env = os.environ.copy() + + existing_pythonpath = env.get("PYTHONPATH") + + if existing_pythonpath: + env["PYTHONPATH"] = f"src:{existing_pythonpath}" + else: + env["PYTHONPATH"] = "src" + + server_params = StdioServerParameters( + command=sys.executable, + args=[ + "-m", + "flowsense.mcp.server", + ], + env=env, + ) + + async with ( + stdio_client(server_params) as (read, write), + ClientSession(read, write) as session, + ): + await session.initialize() + + response = await session.list_tools() + + tool_names = [tool.name for tool in response.tools] + + assert "analyze_airflow_dag" in tool_names + + +def test_serialize_analysis() -> None: + analysis = DAGAnalysis( + dag_id="demo", + runs_analyzed=5, + overall_severity="CRITICAL", + primary_origin=RootCauseResult( + task_id="transform", + classification="COMBINED", + severity="CRITICAL", + propagation_score=0.8, + ), + drift_results={ + "transform": DriftResult( + task_id="transform", + baseline=3.0, + current=9.5, + mad=0.1, + robust_z_score=43.84, + deviation_percent=216.67, + severity="CRITICAL", + ) + }, + handoff_drift_results={ + ("extract", "transform"): DriftResult( + task_id="extract->transform", + baseline=2.0, + current=8.0, + mad=0.2, + robust_z_score=20.24, + deviation_percent=300.0, + severity="CRITICAL", + ) + }, + task_impacts={ + "transform": TaskImpact( + task_id="transform", + classification="COMBINED", + task_severity="CRITICAL", + upstream_handoff_severity="CRITICAL", + ) + }, + propagation_results=[], + dependencies={ + "transform": [], + }, + ) + + result = serialize_analysis(analysis) + + assert result["dag_id"] == "demo" + assert result["runs_analyzed"] == 5 + assert result["overall_severity"] == "CRITICAL" + + primary_origin = result["primary_origin"] + + assert primary_origin["task_id"] == "transform" + assert primary_origin["classification"] == "COMBINED" + assert primary_origin["severity"] == "CRITICAL" + assert primary_origin["propagation_score"] == 0.8 + transform = result["drift_results"]["transform"] + + assert transform["baseline"] == 3.0 + assert transform["current"] == 9.5 + assert transform["mad"] == 0.1 + assert transform["robust_z_score"] == 43.84 + assert transform["deviation_percent"] == 216.67 + assert transform["severity"] == "CRITICAL" + + handoff = result["handoff_drift_results"]["extract->transform"] + + assert handoff["baseline"] == 2.0 + assert handoff["current"] == 8.0 + assert handoff["severity"] == "CRITICAL" + + impact = result["task_impacts"]["transform"] + + assert impact["classification"] == "COMBINED" + assert impact["task_severity"] == "CRITICAL" + assert impact["upstream_handoff_severity"] == "CRITICAL" + + assert result["propagation_results"] == [] + assert result["dependencies"] == { + "transform": [], + } diff --git a/tests/test_propagation.py b/tests/test_propagation.py index 30bddc5..c44ac98 100644 --- a/tests/test_propagation.py +++ b/tests/test_propagation.py @@ -122,3 +122,246 @@ def test_medium_task_is_not_treated_as_origin() -> None: ) assert results == [] + + +def test_detects_multi_hop_propagation() -> None: + drift_results = { + "extract": DriftResult( + task_id="extract", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "transform": DriftResult( + task_id="transform", + baseline=3.0, + current=6.0, + mad=0.5, + robust_z_score=4.0, + deviation_percent=100.0, + severity="HIGH", + ), + "load": DriftResult( + task_id="load", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + } + + dependencies = { + "extract": ["transform"], + "transform": ["load"], + "load": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + extract_result = next( + result for result in results if result.origin_task == "extract" + ) + + assert extract_result.path == [ + "extract", + "transform", + "load", + ] + assert extract_result.affected_tasks == [ + "transform", + "load", + ] + assert extract_result.propagation_score > 0 + + +def test_does_not_report_downstream_task_as_duplicate_origin() -> None: + drift_results = { + "extract": DriftResult( + task_id="extract", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "transform": DriftResult( + task_id="transform", + baseline=3.0, + current=6.0, + mad=0.5, + robust_z_score=4.0, + deviation_percent=100.0, + severity="HIGH", + ), + "load": DriftResult( + task_id="load", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + } + + dependencies = { + "extract": ["transform"], + "transform": ["load"], + "load": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + origins = [result.origin_task for result in results] + + assert origins == ["extract"] + + +def test_detects_branching_propagation_paths() -> None: + drift_results = { + "extract": DriftResult( + task_id="extract", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "transform_a": DriftResult( + task_id="transform_a", + baseline=2.0, + current=5.0, + mad=0.3, + robust_z_score=4.5, + deviation_percent=150.0, + severity="HIGH", + ), + "load_a": DriftResult( + task_id="load_a", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + "transform_b": DriftResult( + task_id="transform_b", + baseline=2.0, + current=4.0, + mad=0.3, + robust_z_score=3.8, + deviation_percent=100.0, + severity="HIGH", + ), + "load_b": DriftResult( + task_id="load_b", + baseline=1.0, + current=1.6, + mad=0.2, + robust_z_score=2.2, + deviation_percent=60.0, + severity="MEDIUM", + ), + } + + dependencies = { + "extract": ["transform_a", "transform_b"], + "transform_a": ["load_a"], + "load_a": [], + "transform_b": ["load_b"], + "load_b": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + extract_results = [result for result in results if result.origin_task == "extract"] + + paths = [result.path for result in extract_results] + + assert len(extract_results) == 2 + + assert [ + "extract", + "transform_a", + "load_a", + ] in paths + + assert [ + "extract", + "transform_b", + "load_b", + ] in paths + + +def test_preserves_origin_after_normal_dependency_gap() -> None: + drift_results = { + "upstream": DriftResult( + task_id="upstream", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "normal_bridge": DriftResult( + task_id="normal_bridge", + baseline=2.0, + current=2.1, + mad=0.3, + robust_z_score=0.2, + deviation_percent=5.0, + severity="NORMAL", + ), + "independent_origin": DriftResult( + task_id="independent_origin", + baseline=3.0, + current=7.0, + mad=0.5, + robust_z_score=5.4, + deviation_percent=133.3, + severity="CRITICAL", + ), + "downstream": DriftResult( + task_id="downstream", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + } + + dependencies = { + "upstream": ["normal_bridge"], + "normal_bridge": ["independent_origin"], + "independent_origin": ["downstream"], + "downstream": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + assert len(results) == 1 + assert results[0].origin_task == "independent_origin" + assert results[0].path == ["independent_origin", "downstream"] diff --git a/tests/test_root_cause.py b/tests/test_root_cause.py new file mode 100644 index 0000000..7b2faa7 --- /dev/null +++ b/tests/test_root_cause.py @@ -0,0 +1,232 @@ +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact +from flowsense.engine.propagation import PropagationResult +from flowsense.engine.root_cause import select_primary_origin + + +def _drift( + task_id: str, + severity: str, +) -> DriftResult: + return DriftResult( + task_id=task_id, + baseline=1.0, + current=2.0, + mad=0.1, + robust_z_score=3.0, + deviation_percent=100.0, + severity=severity, + ) + + +def _impact( + task_id: str, + classification: str, + severity: str, +) -> TaskImpact: + return TaskImpact( + task_id=task_id, + classification=classification, + task_severity=severity, + upstream_handoff_severity=None, + ) + + +def test_selects_single_own_drift_as_primary_origin() -> None: + drift_results = { + "transform": _drift( + "transform", + "CRITICAL", + ), + } + + task_impacts = { + "transform": _impact( + "transform", + "OWN_DRIFT", + "CRITICAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "transform": [], + }, + propagation_results=[], + ) + + assert result is not None + assert result.task_id == "transform" + + +def test_inherited_delay_is_not_primary_origin() -> None: + drift_results = { + "transform": _drift( + "transform", + "NORMAL", + ), + } + + task_impacts = { + "transform": _impact( + "transform", + "INHERITED_DELAY", + "NORMAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "transform": [], + }, + propagation_results=[], + ) + + assert result is None + + +def test_prefers_upstream_root_over_downstream_combined_task() -> None: + drift_results = { + "extract": _drift( + "extract", + "CRITICAL", + ), + "transform": _drift( + "transform", + "CRITICAL", + ), + } + + task_impacts = { + "extract": _impact( + "extract", + "OWN_DRIFT", + "CRITICAL", + ), + "transform": _impact( + "transform", + "COMBINED", + "CRITICAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "extract": ["transform"], + "transform": [], + }, + propagation_results=[], + ) + + assert result is not None + assert result.task_id == "extract" + + +def test_selects_more_severe_independent_root() -> None: + drift_results = { + "extract_a": _drift( + "extract_a", + "MEDIUM", + ), + "extract_b": _drift( + "extract_b", + "CRITICAL", + ), + } + + task_impacts = { + "extract_a": _impact( + "extract_a", + "OWN_DRIFT", + "MEDIUM", + ), + "extract_b": _impact( + "extract_b", + "OWN_DRIFT", + "CRITICAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "extract_a": [], + "extract_b": [], + }, + propagation_results=[], + ) + + assert result is not None + assert result.task_id == "extract_b" + + +def test_prefers_higher_propagation_score_between_independent_roots() -> None: + drift_results = { + "extract_a": _drift( + "extract_a", + "CRITICAL", + ), + "extract_b": _drift( + "extract_b", + "CRITICAL", + ), + } + + task_impacts = { + "extract_a": _impact( + "extract_a", + "OWN_DRIFT", + "CRITICAL", + ), + "extract_b": _impact( + "extract_b", + "OWN_DRIFT", + "CRITICAL", + ), + } + + propagation_results = [ + PropagationResult( + origin_task="extract_a", + affected_tasks=["load_a"], + path=["extract_a", "load_a"], + propagation_score=0.4, + ), + PropagationResult( + origin_task="extract_b", + affected_tasks=[ + "transform_b", + "load_b", + ], + path=[ + "extract_b", + "transform_b", + "load_b", + ], + propagation_score=0.8, + ), + ] + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "extract_a": ["load_a"], + "load_a": [], + "extract_b": ["transform_b"], + "transform_b": ["load_b"], + "load_b": [], + }, + propagation_results=propagation_results, + ) + + assert result is not None + assert result.task_id == "extract_b" + assert result.propagation_score == 0.8 diff --git a/tests/test_timing.py b/tests/test_timing.py new file mode 100644 index 0000000..c694854 --- /dev/null +++ b/tests/test_timing.py @@ -0,0 +1,162 @@ +from datetime import UTC, datetime, timedelta + +import pytest + +from flowsense.engine.timing import ( + build_handoff_history, + calculate_handoff_delay, + calculate_handoff_drift, +) +from flowsense.models import TaskRun + + +def test_calculate_handoff_delay() -> None: + upstream_end = datetime(2026, 8, 20, 10, 0, 5, tzinfo=UTC) + + upstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + state="success", + end_date=upstream_end, + ) + + downstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + start_date=upstream_end + timedelta(seconds=7), + ) + + result = calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + + assert result.upstream_task == "extract" + assert result.downstream_task == "transform" + assert result.dag_run_id == "run_1" + assert result.handoff_delay == 7.0 + + +def test_calculate_handoff_delay_requires_upstream_end_date() -> None: + upstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + ) + + downstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + start_date=datetime(2026, 8, 20, 10, 0, tzinfo=UTC), + ) + + with pytest.raises(ValueError): + calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + + +def test_calculate_handoff_delay_requires_same_dag_run() -> None: + upstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + end_date=datetime(2026, 8, 20, 10, 0, tzinfo=UTC), + ) + + downstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_2", + task_id="transform", + start_date=datetime(2026, 8, 20, 10, 0, 5, tzinfo=UTC), + ) + + with pytest.raises(ValueError): + calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + + +def test_build_handoff_history() -> None: + base_time = datetime( + 2026, + 8, + 20, + 10, + 0, + tzinfo=UTC, + ) + + task_runs = [] + + for index, delay in enumerate( + [2.0, 3.0, 4.0], + start=1, + ): + dag_run_id = f"run_{index}" + + upstream_end = base_time + timedelta( + minutes=index, + ) + + task_runs.extend( + [ + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="extract", + state="success", + end_date=upstream_end, + ), + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="transform", + state="success", + start_date=upstream_end + timedelta(seconds=delay), + ), + ] + ) + + dependencies = { + "extract": ["transform"], + "transform": [], + } + + history = build_handoff_history( + task_runs=task_runs, + dependencies=dependencies, + ) + + assert history == { + ("extract", "transform"): [ + 2.0, + 3.0, + 4.0, + ] + } + + +def test_calculate_handoff_drift_detects_anomaly() -> None: + result = calculate_handoff_drift( + upstream_task="extract", + downstream_task="transform", + handoff_delays=[ + 2.0, + 2.1, + 1.9, + 2.0, + 8.0, + ], + ) + + assert result.task_id == "extract->transform" + assert result.severity == "CRITICAL" + assert result.current == 8.0 + assert result.baseline == 2.0 diff --git a/uv.lock b/uv.lock index e0a7539..64a7413 100644 --- a/uv.lock +++ b/uv.lock @@ -235,13 +235,13 @@ dev = [ { name = "ruff" }, ] mcp = [ - { name = "mcp" }, + { name = "mcp", extra = ["cli"] }, ] [package.metadata] requires-dist = [ { name = "httpx", specifier = ">=0.27" }, - { name = "mcp", marker = "extra == 'mcp'" }, + { name = "mcp", extras = ["cli"], marker = "extra == 'mcp'", specifier = ">=2.0" }, { name = "numpy", specifier = ">=2.0" }, { name = "pandas", specifier = ">=2.2" }, { name = "pydantic", specifier = ">=2.8" }, @@ -411,6 +411,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/67/72/7d7897418912c1d12e87556630dfb7bf0eac71160e9bef8b447960804ee3/mcp-2.0.0-py3-none-any.whl", hash = "sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6", size = 349980, upload-time = "2026-07-28T13:45:28.853Z" }, ] +[package.optional-dependencies] +cli = [ + { name = "python-dotenv" }, + { name = "typer" }, +] + [[package]] name = "mcp-types" version = "2.0.0" @@ -732,6 +738,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/53/ed9d74092561d4b01a2ef1349d52cdbc135e526c245f366b089cfca6de49/python_dotenv-1.2.3.tar.gz", hash = "sha256:a20a594dabeaa385725aa239d5244871c143ecb356add8a20fcf23773a6c3a35", size = 58945, upload-time = "2026-08-16T16:54:54.067Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/17/c5c6b53ddc18f297992099b3d9ec16c855c0ccc83263a21fe4d1c625ec6c/python_dotenv-1.2.3-py3-none-any.whl", hash = "sha256:904552145e8bfed22162c09dab1c2b9b54fefa7b23ba780f4f26ca0316b0f0d9", size = 22780, upload-time = "2026-08-16T16:54:52.473Z" }, +] + [[package]] name = "python-multipart" version = "0.0.32"