From 8930af47a8e17c2ec46ca05e7a1be61183d8de8f Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 22:47:44 +0300 Subject: [PATCH 01/16] fix: detect drift when baseline MAD is zero --- src/flowsense/engine/drift.py | 6 +++++- tests/test_drift.py | 22 ++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/flowsense/engine/drift.py b/src/flowsense/engine/drift.py index 7dce2b9..e1e33e7 100644 --- a/src/flowsense/engine/drift.py +++ b/src/flowsense/engine/drift.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math from dataclasses import dataclass import numpy as np @@ -37,7 +38,10 @@ def calculate_drift( mad = float(np.median(absolute_deviations)) if mad == 0: - robust_z_score = 0.0 + if math.isclose(current, median): + robust_z_score = 0.0 + else: + robust_z_score = math.copysign(5.0, current - median) else: robust_z_score = 0.6745 * (current - median) / mad diff --git a/tests/test_drift.py b/tests/test_drift.py index c176921..5495d50 100644 --- a/tests/test_drift.py +++ b/tests/test_drift.py @@ -41,6 +41,28 @@ def test_calculate_drift_critical() -> None: assert result.deviation_percent > 100 +def test_calculate_drift_detects_change_when_mad_is_zero() -> None: + result = calculate_drift( + "transform", + [3.0, 3.0, 3.0, 3.0, 9.0], + ) + + assert result.mad == 0.0 + assert result.robust_z_score == 5.0 + assert result.severity == "CRITICAL" + + +def test_calculate_drift_remains_normal_when_mad_and_change_are_zero() -> None: + result = calculate_drift( + "transform", + [3.0, 3.0, 3.0, 3.0, 3.0], + ) + + assert result.mad == 0.0 + assert result.robust_z_score == 0.0 + assert result.severity == "NORMAL" + + def test_calculate_drift_requires_minimum_history() -> None: durations = [ 3.0, From 76d10ca474c483bcf5d71453f0e7b8db66190aa5 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 22:53:07 +0300 Subject: [PATCH 02/16] feat: add pagination support for Airflow API --- src/flowsense/collector/airflow_client.py | 80 ++++++++++----- tests/test_airflow_client.py | 117 ++++++++++++++++++++++ 2 files changed, 173 insertions(+), 24 deletions(-) create mode 100644 tests/test_airflow_client.py diff --git a/src/flowsense/collector/airflow_client.py b/src/flowsense/collector/airflow_client.py index f98a45e..33d7061 100644 --- a/src/flowsense/collector/airflow_client.py +++ b/src/flowsense/collector/airflow_client.py @@ -5,6 +5,8 @@ from flowsense.config import get_airflow_config from flowsense.models import TaskRun +PAGE_SIZE = 100 + class AirflowClient: def __init__( @@ -50,19 +52,59 @@ def _headers(self) -> dict[str, str]: "Accept": "application/json", } + def _get_paginated( + self, + url: str, + collection_key: str, + ) -> dict: + items: list[dict] = [] + offset = 0 + last_page: dict = {} + + while True: + response = httpx.get( + url, + headers=self._headers(), + params={ + "limit": PAGE_SIZE, + "offset": offset, + }, + timeout=10.0, + ) + + response.raise_for_status() + + last_page = response.json() + page_items = last_page[collection_key] + items.extend(page_items) + + total_entries = last_page.get("total_entries") + + if not page_items: + break + + if total_entries is not None and len(items) >= total_entries: + break + + if total_entries is None and len(page_items) < PAGE_SIZE: + break + + offset += len(page_items) + + return { + **last_page, + collection_key: items, + "total_entries": last_page.get("total_entries", len(items)), + } + def get_dag_runs(self, dag_id: str) -> dict: url = f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns" - response = httpx.get( - url, - headers=self._headers(), - timeout=10.0, + return self._get_paginated( + url=url, + collection_key="dag_runs", ) - response.raise_for_status() - - return response.json() - def get_task_instances( self, dag_id: str, @@ -70,16 +112,11 @@ def get_task_instances( ) -> dict: url = f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances" - response = httpx.get( - url, - headers=self._headers(), - timeout=10.0, + return self._get_paginated( + url=url, + collection_key="task_instances", ) - response.raise_for_status() - - return response.json() - def collect_task_runs( self, dag_id: str, @@ -130,16 +167,11 @@ def get_dag_tasks( ) -> dict: url = f"{self.base_url}/api/v2/dags/{dag_id}/tasks" - response = httpx.get( - url, - headers=self._headers(), - timeout=10.0, + return self._get_paginated( + url=url, + collection_key="tasks", ) - response.raise_for_status() - - return response.json() - def get_dag_dependencies( self, dag_id: str, diff --git a/tests/test_airflow_client.py b/tests/test_airflow_client.py new file mode 100644 index 0000000..112b1e0 --- /dev/null +++ b/tests/test_airflow_client.py @@ -0,0 +1,117 @@ +from unittest.mock import MagicMock, call, patch + +import pytest + +from flowsense.collector.airflow_client import PAGE_SIZE, AirflowClient +from flowsense.config import AirflowConfig + + +@pytest.fixture +def client() -> AirflowClient: + with patch( + "flowsense.collector.airflow_client.get_airflow_config", + return_value=AirflowConfig( + base_url="http://airflow.test", + username="airflow", + password="airflow", + ), + ): + airflow_client = AirflowClient() + + airflow_client._token = "token" + return airflow_client + + +@patch("flowsense.collector.airflow_client.httpx.get") +def test_get_dag_runs_collects_all_pages( + mock_get: MagicMock, + client: AirflowClient, +) -> None: + first_page = MagicMock() + first_page.json.return_value = { + "dag_runs": [{"dag_run_id": f"run_{index}"} for index in range(PAGE_SIZE)], + "total_entries": PAGE_SIZE + 1, + } + + second_page = MagicMock() + second_page.json.return_value = { + "dag_runs": [{"dag_run_id": f"run_{PAGE_SIZE}"}], + "total_entries": PAGE_SIZE + 1, + } + + mock_get.side_effect = [first_page, second_page] + + result = client.get_dag_runs("demo") + + assert len(result["dag_runs"]) == PAGE_SIZE + 1 + assert result["total_entries"] == PAGE_SIZE + 1 + assert mock_get.call_args_list == [ + call( + "http://airflow.test/api/v2/dags/demo/dagRuns", + headers={ + "Authorization": "Bearer token", + "Accept": "application/json", + }, + params={"limit": PAGE_SIZE, "offset": 0}, + timeout=10.0, + ), + call( + "http://airflow.test/api/v2/dags/demo/dagRuns", + headers={ + "Authorization": "Bearer token", + "Accept": "application/json", + }, + params={"limit": PAGE_SIZE, "offset": PAGE_SIZE}, + timeout=10.0, + ), + ] + + first_page.raise_for_status.assert_called_once_with() + second_page.raise_for_status.assert_called_once_with() + + +@pytest.mark.parametrize( + ("method_name", "collection_key", "expected_url", "args"), + [ + ( + "get_task_instances", + "task_instances", + "http://airflow.test/api/v2/dags/demo/dagRuns/run_1/taskInstances", + ("demo", "run_1"), + ), + ( + "get_dag_tasks", + "tasks", + "http://airflow.test/api/v2/dags/demo/tasks", + ("demo",), + ), + ], +) +@patch("flowsense.collector.airflow_client.httpx.get") +def test_paginated_endpoints_use_their_collection_key( + mock_get: MagicMock, + method_name: str, + collection_key: str, + expected_url: str, + args: tuple[str, ...], + client: AirflowClient, +) -> None: + response = MagicMock() + response.json.return_value = { + collection_key: [{"id": "item_1"}], + "total_entries": 1, + } + mock_get.return_value = response + + result = getattr(client, method_name)(*args) + + assert result[collection_key] == [{"id": "item_1"}] + mock_get.assert_called_once_with( + expected_url, + headers={ + "Authorization": "Bearer token", + "Accept": "application/json", + }, + params={"limit": PAGE_SIZE, "offset": 0}, + timeout=10.0, + ) From 24d246099f8739944f5e4282f49ea4a1e34a1fd2 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 22:54:52 +0300 Subject: [PATCH 03/16] fix: aggregate dynamically mapped task runs --- src/flowsense/engine/history.py | 18 ++++++++++-- src/flowsense/engine/timing.py | 38 ++++++++++++++++++++----- tests/test_history.py | 41 +++++++++++++++++++++++++++ tests/test_timing.py | 49 +++++++++++++++++++++++++++++++++ 4 files changed, 136 insertions(+), 10 deletions(-) diff --git a/src/flowsense/engine/history.py b/src/flowsense/engine/history.py index 4b5aeba..0e56d5f 100644 --- a/src/flowsense/engine/history.py +++ b/src/flowsense/engine/history.py @@ -8,10 +8,22 @@ def build_duration_history( task_runs: list[TaskRun], ) -> dict[str, list[float]]: - history: dict[str, list[float]] = defaultdict(list) + """Build logical-task history using the slowest mapped instance per DAG run.""" + durations_by_run_and_task: dict[tuple[str, str], float] = {} for task_run in task_runs: - if task_run.state == "success" and task_run.duration is not None: - history[task_run.task_id].append(task_run.duration) + if task_run.state != "success" or task_run.duration is None: + continue + + key = (task_run.dag_run_id, task_run.task_id) + current_duration = durations_by_run_and_task.get(key) + + if current_duration is None or task_run.duration > current_duration: + durations_by_run_and_task[key] = task_run.duration + + history: dict[str, list[float]] = defaultdict(list) + + for (_dag_run_id, task_id), duration in durations_by_run_and_task.items(): + history[task_id].append(duration) return dict(history) diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py index 7c43cac..edf942f 100644 --- a/src/flowsense/engine/timing.py +++ b/src/flowsense/engine/timing.py @@ -41,26 +41,50 @@ 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]] = {} + """Build logical-edge history across regular and dynamically mapped tasks. + + A mapped upstream is complete at its latest instance end, while a mapped + downstream starts at its earliest instance start. + """ + runs_by_id: dict[str, dict[str, list[TaskRun]]] = {} for task_run in task_runs: - runs_by_id.setdefault(task_run.dag_run_id, {})[task_run.task_id] = task_run + tasks_in_run = runs_by_id.setdefault(task_run.dag_run_id, {}) + tasks_in_run.setdefault(task_run.task_id, []).append(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) + upstream_runs = tasks_in_run.get(upstream_task, []) + upstream_runs_with_end = [ + task_run for task_run in upstream_runs if task_run.end_date is not None + ] - if upstream_run is None: + if not upstream_runs_with_end: continue - for downstream_task in downstream_tasks: - downstream_run = tasks_in_run.get(downstream_task) + upstream_run = max( + upstream_runs_with_end, + key=lambda task_run: task_run.end_date, + ) - if downstream_run is None: + for downstream_task in downstream_tasks: + downstream_runs = tasks_in_run.get(downstream_task, []) + downstream_runs_with_start = [ + task_run + for task_run in downstream_runs + if task_run.start_date is not None + ] + + if not downstream_runs_with_start: continue + downstream_run = min( + downstream_runs_with_start, + key=lambda task_run: task_run.start_date, + ) + try: timing = calculate_handoff_delay( upstream_run=upstream_run, diff --git a/tests/test_history.py b/tests/test_history.py index a8561e3..c9ccab6 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -85,3 +85,44 @@ def test_build_duration_history_ignores_missing_duration() -> None: history = build_duration_history(task_runs) assert history == {} + + +def test_build_duration_history_uses_slowest_mapped_instance_per_run() -> None: + task_runs = [ + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + duration=2.0, + map_index=0, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + duration=5.0, + map_index=1, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_2", + task_id="transform", + state="success", + duration=4.0, + map_index=0, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_2", + task_id="transform", + state="success", + duration=3.0, + map_index=1, + ), + ] + + history = build_duration_history(task_runs) + + assert history == {"transform": [5.0, 4.0]} diff --git a/tests/test_timing.py b/tests/test_timing.py index c694854..d629861 100644 --- a/tests/test_timing.py +++ b/tests/test_timing.py @@ -143,6 +143,55 @@ def test_build_handoff_history() -> None: } +def test_build_handoff_history_aggregates_mapped_task_boundaries() -> None: + base_time = datetime(2026, 8, 20, 10, 0, tzinfo=UTC) + + task_runs = [ + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + state="success", + end_date=base_time + timedelta(seconds=10), + map_index=0, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + state="success", + end_date=base_time + timedelta(seconds=5), + map_index=1, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + start_date=base_time + timedelta(seconds=12), + map_index=0, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + start_date=base_time + timedelta(seconds=20), + map_index=1, + ), + ] + + history = build_handoff_history( + task_runs=task_runs, + dependencies={ + "extract": ["transform"], + "transform": [], + }, + ) + + assert history == {("extract", "transform"): [2.0]} + + def test_calculate_handoff_drift_detects_anomaly() -> None: result = calculate_handoff_drift( upstream_task="extract", From 428137de1b1c55ec08d581168814c0c377ab092e Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 22:56:20 +0300 Subject: [PATCH 04/16] fix: align root cause and propagation chain rules --- src/flowsense/engine/root_cause.py | 8 +++++ tests/test_root_cause.py | 49 ++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/src/flowsense/engine/root_cause.py b/src/flowsense/engine/root_cause.py index 4e2e03f..a018222 100644 --- a/src/flowsense/engine/root_cause.py +++ b/src/flowsense/engine/root_cause.py @@ -43,6 +43,7 @@ def _build_reverse_dependencies( def _has_candidate_upstream( task_id: str, candidate_tasks: set[str], + drift_results: dict[str, DriftResult], reverse_dependencies: dict[str, list[str]], visited: set[str] | None = None, ) -> bool: @@ -55,12 +56,18 @@ def _has_candidate_upstream( visited.add(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_task in candidate_tasks: return True if _has_candidate_upstream( task_id=upstream_task, candidate_tasks=candidate_tasks, + drift_results=drift_results, reverse_dependencies=reverse_dependencies, visited=visited, ): @@ -98,6 +105,7 @@ def select_primary_origin( if not _has_candidate_upstream( task_id=task_id, candidate_tasks=candidate_tasks, + drift_results=drift_results, reverse_dependencies=reverse_dependencies, ) } diff --git a/tests/test_root_cause.py b/tests/test_root_cause.py index 7b2faa7..0b43010 100644 --- a/tests/test_root_cause.py +++ b/tests/test_root_cause.py @@ -230,3 +230,52 @@ def test_prefers_higher_propagation_score_between_independent_roots() -> None: assert result is not None assert result.task_id == "extract_b" assert result.propagation_score == 0.8 + + +def test_treats_candidate_after_normal_dependency_gap_as_independent() -> None: + drift_results = { + "upstream": _drift( + "upstream", + "HIGH", + ), + "normal_bridge": _drift( + "normal_bridge", + "NORMAL", + ), + "independent_origin": _drift( + "independent_origin", + "CRITICAL", + ), + } + + task_impacts = { + "upstream": _impact( + "upstream", + "OWN_DRIFT", + "HIGH", + ), + "normal_bridge": _impact( + "normal_bridge", + "NORMAL", + "NORMAL", + ), + "independent_origin": _impact( + "independent_origin", + "OWN_DRIFT", + "CRITICAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "upstream": ["normal_bridge"], + "normal_bridge": ["independent_origin"], + "independent_origin": [], + }, + propagation_results=[], + ) + + assert result is not None + assert result.task_id == "independent_origin" From 68ef6d96a3870c58868df7ee08dc807cc5188a96 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 22:58:10 +0300 Subject: [PATCH 05/16] feat: report skipped analysis diagnostics --- src/flowsense/cli/main.py | 8 +++++ src/flowsense/engine/analyzer.py | 24 ++++++++++--- src/flowsense/mcp/server.py | 8 +++++ src/flowsense/models/__init__.py | 3 +- src/flowsense/models/dag_analysis.py | 10 +++++- tests/test_analyzer.py | 52 ++++++++++++++++++++++++++++ tests/test_mcp_server.py | 16 ++++++++- 7 files changed, 113 insertions(+), 8 deletions(-) diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index 33d5604..585fb99 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -74,3 +74,11 @@ def analyze( console.print(f"Origin: {result.origin_task}") console.print(f"Path: {' -> '.join(result.path)}") console.print(f"Propagation Score: {result.propagation_score:.2f}") + + if analysis.diagnostics: + console.print("\n[bold yellow]Diagnostics[/bold yellow]\n") + + for diagnostic in analysis.diagnostics: + console.print( + f"[{diagnostic.code}] {diagnostic.subject_id}: {diagnostic.message}" + ) diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py index aa45c3c..56d1faf 100644 --- a/src/flowsense/engine/analyzer.py +++ b/src/flowsense/engine/analyzer.py @@ -10,7 +10,7 @@ build_handoff_history, calculate_handoff_drift, ) -from flowsense.models import DAGAnalysis +from flowsense.models import AnalysisDiagnostic, DAGAnalysis def analyze_dag(dag_id: str) -> DAGAnalysis: @@ -18,6 +18,7 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: task_runs = client.collect_task_runs(dag_id) duration_history = build_duration_history(task_runs) + diagnostics: list[AnalysisDiagnostic] = [] drift_results = {} @@ -27,8 +28,14 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: task_id=task_id, durations=durations, ) - except ValueError: - continue + except ValueError as exc: + diagnostics.append( + AnalysisDiagnostic( + code="INSUFFICIENT_TASK_HISTORY", + subject_id=task_id, + message=str(exc), + ) + ) dependencies = client.get_dag_dependencies(dag_id) @@ -48,8 +55,14 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: downstream_task=downstream_task, handoff_delays=delays, ) - except ValueError: - continue + except ValueError as exc: + diagnostics.append( + AnalysisDiagnostic( + code="INSUFFICIENT_HANDOFF_HISTORY", + subject_id=f"{_upstream_task}->{downstream_task}", + message=str(exc), + ) + ) task_impacts = {} @@ -111,4 +124,5 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: task_impacts=task_impacts, propagation_results=propagation_results, dependencies=dependencies, + diagnostics=diagnostics, ) diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py index af8d483..620d538 100644 --- a/src/flowsense/mcp/server.py +++ b/src/flowsense/mcp/server.py @@ -66,6 +66,14 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: for result in analysis.propagation_results ], "dependencies": analysis.dependencies, + "diagnostics": [ + { + "code": diagnostic.code, + "subject_id": diagnostic.subject_id, + "message": diagnostic.message, + } + for diagnostic in analysis.diagnostics + ], } diff --git a/src/flowsense/models/__init__.py b/src/flowsense/models/__init__.py index 133820d..0c373d7 100644 --- a/src/flowsense/models/__init__.py +++ b/src/flowsense/models/__init__.py @@ -1,7 +1,8 @@ -from flowsense.models.dag_analysis import DAGAnalysis +from flowsense.models.dag_analysis import AnalysisDiagnostic, DAGAnalysis from flowsense.models.task_run import TaskRun __all__ = [ + "AnalysisDiagnostic", "DAGAnalysis", "TaskRun", ] diff --git a/src/flowsense/models/dag_analysis.py b/src/flowsense/models/dag_analysis.py index e963ffb..71bed5a 100644 --- a/src/flowsense/models/dag_analysis.py +++ b/src/flowsense/models/dag_analysis.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from flowsense.engine.drift import DriftResult from flowsense.engine.impact import TaskImpact @@ -8,6 +8,13 @@ from flowsense.engine.root_cause import RootCauseResult +@dataclass(frozen=True) +class AnalysisDiagnostic: + code: str + subject_id: str + message: str + + @dataclass class DAGAnalysis: dag_id: str @@ -19,3 +26,4 @@ class DAGAnalysis: task_impacts: dict[str, TaskImpact] propagation_results: list[PropagationResult] dependencies: dict[str, list[str]] + diagnostics: list[AnalysisDiagnostic] = field(default_factory=list) diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index f76d9b8..63930b6 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -240,3 +240,55 @@ def test_analyze_dag_uses_handoff_for_overall_severity( ) assert analysis.overall_severity == "CRITICAL" + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_reports_insufficient_history_diagnostics( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + base_time = datetime(2026, 8, 20, 10, 0, tzinfo=UTC) + task_runs = [] + + for index in range(4): + dag_run_id = f"run_{index}" + extract_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=extract_end, + duration=1.0, + ), + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="transform", + state="success", + start_date=extract_end + timedelta(seconds=2), + duration=3.0, + ), + ] + ) + + client.collect_task_runs.return_value = task_runs + client.get_dag_dependencies.return_value = { + "extract": ["transform"], + "transform": [], + } + + analysis = analyze_dag("demo") + + diagnostics = { + (diagnostic.code, diagnostic.subject_id) for diagnostic in analysis.diagnostics + } + + assert diagnostics == { + ("INSUFFICIENT_TASK_HISTORY", "extract"), + ("INSUFFICIENT_TASK_HISTORY", "transform"), + ("INSUFFICIENT_HANDOFF_HISTORY", "extract->transform"), + } diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index bbee8f1..a01aa1e 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -9,7 +9,7 @@ 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 +from flowsense.models import AnalysisDiagnostic, DAGAnalysis @pytest.fixture @@ -95,6 +95,13 @@ def test_serialize_analysis() -> None: dependencies={ "transform": [], }, + diagnostics=[ + AnalysisDiagnostic( + code="INSUFFICIENT_TASK_HISTORY", + subject_id="load", + message="load requires at least 5 runs.", + ) + ], ) result = serialize_analysis(analysis) @@ -134,3 +141,10 @@ def test_serialize_analysis() -> None: assert result["dependencies"] == { "transform": [], } + assert result["diagnostics"] == [ + { + "code": "INSUFFICIENT_TASK_HISTORY", + "subject_id": "load", + "message": "load requires at least 5 runs.", + } + ] From 936e8e82cb4578fdb19e153f73700f9914bbdc30 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:01:55 +0300 Subject: [PATCH 06/16] feat: report invalid handoff timing diagnostics --- src/flowsense/engine/analyzer.py | 14 +++- src/flowsense/engine/timing.py | 114 +++++++++++++++++++++++++++---- tests/test_analyzer.py | 35 ++++++++++ tests/test_timing.py | 48 +++++++++++++ 4 files changed, 194 insertions(+), 17 deletions(-) diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py index 56d1faf..2050d52 100644 --- a/src/flowsense/engine/analyzer.py +++ b/src/flowsense/engine/analyzer.py @@ -7,7 +7,7 @@ from flowsense.engine.propagation import analyze_propagation from flowsense.engine.root_cause import select_primary_origin from flowsense.engine.timing import ( - build_handoff_history, + build_handoff_history_with_diagnostics, calculate_handoff_drift, ) from flowsense.models import AnalysisDiagnostic, DAGAnalysis @@ -39,10 +39,20 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: dependencies = client.get_dag_dependencies(dag_id) - handoff_history = build_handoff_history( + handoff_history_result = build_handoff_history_with_diagnostics( task_runs=task_runs, dependencies=dependencies, ) + handoff_history = handoff_history_result.history + + diagnostics.extend( + AnalysisDiagnostic( + code=diagnostic.code, + subject_id=(f"{diagnostic.upstream_task}->{diagnostic.downstream_task}"), + message=diagnostic.message, + ) + for diagnostic in handoff_history_result.diagnostics + ) handoff_drift_results = {} diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py index edf942f..0cc842b 100644 --- a/src/flowsense/engine/timing.py +++ b/src/flowsense/engine/timing.py @@ -14,6 +14,21 @@ class HandoffTiming: handoff_delay: float +@dataclass(frozen=True) +class HandoffHistoryDiagnostic: + code: str + upstream_task: str + downstream_task: str + dag_run_id: str + message: str + + +@dataclass +class HandoffHistoryResult: + history: dict[tuple[str, str], list[float]] + diagnostics: list[HandoffHistoryDiagnostic] + + def calculate_handoff_delay( upstream_run: TaskRun, downstream_run: TaskRun, @@ -41,6 +56,16 @@ def build_handoff_history( task_runs: list[TaskRun], dependencies: dict[str, list[str]], ) -> dict[tuple[str, str], list[float]]: + return build_handoff_history_with_diagnostics( + task_runs=task_runs, + dependencies=dependencies, + ).history + + +def build_handoff_history_with_diagnostics( + task_runs: list[TaskRun], + dependencies: dict[str, list[str]], +) -> HandoffHistoryResult: """Build logical-edge history across regular and dynamically mapped tasks. A mapped upstream is complete at its latest instance end, while a mapped @@ -53,35 +78,82 @@ def build_handoff_history( tasks_in_run.setdefault(task_run.task_id, []).append(task_run) history: dict[tuple[str, str], list[float]] = {} + diagnostics: list[HandoffHistoryDiagnostic] = [] - for tasks_in_run in runs_by_id.values(): + for dag_run_id, tasks_in_run in runs_by_id.items(): for upstream_task, downstream_tasks in dependencies.items(): upstream_runs = tasks_in_run.get(upstream_task, []) - upstream_runs_with_end = [ - task_run for task_run in upstream_runs if task_run.end_date is not None - ] - if not upstream_runs_with_end: + if not upstream_runs: + for downstream_task in downstream_tasks: + diagnostics.append( + HandoffHistoryDiagnostic( + code="MISSING_UPSTREAM_TASK_RUN", + upstream_task=upstream_task, + downstream_task=downstream_task, + dag_run_id=dag_run_id, + message=( + f"{upstream_task} has no task run in {dag_run_id}." + ), + ) + ) + continue + + if any(task_run.end_date is None for task_run in upstream_runs): + for downstream_task in downstream_tasks: + diagnostics.append( + HandoffHistoryDiagnostic( + code="MISSING_UPSTREAM_END_DATE", + upstream_task=upstream_task, + downstream_task=downstream_task, + dag_run_id=dag_run_id, + message=( + f"{upstream_task} has no complete end_date in " + f"{dag_run_id}." + ), + ) + ) continue upstream_run = max( - upstream_runs_with_end, + upstream_runs, key=lambda task_run: task_run.end_date, ) for downstream_task in downstream_tasks: downstream_runs = tasks_in_run.get(downstream_task, []) - downstream_runs_with_start = [ - task_run - for task_run in downstream_runs - if task_run.start_date is not None - ] - if not downstream_runs_with_start: + if not downstream_runs: + diagnostics.append( + HandoffHistoryDiagnostic( + code="MISSING_DOWNSTREAM_TASK_RUN", + upstream_task=upstream_task, + downstream_task=downstream_task, + dag_run_id=dag_run_id, + message=( + f"{downstream_task} has no task run in {dag_run_id}." + ), + ) + ) + continue + + if any(task_run.start_date is None for task_run in downstream_runs): + diagnostics.append( + HandoffHistoryDiagnostic( + code="MISSING_DOWNSTREAM_START_DATE", + upstream_task=upstream_task, + downstream_task=downstream_task, + dag_run_id=dag_run_id, + message=( + f"{downstream_task} has no complete start_date in " + f"{dag_run_id}." + ), + ) + ) continue downstream_run = min( - downstream_runs_with_start, + downstream_runs, key=lambda task_run: task_run.start_date, ) @@ -90,7 +162,16 @@ def build_handoff_history( upstream_run=upstream_run, downstream_run=downstream_run, ) - except ValueError: + except ValueError as exc: + diagnostics.append( + HandoffHistoryDiagnostic( + code="INVALID_HANDOFF_TIMING", + upstream_task=upstream_task, + downstream_task=downstream_task, + dag_run_id=dag_run_id, + message=str(exc), + ) + ) continue edge = ( @@ -100,7 +181,10 @@ def build_handoff_history( history.setdefault(edge, []).append(timing.handoff_delay) - return history + return HandoffHistoryResult( + history=history, + diagnostics=diagnostics, + ) def calculate_handoff_drift( diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index 63930b6..8b5dd41 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -292,3 +292,38 @@ def test_analyze_dag_reports_insufficient_history_diagnostics( ("INSUFFICIENT_TASK_HISTORY", "transform"), ("INSUFFICIENT_HANDOFF_HISTORY", "extract->transform"), } + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_reports_missing_handoff_timestamp( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + client.collect_task_runs.return_value = [ + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + state="success", + end_date=None, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + start_date=datetime(2026, 8, 20, 10, 0, tzinfo=UTC), + ), + ] + client.get_dag_dependencies.return_value = { + "extract": ["transform"], + "transform": [], + } + + analysis = analyze_dag("demo") + + assert any( + diagnostic.code == "MISSING_UPSTREAM_END_DATE" + and diagnostic.subject_id == "extract->transform" + for diagnostic in analysis.diagnostics + ) diff --git a/tests/test_timing.py b/tests/test_timing.py index d629861..b8aa808 100644 --- a/tests/test_timing.py +++ b/tests/test_timing.py @@ -4,6 +4,7 @@ from flowsense.engine.timing import ( build_handoff_history, + build_handoff_history_with_diagnostics, calculate_handoff_delay, calculate_handoff_drift, ) @@ -192,6 +193,53 @@ def test_build_handoff_history_aggregates_mapped_task_boundaries() -> None: assert history == {("extract", "transform"): [2.0]} +def test_build_handoff_history_reports_missing_timestamps() -> None: + task_runs = [ + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + state="success", + end_date=None, + ), + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + start_date=datetime(2026, 8, 20, 10, 0, tzinfo=UTC), + ), + TaskRun( + dag_id="demo", + dag_run_id="run_2", + task_id="extract", + state="success", + end_date=datetime(2026, 8, 20, 11, 0, tzinfo=UTC), + ), + TaskRun( + dag_id="demo", + dag_run_id="run_2", + task_id="transform", + state="success", + start_date=None, + ), + ] + + result = build_handoff_history_with_diagnostics( + task_runs=task_runs, + dependencies={ + "extract": ["transform"], + "transform": [], + }, + ) + + assert result.history == {} + assert [diagnostic.code for diagnostic in result.diagnostics] == [ + "MISSING_UPSTREAM_END_DATE", + "MISSING_DOWNSTREAM_START_DATE", + ] + + def test_calculate_handoff_drift_detects_anomaly() -> None: result = calculate_handoff_drift( upstream_task="extract", From 46c53feb6744b6adf68f15bfa7b400a96e0f7620 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:05:54 +0300 Subject: [PATCH 07/16] refactor: centralize domain models and enums --- src/flowsense/collector/airflow_client.py | 2 +- src/flowsense/domain/__init__.py | 22 ++++++++ src/flowsense/domain/enums.py | 23 +++++++++ src/flowsense/domain/models.py | 19 +++++++ src/flowsense/domain/results.py | 61 +++++++++++++++++++++++ src/flowsense/engine/analyzer.py | 14 ++---- src/flowsense/engine/drift.py | 20 ++------ src/flowsense/engine/history.py | 2 +- src/flowsense/engine/impact.py | 48 ++++++------------ src/flowsense/engine/propagation.py | 30 +++-------- src/flowsense/engine/root_cause.py | 36 +++++-------- src/flowsense/engine/timing.py | 4 +- src/flowsense/mcp/server.py | 2 +- src/flowsense/models/__init__.py | 3 +- src/flowsense/models/dag_analysis.py | 30 +---------- src/flowsense/models/task_run.py | 20 +------- tests/test_domain.py | 39 +++++++++++++++ 17 files changed, 218 insertions(+), 157 deletions(-) create mode 100644 src/flowsense/domain/__init__.py create mode 100644 src/flowsense/domain/enums.py create mode 100644 src/flowsense/domain/models.py create mode 100644 src/flowsense/domain/results.py create mode 100644 tests/test_domain.py diff --git a/src/flowsense/collector/airflow_client.py b/src/flowsense/collector/airflow_client.py index 33d7061..412d1ff 100644 --- a/src/flowsense/collector/airflow_client.py +++ b/src/flowsense/collector/airflow_client.py @@ -3,7 +3,7 @@ import httpx from flowsense.config import get_airflow_config -from flowsense.models import TaskRun +from flowsense.domain import TaskRun PAGE_SIZE = 100 diff --git a/src/flowsense/domain/__init__.py b/src/flowsense/domain/__init__.py new file mode 100644 index 0000000..8f78cc1 --- /dev/null +++ b/src/flowsense/domain/__init__.py @@ -0,0 +1,22 @@ +from flowsense.domain.enums import ImpactClassification, Severity +from flowsense.domain.models import TaskRun +from flowsense.domain.results import ( + AnalysisDiagnostic, + DAGAnalysis, + DriftResult, + PropagationResult, + RootCauseResult, + TaskImpact, +) + +__all__ = [ + "AnalysisDiagnostic", + "DAGAnalysis", + "DriftResult", + "ImpactClassification", + "PropagationResult", + "RootCauseResult", + "Severity", + "TaskImpact", + "TaskRun", +] diff --git a/src/flowsense/domain/enums.py b/src/flowsense/domain/enums.py new file mode 100644 index 0000000..c86a7ab --- /dev/null +++ b/src/flowsense/domain/enums.py @@ -0,0 +1,23 @@ +from enum import StrEnum + + +class Severity(StrEnum): + NORMAL = "NORMAL" + MEDIUM = "MEDIUM" + HIGH = "HIGH" + CRITICAL = "CRITICAL" + + +class ImpactClassification(StrEnum): + NORMAL = "NORMAL" + OWN_DRIFT = "OWN_DRIFT" + INHERITED_DELAY = "INHERITED_DELAY" + COMBINED = "COMBINED" + + +SEVERITY_SCORE: dict[Severity, int] = { + Severity.NORMAL: 0, + Severity.MEDIUM: 1, + Severity.HIGH: 2, + Severity.CRITICAL: 3, +} diff --git a/src/flowsense/domain/models.py b/src/flowsense/domain/models.py new file mode 100644 index 0000000..3e984cc --- /dev/null +++ b/src/flowsense/domain/models.py @@ -0,0 +1,19 @@ +from datetime import datetime + +from pydantic import BaseModel + + +class TaskRun(BaseModel): + dag_id: str + dag_run_id: str + task_id: str + + state: str | None = None + + start_date: datetime | None = None + end_date: datetime | None = None + + duration: float | None = None + try_number: int = 0 + + map_index: int = -1 diff --git a/src/flowsense/domain/results.py b/src/flowsense/domain/results.py new file mode 100644 index 0000000..df4ef7e --- /dev/null +++ b/src/flowsense/domain/results.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +from flowsense.domain.enums import ImpactClassification, Severity + + +@dataclass(frozen=True) +class DriftResult: + task_id: str + baseline: float + current: float + mad: float + robust_z_score: float + deviation_percent: float + severity: Severity + + +@dataclass(frozen=True) +class TaskImpact: + task_id: str + classification: ImpactClassification + task_severity: Severity + upstream_handoff_severity: Severity | None + + +@dataclass(frozen=True) +class PropagationResult: + origin_task: str + affected_tasks: list[str] + path: list[str] + propagation_score: float + + +@dataclass(frozen=True) +class RootCauseResult: + task_id: str + classification: ImpactClassification + severity: Severity + propagation_score: float + + +@dataclass(frozen=True) +class AnalysisDiagnostic: + code: str + subject_id: str + message: str + + +@dataclass +class DAGAnalysis: + dag_id: str + runs_analyzed: int + overall_severity: Severity + 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]] + diagnostics: list[AnalysisDiagnostic] = field(default_factory=list) diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py index 2050d52..b78a5be 100644 --- a/src/flowsense/engine/analyzer.py +++ b/src/flowsense/engine/analyzer.py @@ -1,6 +1,8 @@ from __future__ import annotations from flowsense.collector.airflow_client import AirflowClient +from flowsense.domain import AnalysisDiagnostic, DAGAnalysis, Severity +from flowsense.domain.enums import SEVERITY_SCORE from flowsense.engine.drift import calculate_drift from flowsense.engine.history import build_duration_history from flowsense.engine.impact import classify_task_impact @@ -10,7 +12,6 @@ build_handoff_history_with_diagnostics, calculate_handoff_drift, ) -from flowsense.models import AnalysisDiagnostic, DAGAnalysis def analyze_dag(dag_id: str) -> DAGAnalysis: @@ -97,14 +98,7 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: dependencies=dependencies, ) - severity_order = { - "NORMAL": 0, - "MEDIUM": 1, - "HIGH": 2, - "CRITICAL": 3, - } - - overall_severity = "NORMAL" + overall_severity = Severity.NORMAL all_drift_results = [ *drift_results.values(), @@ -114,7 +108,7 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: if all_drift_results: overall_severity = max( all_drift_results, - key=lambda result: severity_order[result.severity], + key=lambda result: SEVERITY_SCORE[result.severity], ).severity primary_origin = select_primary_origin( diff --git a/src/flowsense/engine/drift.py b/src/flowsense/engine/drift.py index e1e33e7..d44e357 100644 --- a/src/flowsense/engine/drift.py +++ b/src/flowsense/engine/drift.py @@ -1,20 +1,10 @@ from __future__ import annotations import math -from dataclasses import dataclass import numpy as np - -@dataclass -class DriftResult: - task_id: str - baseline: float - current: float - mad: float - robust_z_score: float - deviation_percent: float - severity: str +from flowsense.domain import DriftResult, Severity def calculate_drift( @@ -53,13 +43,13 @@ def calculate_drift( absolute_z = abs(robust_z_score) if absolute_z >= 5: - severity = "CRITICAL" + severity = Severity.CRITICAL elif absolute_z >= 3.5: - severity = "HIGH" + severity = Severity.HIGH elif absolute_z >= 2: - severity = "MEDIUM" + severity = Severity.MEDIUM else: - severity = "NORMAL" + severity = Severity.NORMAL return DriftResult( task_id=task_id, diff --git a/src/flowsense/engine/history.py b/src/flowsense/engine/history.py index 0e56d5f..612b60e 100644 --- a/src/flowsense/engine/history.py +++ b/src/flowsense/engine/history.py @@ -2,7 +2,7 @@ from collections import defaultdict -from flowsense.models import TaskRun +from flowsense.domain import TaskRun def build_duration_history( diff --git a/src/flowsense/engine/impact.py b/src/flowsense/engine/impact.py index fd69ca0..2fdbbc7 100644 --- a/src/flowsense/engine/impact.py +++ b/src/flowsense/engine/impact.py @@ -1,16 +1,12 @@ 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 +from flowsense.domain import ( + DriftResult, + ImpactClassification, + Severity, + TaskImpact, +) +from flowsense.domain.enums import SEVERITY_SCORE def classify_task_impact( @@ -21,42 +17,30 @@ def classify_task_impact( anomalous_handoffs = [ drift for drift in upstream_handoff_drifts - if drift.severity - in { - "MEDIUM", - "HIGH", - "CRITICAL", - } + if drift.severity in {Severity.MEDIUM, Severity.HIGH, Severity.CRITICAL} ] task_is_anomalous = task_drift.severity in { - "MEDIUM", - "HIGH", - "CRITICAL", + Severity.MEDIUM, + Severity.HIGH, + Severity.CRITICAL, } if task_is_anomalous and anomalous_handoffs: - classification = "COMBINED" + classification = ImpactClassification.COMBINED elif task_is_anomalous: - classification = "OWN_DRIFT" + classification = ImpactClassification.OWN_DRIFT elif anomalous_handoffs: - classification = "INHERITED_DELAY" + classification = ImpactClassification.INHERITED_DELAY else: - classification = "NORMAL" + classification = ImpactClassification.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], + key=lambda result: SEVERITY_SCORE[result.severity], ).severity return TaskImpact( diff --git a/src/flowsense/engine/propagation.py b/src/flowsense/engine/propagation.py index 507cdbc..6ebcd34 100644 --- a/src/flowsense/engine/propagation.py +++ b/src/flowsense/engine/propagation.py @@ -1,23 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass - -from flowsense.engine.drift import DriftResult - -SEVERITY_SCORE = { - "NORMAL": 0, - "MEDIUM": 1, - "HIGH": 2, - "CRITICAL": 3, -} - - -@dataclass -class PropagationResult: - origin_task: str - affected_tasks: list[str] - path: list[str] - propagation_score: float +from flowsense.domain import DriftResult, PropagationResult, Severity +from flowsense.domain.enums import SEVERITY_SCORE def _build_reverse_dependencies( @@ -51,12 +35,12 @@ def _has_anomalous_upstream( 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": + if upstream_drift is None or upstream_drift.severity == Severity.NORMAL: continue if upstream_drift.severity in { - "HIGH", - "CRITICAL", + Severity.HIGH, + Severity.CRITICAL, }: return True @@ -90,7 +74,7 @@ def _find_propagation_paths( if downstream_drift is None: continue - if downstream_drift.severity == "NORMAL": + if downstream_drift.severity == Severity.NORMAL: continue next_path = [*path, downstream_task] @@ -122,7 +106,7 @@ def analyze_propagation( reverse_dependencies = _build_reverse_dependencies(dependencies) for task_id, drift in drift_results.items(): - if drift.severity not in {"HIGH", "CRITICAL"}: + if drift.severity not in {Severity.HIGH, Severity.CRITICAL}: continue if _has_anomalous_upstream( diff --git a/src/flowsense/engine/root_cause.py b/src/flowsense/engine/root_cause.py index a018222..228cf4d 100644 --- a/src/flowsense/engine/root_cause.py +++ b/src/flowsense/engine/root_cause.py @@ -1,26 +1,14 @@ 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, -} +from flowsense.domain import ( + DriftResult, + ImpactClassification, + PropagationResult, + RootCauseResult, + Severity, + TaskImpact, +) +from flowsense.domain.enums import SEVERITY_SCORE def _build_reverse_dependencies( @@ -58,7 +46,7 @@ def _has_candidate_upstream( 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": + if upstream_drift is None or upstream_drift.severity == Severity.NORMAL: continue if upstream_task in candidate_tasks: @@ -87,8 +75,8 @@ def select_primary_origin( for task_id, impact in task_impacts.items() if impact.classification in { - "OWN_DRIFT", - "COMBINED", + ImpactClassification.OWN_DRIFT, + ImpactClassification.COMBINED, } } diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py index 0cc842b..efdc44c 100644 --- a/src/flowsense/engine/timing.py +++ b/src/flowsense/engine/timing.py @@ -2,8 +2,8 @@ from dataclasses import dataclass -from flowsense.engine.drift import DriftResult, calculate_drift -from flowsense.models import TaskRun +from flowsense.domain import DriftResult, TaskRun +from flowsense.engine.drift import calculate_drift @dataclass diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py index 620d538..7597a5b 100644 --- a/src/flowsense/mcp/server.py +++ b/src/flowsense/mcp/server.py @@ -2,8 +2,8 @@ from mcp.server import MCPServer +from flowsense.domain import DAGAnalysis from flowsense.engine.analyzer import analyze_dag -from flowsense.models import DAGAnalysis mcp = MCPServer("FlowSense Engine") diff --git a/src/flowsense/models/__init__.py b/src/flowsense/models/__init__.py index 0c373d7..ffa67f1 100644 --- a/src/flowsense/models/__init__.py +++ b/src/flowsense/models/__init__.py @@ -1,5 +1,4 @@ -from flowsense.models.dag_analysis import AnalysisDiagnostic, DAGAnalysis -from flowsense.models.task_run import TaskRun +from flowsense.domain import AnalysisDiagnostic, DAGAnalysis, TaskRun __all__ = [ "AnalysisDiagnostic", diff --git a/src/flowsense/models/dag_analysis.py b/src/flowsense/models/dag_analysis.py index 71bed5a..6592f58 100644 --- a/src/flowsense/models/dag_analysis.py +++ b/src/flowsense/models/dag_analysis.py @@ -1,29 +1,3 @@ -from __future__ import annotations +from flowsense.domain import AnalysisDiagnostic, DAGAnalysis -from dataclasses import dataclass, field - -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(frozen=True) -class AnalysisDiagnostic: - code: str - subject_id: str - message: str - - -@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]] - diagnostics: list[AnalysisDiagnostic] = field(default_factory=list) +__all__ = ["AnalysisDiagnostic", "DAGAnalysis"] diff --git a/src/flowsense/models/task_run.py b/src/flowsense/models/task_run.py index 3e984cc..cd9241c 100644 --- a/src/flowsense/models/task_run.py +++ b/src/flowsense/models/task_run.py @@ -1,19 +1,3 @@ -from datetime import datetime +from flowsense.domain import TaskRun -from pydantic import BaseModel - - -class TaskRun(BaseModel): - dag_id: str - dag_run_id: str - task_id: str - - state: str | None = None - - start_date: datetime | None = None - end_date: datetime | None = None - - duration: float | None = None - try_number: int = 0 - - map_index: int = -1 +__all__ = ["TaskRun"] diff --git a/tests/test_domain.py b/tests/test_domain.py new file mode 100644 index 0000000..8b4f367 --- /dev/null +++ b/tests/test_domain.py @@ -0,0 +1,39 @@ +from flowsense.domain import ( + AnalysisDiagnostic, + DAGAnalysis, + DriftResult, + ImpactClassification, + PropagationResult, + RootCauseResult, + Severity, + TaskImpact, + TaskRun, +) +from flowsense.engine.drift import DriftResult as LegacyDriftResult +from flowsense.engine.impact import TaskImpact as LegacyTaskImpact +from flowsense.engine.propagation import PropagationResult as LegacyPropagationResult +from flowsense.engine.root_cause import RootCauseResult as LegacyRootCauseResult +from flowsense.models import ( + AnalysisDiagnostic as LegacyAnalysisDiagnostic, +) +from flowsense.models import ( + DAGAnalysis as LegacyDAGAnalysis, +) +from flowsense.models import ( + TaskRun as LegacyTaskRun, +) + + +def test_domain_enums_are_string_compatible() -> None: + assert Severity.CRITICAL == "CRITICAL" + assert ImpactClassification.OWN_DRIFT == "OWN_DRIFT" + + +def test_legacy_model_imports_reexport_domain_types() -> None: + assert LegacyAnalysisDiagnostic is AnalysisDiagnostic + assert LegacyDAGAnalysis is DAGAnalysis + assert LegacyDriftResult is DriftResult + assert LegacyPropagationResult is PropagationResult + assert LegacyRootCauseResult is RootCauseResult + assert LegacyTaskImpact is TaskImpact + assert LegacyTaskRun is TaskRun From 7ffe6dcb85dae7826237363231de8acd4f744c03 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:18:04 +0300 Subject: [PATCH 08/16] refactor: inject DAG data source into analyzer --- src/flowsense/application/__init__.py | 4 + src/flowsense/application/analyzer.py | 132 +++++++++++++++++++++++++ src/flowsense/application/ports.py | 9 ++ src/flowsense/cli/main.py | 8 +- src/flowsense/engine/analyzer.py | 134 ++------------------------ src/flowsense/mcp/server.py | 8 +- tests/test_analyzer.py | 58 ++++------- 7 files changed, 186 insertions(+), 167 deletions(-) create mode 100644 src/flowsense/application/__init__.py create mode 100644 src/flowsense/application/analyzer.py create mode 100644 src/flowsense/application/ports.py diff --git a/src/flowsense/application/__init__.py b/src/flowsense/application/__init__.py new file mode 100644 index 0000000..97577db --- /dev/null +++ b/src/flowsense/application/__init__.py @@ -0,0 +1,4 @@ +from flowsense.application.analyzer import analyze_dag +from flowsense.application.ports import DAGDataSource + +__all__ = ["DAGDataSource", "analyze_dag"] diff --git a/src/flowsense/application/analyzer.py b/src/flowsense/application/analyzer.py new file mode 100644 index 0000000..48875dd --- /dev/null +++ b/src/flowsense/application/analyzer.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from flowsense.application.ports import DAGDataSource +from flowsense.domain import AnalysisDiagnostic, DAGAnalysis, Severity +from flowsense.domain.enums import SEVERITY_SCORE +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_with_diagnostics, + calculate_handoff_drift, +) + + +def analyze_dag( + dag_id: str, + source: DAGDataSource, +) -> DAGAnalysis: + task_runs = source.collect_task_runs(dag_id) + duration_history = build_duration_history(task_runs) + diagnostics: list[AnalysisDiagnostic] = [] + + 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 as exc: + diagnostics.append( + AnalysisDiagnostic( + code="INSUFFICIENT_TASK_HISTORY", + subject_id=task_id, + message=str(exc), + ) + ) + + dependencies = source.get_dag_dependencies(dag_id) + + handoff_history_result = build_handoff_history_with_diagnostics( + task_runs=task_runs, + dependencies=dependencies, + ) + handoff_history = handoff_history_result.history + + diagnostics.extend( + AnalysisDiagnostic( + code=diagnostic.code, + subject_id=f"{diagnostic.upstream_task}->{diagnostic.downstream_task}", + message=diagnostic.message, + ) + for diagnostic in handoff_history_result.diagnostics + ) + + 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 as exc: + diagnostics.append( + AnalysisDiagnostic( + code="INSUFFICIENT_HANDOFF_HISTORY", + subject_id=f"{upstream_task}->{downstream_task}", + message=str(exc), + ) + ) + + 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, + ) + + overall_severity = 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_SCORE[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, + diagnostics=diagnostics, + ) diff --git a/src/flowsense/application/ports.py b/src/flowsense/application/ports.py new file mode 100644 index 0000000..b36b60a --- /dev/null +++ b/src/flowsense/application/ports.py @@ -0,0 +1,9 @@ +from typing import Protocol + +from flowsense.domain import TaskRun + + +class DAGDataSource(Protocol): + def collect_task_runs(self, dag_id: str) -> list[TaskRun]: ... + + def get_dag_dependencies(self, dag_id: str) -> dict[str, list[str]]: ... diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index 585fb99..fc25596 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -4,7 +4,8 @@ from rich.console import Console from rich.table import Table -from flowsense.engine.analyzer import analyze_dag +from flowsense.application import analyze_dag +from flowsense.collector.airflow_client import AirflowClient app = typer.Typer( name="flowsense", @@ -27,7 +28,10 @@ def analyze( help="Airflow DAG id to analyze.", ), ) -> None: - analysis = analyze_dag(dag_id) + analysis = analyze_dag( + dag_id=dag_id, + source=AirflowClient(), + ) console.print(f"\n[bold]FlowSense Analysis — {analysis.dag_id}[/bold]\n") diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py index b78a5be..7392785 100644 --- a/src/flowsense/engine/analyzer.py +++ b/src/flowsense/engine/analyzer.py @@ -1,132 +1,16 @@ -from __future__ import annotations +"""Backward-compatible analyzer entry point. +New integrations should inject a data source into +``flowsense.application.analyze_dag``. +""" + +from flowsense.application import analyze_dag as analyze_dag_with_source from flowsense.collector.airflow_client import AirflowClient -from flowsense.domain import AnalysisDiagnostic, DAGAnalysis, Severity -from flowsense.domain.enums import SEVERITY_SCORE -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_with_diagnostics, - calculate_handoff_drift, -) +from flowsense.domain 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) - diagnostics: list[AnalysisDiagnostic] = [] - - 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 as exc: - diagnostics.append( - AnalysisDiagnostic( - code="INSUFFICIENT_TASK_HISTORY", - subject_id=task_id, - message=str(exc), - ) - ) - - dependencies = client.get_dag_dependencies(dag_id) - - handoff_history_result = build_handoff_history_with_diagnostics( - task_runs=task_runs, - dependencies=dependencies, - ) - handoff_history = handoff_history_result.history - - diagnostics.extend( - AnalysisDiagnostic( - code=diagnostic.code, - subject_id=(f"{diagnostic.upstream_task}->{diagnostic.downstream_task}"), - message=diagnostic.message, - ) - for diagnostic in handoff_history_result.diagnostics - ) - - 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 as exc: - diagnostics.append( - AnalysisDiagnostic( - code="INSUFFICIENT_HANDOFF_HISTORY", - subject_id=f"{_upstream_task}->{downstream_task}", - message=str(exc), - ) - ) - - 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, - ) - - overall_severity = 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_SCORE[result.severity], - ).severity - - primary_origin = select_primary_origin( - drift_results=drift_results, - task_impacts=task_impacts, - dependencies=dependencies, - propagation_results=propagation_results, - ) - - return DAGAnalysis( + return analyze_dag_with_source( 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, - diagnostics=diagnostics, + source=AirflowClient(), ) diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py index 7597a5b..59cb2b4 100644 --- a/src/flowsense/mcp/server.py +++ b/src/flowsense/mcp/server.py @@ -2,8 +2,9 @@ from mcp.server import MCPServer +from flowsense.application import analyze_dag +from flowsense.collector.airflow_client import AirflowClient from flowsense.domain import DAGAnalysis -from flowsense.engine.analyzer import analyze_dag mcp = MCPServer("FlowSense Engine") @@ -80,7 +81,10 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: @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) + analysis = analyze_dag( + dag_id=dag_id, + source=AirflowClient(), + ) return serialize_analysis(analysis) diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index 8b5dd41..2edd4ff 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -1,15 +1,12 @@ from datetime import UTC, datetime, timedelta -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock -from flowsense.engine.analyzer import analyze_dag +from flowsense.application import DAGDataSource, 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 +def test_analyze_dag_identifies_primary_origin() -> None: + client = MagicMock(spec=DAGDataSource) client.collect_task_runs.return_value = [ TaskRun( @@ -42,7 +39,7 @@ def test_analyze_dag_identifies_primary_origin( "load": [], } - analysis = analyze_dag("demo") + analysis = analyze_dag("demo", client) assert analysis.overall_severity == "CRITICAL" assert analysis.primary_origin is not None @@ -52,11 +49,8 @@ def test_analyze_dag_identifies_primary_origin( 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 +def test_analyze_dag_identifies_isolated_primary_origin() -> None: + client = MagicMock(spec=DAGDataSource) client.collect_task_runs.return_value = [ TaskRun( @@ -76,18 +70,15 @@ def test_analyze_dag_identifies_isolated_primary_origin( "transform": [], } - analysis = analyze_dag("demo") + analysis = analyze_dag("demo", client) 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 +def test_analyze_dag_calculates_handoff_drift() -> None: + client = MagicMock(spec=DAGDataSource) base_time = datetime( 2026, @@ -150,7 +141,7 @@ def test_analyze_dag_calculates_handoff_drift( "transform": [], } - analysis = analyze_dag("demo") + analysis = analyze_dag("demo", client) edge = ("extract", "transform") @@ -163,11 +154,8 @@ def test_analyze_dag_calculates_handoff_drift( 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 +def test_analyze_dag_uses_handoff_for_overall_severity() -> None: + client = MagicMock(spec=DAGDataSource) base_time = datetime( 2026, @@ -230,7 +218,7 @@ def test_analyze_dag_uses_handoff_for_overall_severity( "transform": [], } - analysis = analyze_dag("demo") + analysis = analyze_dag("demo", client) assert analysis.drift_results["extract"].severity == "NORMAL" assert analysis.drift_results["transform"].severity == "NORMAL" @@ -242,11 +230,8 @@ def test_analyze_dag_uses_handoff_for_overall_severity( assert analysis.overall_severity == "CRITICAL" -@patch("flowsense.engine.analyzer.AirflowClient") -def test_analyze_dag_reports_insufficient_history_diagnostics( - mock_client_class: MagicMock, -) -> None: - client = mock_client_class.return_value +def test_analyze_dag_reports_insufficient_history_diagnostics() -> None: + client = MagicMock(spec=DAGDataSource) base_time = datetime(2026, 8, 20, 10, 0, tzinfo=UTC) task_runs = [] @@ -281,7 +266,7 @@ def test_analyze_dag_reports_insufficient_history_diagnostics( "transform": [], } - analysis = analyze_dag("demo") + analysis = analyze_dag("demo", client) diagnostics = { (diagnostic.code, diagnostic.subject_id) for diagnostic in analysis.diagnostics @@ -294,11 +279,8 @@ def test_analyze_dag_reports_insufficient_history_diagnostics( } -@patch("flowsense.engine.analyzer.AirflowClient") -def test_analyze_dag_reports_missing_handoff_timestamp( - mock_client_class: MagicMock, -) -> None: - client = mock_client_class.return_value +def test_analyze_dag_reports_missing_handoff_timestamp() -> None: + client = MagicMock(spec=DAGDataSource) client.collect_task_runs.return_value = [ TaskRun( dag_id="demo", @@ -320,7 +302,7 @@ def test_analyze_dag_reports_missing_handoff_timestamp( "transform": [], } - analysis = analyze_dag("demo") + analysis = analyze_dag("demo", client) assert any( diagnostic.code == "MISSING_UPSTREAM_END_DATE" From 7738a9e294a31ce6ecdab14edf4eae55eb7daa02 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:18:09 +0300 Subject: [PATCH 09/16] refactor: isolate Airflow DTO mapping --- src/flowsense/collector/airflow_client.py | 192 +----------------- src/flowsense/infrastructure/__init__.py | 1 + .../infrastructure/airflow/__init__.py | 3 + .../infrastructure/airflow/client.py | 165 +++++++++++++++ src/flowsense/infrastructure/airflow/dto.py | 29 +++ .../infrastructure/airflow/mapper.py | 27 +++ tests/test_airflow_client.py | 8 +- tests/test_airflow_mapper.py | 58 ++++++ 8 files changed, 290 insertions(+), 193 deletions(-) create mode 100644 src/flowsense/infrastructure/__init__.py create mode 100644 src/flowsense/infrastructure/airflow/__init__.py create mode 100644 src/flowsense/infrastructure/airflow/client.py create mode 100644 src/flowsense/infrastructure/airflow/dto.py create mode 100644 src/flowsense/infrastructure/airflow/mapper.py create mode 100644 tests/test_airflow_mapper.py diff --git a/src/flowsense/collector/airflow_client.py b/src/flowsense/collector/airflow_client.py index 412d1ff..fd8defc 100644 --- a/src/flowsense/collector/airflow_client.py +++ b/src/flowsense/collector/airflow_client.py @@ -1,191 +1,5 @@ -from __future__ import annotations +"""Backward-compatible imports for the Airflow infrastructure adapter.""" -import httpx +from flowsense.infrastructure.airflow.client import PAGE_SIZE, AirflowClient -from flowsense.config import get_airflow_config -from flowsense.domain import TaskRun - -PAGE_SIZE = 100 - - -class AirflowClient: - def __init__( - self, - base_url: str | None = None, - username: str | None = None, - password: str | None = None, - ): - config = get_airflow_config() - - self.base_url = (base_url or config.base_url).rstrip("/") - - self.username = username or config.username - self.password = password or config.password - - self._token: str | None = None - - def _get_token(self) -> str: - if self._token: - return self._token - - response = httpx.post( - f"{self.base_url}/auth/token", - json={ - "username": self.username, - "password": self.password, - }, - timeout=10.0, - ) - - response.raise_for_status() - - data = response.json() - self._token = data["access_token"] - - return self._token - - def _headers(self) -> dict[str, str]: - token = self._get_token() - - return { - "Authorization": f"Bearer {token}", - "Accept": "application/json", - } - - def _get_paginated( - self, - url: str, - collection_key: str, - ) -> dict: - items: list[dict] = [] - offset = 0 - last_page: dict = {} - - while True: - response = httpx.get( - url, - headers=self._headers(), - params={ - "limit": PAGE_SIZE, - "offset": offset, - }, - timeout=10.0, - ) - - response.raise_for_status() - - last_page = response.json() - page_items = last_page[collection_key] - items.extend(page_items) - - total_entries = last_page.get("total_entries") - - if not page_items: - break - - if total_entries is not None and len(items) >= total_entries: - break - - if total_entries is None and len(page_items) < PAGE_SIZE: - break - - offset += len(page_items) - - return { - **last_page, - collection_key: items, - "total_entries": last_page.get("total_entries", len(items)), - } - - def get_dag_runs(self, dag_id: str) -> dict: - url = f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns" - - return self._get_paginated( - url=url, - collection_key="dag_runs", - ) - - def get_task_instances( - self, - dag_id: str, - dag_run_id: str, - ) -> dict: - url = f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances" - - return self._get_paginated( - url=url, - collection_key="task_instances", - ) - - def collect_task_runs( - self, - dag_id: str, - ) -> list[TaskRun]: - runs = self.get_dag_runs(dag_id) - - task_runs: list[TaskRun] = [] - - runs_data = sorted( - runs["dag_runs"], - key=lambda run: run.get("run_after") or run.get("queued_at") or "", - ) - - for run in runs_data: - if run.get("state") != "success": - continue - - run_id = run["dag_run_id"] - - tasks = self.get_task_instances( - dag_id, - run_id, - ) - - for task in tasks["task_instances"]: - if task.get("state") != "success": - continue - - task_runs.append( - TaskRun( - dag_id=dag_id, - dag_run_id=run_id, - task_id=task["task_id"], - state=task.get("state"), - start_date=task.get("start_date"), - end_date=task.get("end_date"), - duration=task.get("duration"), - try_number=task.get("try_number", 0), - map_index=task.get("map_index", -1), - ) - ) - - return task_runs - - def get_dag_tasks( - self, - dag_id: str, - ) -> dict: - url = f"{self.base_url}/api/v2/dags/{dag_id}/tasks" - - return self._get_paginated( - url=url, - collection_key="tasks", - ) - - def get_dag_dependencies( - self, - dag_id: str, - ) -> dict[str, list[str]]: - data = self.get_dag_tasks(dag_id) - - dependencies: dict[str, list[str]] = {} - - for task in data["tasks"]: - task_id = task["task_id"] - - dependencies[task_id] = task.get( - "downstream_task_ids", - [], - ) - - return dependencies +__all__ = ["PAGE_SIZE", "AirflowClient"] diff --git a/src/flowsense/infrastructure/__init__.py b/src/flowsense/infrastructure/__init__.py new file mode 100644 index 0000000..e971391 --- /dev/null +++ b/src/flowsense/infrastructure/__init__.py @@ -0,0 +1 @@ +"""Infrastructure adapters for external systems.""" diff --git a/src/flowsense/infrastructure/airflow/__init__.py b/src/flowsense/infrastructure/airflow/__init__.py new file mode 100644 index 0000000..f096438 --- /dev/null +++ b/src/flowsense/infrastructure/airflow/__init__.py @@ -0,0 +1,3 @@ +from flowsense.infrastructure.airflow.client import AirflowClient + +__all__ = ["AirflowClient"] diff --git a/src/flowsense/infrastructure/airflow/client.py b/src/flowsense/infrastructure/airflow/client.py new file mode 100644 index 0000000..4ead123 --- /dev/null +++ b/src/flowsense/infrastructure/airflow/client.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import httpx + +from flowsense.config import get_airflow_config +from flowsense.domain import TaskRun +from flowsense.infrastructure.airflow.dto import ( + AirflowDagRunDTO, + AirflowTaskDTO, + AirflowTaskInstanceDTO, +) +from flowsense.infrastructure.airflow.mapper import ( + map_dependencies, + map_task_instance, +) + +PAGE_SIZE = 100 + + +class AirflowClient: + def __init__( + self, + base_url: str | None = None, + username: str | None = None, + password: str | None = None, + ): + config = get_airflow_config() + + self.base_url = (base_url or config.base_url).rstrip("/") + self.username = username or config.username + self.password = password or config.password + self._token: str | None = None + + def _get_token(self) -> str: + if self._token: + return self._token + + response = httpx.post( + f"{self.base_url}/auth/token", + json={ + "username": self.username, + "password": self.password, + }, + timeout=10.0, + ) + response.raise_for_status() + + self._token = response.json()["access_token"] + return self._token + + def _headers(self) -> dict[str, str]: + return { + "Authorization": f"Bearer {self._get_token()}", + "Accept": "application/json", + } + + def _get_paginated( + self, + url: str, + collection_key: str, + ) -> dict: + items: list[dict] = [] + offset = 0 + last_page: dict = {} + + while True: + response = httpx.get( + url, + headers=self._headers(), + params={"limit": PAGE_SIZE, "offset": offset}, + timeout=10.0, + ) + response.raise_for_status() + + last_page = response.json() + page_items = last_page[collection_key] + items.extend(page_items) + total_entries = last_page.get("total_entries") + + if not page_items: + break + + if total_entries is not None and len(items) >= total_entries: + break + + if total_entries is None and len(page_items) < PAGE_SIZE: + break + + offset += len(page_items) + + return { + **last_page, + collection_key: items, + "total_entries": last_page.get("total_entries", len(items)), + } + + def get_dag_runs(self, dag_id: str) -> dict: + return self._get_paginated( + url=f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns", + collection_key="dag_runs", + ) + + def get_task_instances( + self, + dag_id: str, + dag_run_id: str, + ) -> dict: + return self._get_paginated( + url=( + f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns/" + f"{dag_run_id}/taskInstances" + ), + collection_key="task_instances", + ) + + def collect_task_runs(self, dag_id: str) -> list[TaskRun]: + response = self.get_dag_runs(dag_id) + dag_runs = [ + AirflowDagRunDTO.model_validate(item) for item in response["dag_runs"] + ] + dag_runs.sort( + key=lambda run: ( + (run.run_after or run.queued_at).timestamp() + if run.run_after or run.queued_at + else float("-inf") + ) + ) + + task_runs: list[TaskRun] = [] + + for dag_run in dag_runs: + if dag_run.state != "success": + continue + + response = self.get_task_instances( + dag_id=dag_id, + dag_run_id=dag_run.dag_run_id, + ) + task_instances = [ + AirflowTaskInstanceDTO.model_validate(item) + for item in response["task_instances"] + ] + + task_runs.extend( + map_task_instance( + dag_id=dag_id, + dag_run_id=dag_run.dag_run_id, + task=task, + ) + for task in task_instances + if task.state == "success" + ) + + return task_runs + + def get_dag_tasks(self, dag_id: str) -> dict: + return self._get_paginated( + url=f"{self.base_url}/api/v2/dags/{dag_id}/tasks", + collection_key="tasks", + ) + + def get_dag_dependencies(self, dag_id: str) -> dict[str, list[str]]: + response = self.get_dag_tasks(dag_id) + tasks = [AirflowTaskDTO.model_validate(item) for item in response["tasks"]] + return map_dependencies(tasks) diff --git a/src/flowsense/infrastructure/airflow/dto.py b/src/flowsense/infrastructure/airflow/dto.py new file mode 100644 index 0000000..043b4de --- /dev/null +++ b/src/flowsense/infrastructure/airflow/dto.py @@ -0,0 +1,29 @@ +from datetime import datetime + +from pydantic import BaseModel, ConfigDict, Field + + +class AirflowDTO(BaseModel): + model_config = ConfigDict(extra="ignore") + + +class AirflowDagRunDTO(AirflowDTO): + dag_run_id: str + state: str | None = None + run_after: datetime | None = None + queued_at: datetime | None = None + + +class AirflowTaskInstanceDTO(AirflowDTO): + task_id: str + state: str | None = None + start_date: datetime | None = None + end_date: datetime | None = None + duration: float | None = None + try_number: int = 0 + map_index: int = -1 + + +class AirflowTaskDTO(AirflowDTO): + task_id: str + downstream_task_ids: list[str] = Field(default_factory=list) diff --git a/src/flowsense/infrastructure/airflow/mapper.py b/src/flowsense/infrastructure/airflow/mapper.py new file mode 100644 index 0000000..c3ce2dd --- /dev/null +++ b/src/flowsense/infrastructure/airflow/mapper.py @@ -0,0 +1,27 @@ +from flowsense.domain import TaskRun +from flowsense.infrastructure.airflow.dto import ( + AirflowTaskDTO, + AirflowTaskInstanceDTO, +) + + +def map_task_instance( + dag_id: str, + dag_run_id: str, + task: AirflowTaskInstanceDTO, +) -> TaskRun: + return TaskRun( + dag_id=dag_id, + dag_run_id=dag_run_id, + task_id=task.task_id, + state=task.state, + start_date=task.start_date, + end_date=task.end_date, + duration=task.duration, + try_number=task.try_number, + map_index=task.map_index, + ) + + +def map_dependencies(tasks: list[AirflowTaskDTO]) -> dict[str, list[str]]: + return {task.task_id: list(task.downstream_task_ids) for task in tasks} diff --git a/tests/test_airflow_client.py b/tests/test_airflow_client.py index 112b1e0..371459d 100644 --- a/tests/test_airflow_client.py +++ b/tests/test_airflow_client.py @@ -2,14 +2,14 @@ import pytest -from flowsense.collector.airflow_client import PAGE_SIZE, AirflowClient from flowsense.config import AirflowConfig +from flowsense.infrastructure.airflow.client import PAGE_SIZE, AirflowClient @pytest.fixture def client() -> AirflowClient: with patch( - "flowsense.collector.airflow_client.get_airflow_config", + "flowsense.infrastructure.airflow.client.get_airflow_config", return_value=AirflowConfig( base_url="http://airflow.test", username="airflow", @@ -22,7 +22,7 @@ def client() -> AirflowClient: return airflow_client -@patch("flowsense.collector.airflow_client.httpx.get") +@patch("flowsense.infrastructure.airflow.client.httpx.get") def test_get_dag_runs_collects_all_pages( mock_get: MagicMock, client: AirflowClient, @@ -87,7 +87,7 @@ def test_get_dag_runs_collects_all_pages( ), ], ) -@patch("flowsense.collector.airflow_client.httpx.get") +@patch("flowsense.infrastructure.airflow.client.httpx.get") def test_paginated_endpoints_use_their_collection_key( mock_get: MagicMock, method_name: str, diff --git a/tests/test_airflow_mapper.py b/tests/test_airflow_mapper.py new file mode 100644 index 0000000..93cb17b --- /dev/null +++ b/tests/test_airflow_mapper.py @@ -0,0 +1,58 @@ +from datetime import UTC, datetime + +from flowsense.infrastructure.airflow.dto import ( + AirflowTaskDTO, + AirflowTaskInstanceDTO, +) +from flowsense.infrastructure.airflow.mapper import ( + map_dependencies, + map_task_instance, +) + + +def test_maps_airflow_task_instance_to_domain_task_run() -> None: + started_at = datetime(2026, 9, 1, 10, 0, tzinfo=UTC) + ended_at = datetime(2026, 9, 1, 10, 0, 3, tzinfo=UTC) + dto = AirflowTaskInstanceDTO.model_validate( + { + "task_id": "transform", + "state": "success", + "start_date": started_at.isoformat(), + "end_date": ended_at.isoformat(), + "duration": 3.0, + "try_number": 2, + "map_index": 4, + "airflow_only_field": "ignored", + } + ) + + task_run = map_task_instance( + dag_id="demo", + dag_run_id="run_1", + task=dto, + ) + + assert task_run.task_id == "transform" + assert task_run.start_date == started_at + assert task_run.end_date == ended_at + assert task_run.try_number == 2 + assert task_run.map_index == 4 + + +def test_maps_airflow_tasks_to_dependency_graph() -> None: + tasks = [ + AirflowTaskDTO( + task_id="extract", + downstream_task_ids=["transform"], + ), + AirflowTaskDTO( + task_id="transform", + ), + ] + + dependencies = map_dependencies(tasks) + + assert dependencies == { + "extract": ["transform"], + "transform": [], + } From badfa1b14f0f87e8290abb1111f9cfbc05561c87 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:20:45 +0300 Subject: [PATCH 10/16] refactor: reuse HTTP client for Airflow requests --- src/flowsense/cli/main.py | 11 +++--- src/flowsense/engine/analyzer.py | 11 +++--- .../infrastructure/airflow/client.py | 21 ++++++++--- src/flowsense/mcp/server.py | 11 +++--- tests/test_airflow_client.py | 36 ++++++++++++------- 5 files changed, 58 insertions(+), 32 deletions(-) diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index fc25596..fa2164b 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -5,7 +5,7 @@ from rich.table import Table from flowsense.application import analyze_dag -from flowsense.collector.airflow_client import AirflowClient +from flowsense.infrastructure.airflow import AirflowClient app = typer.Typer( name="flowsense", @@ -28,10 +28,11 @@ def analyze( help="Airflow DAG id to analyze.", ), ) -> None: - analysis = analyze_dag( - dag_id=dag_id, - source=AirflowClient(), - ) + with AirflowClient() as source: + analysis = analyze_dag( + dag_id=dag_id, + source=source, + ) console.print(f"\n[bold]FlowSense Analysis — {analysis.dag_id}[/bold]\n") diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py index 7392785..8ae1ebd 100644 --- a/src/flowsense/engine/analyzer.py +++ b/src/flowsense/engine/analyzer.py @@ -5,12 +5,13 @@ """ from flowsense.application import analyze_dag as analyze_dag_with_source -from flowsense.collector.airflow_client import AirflowClient from flowsense.domain import DAGAnalysis +from flowsense.infrastructure.airflow import AirflowClient def analyze_dag(dag_id: str) -> DAGAnalysis: - return analyze_dag_with_source( - dag_id=dag_id, - source=AirflowClient(), - ) + with AirflowClient() as source: + return analyze_dag_with_source( + dag_id=dag_id, + source=source, + ) diff --git a/src/flowsense/infrastructure/airflow/client.py b/src/flowsense/infrastructure/airflow/client.py index 4ead123..5d84b04 100644 --- a/src/flowsense/infrastructure/airflow/client.py +++ b/src/flowsense/infrastructure/airflow/client.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Self + import httpx from flowsense.config import get_airflow_config @@ -23,25 +25,37 @@ def __init__( base_url: str | None = None, username: str | None = None, password: str | None = None, + http_client: httpx.Client | None = None, ): config = get_airflow_config() self.base_url = (base_url or config.base_url).rstrip("/") self.username = username or config.username self.password = password or config.password + self._owns_http_client = http_client is None + self._http_client = http_client or httpx.Client(timeout=10.0) self._token: str | None = None + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def close(self) -> None: + if self._owns_http_client: + self._http_client.close() + def _get_token(self) -> str: if self._token: return self._token - response = httpx.post( + response = self._http_client.post( f"{self.base_url}/auth/token", json={ "username": self.username, "password": self.password, }, - timeout=10.0, ) response.raise_for_status() @@ -64,11 +78,10 @@ def _get_paginated( last_page: dict = {} while True: - response = httpx.get( + response = self._http_client.get( url, headers=self._headers(), params={"limit": PAGE_SIZE, "offset": offset}, - timeout=10.0, ) response.raise_for_status() diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py index 59cb2b4..3e4e74d 100644 --- a/src/flowsense/mcp/server.py +++ b/src/flowsense/mcp/server.py @@ -3,8 +3,8 @@ from mcp.server import MCPServer from flowsense.application import analyze_dag -from flowsense.collector.airflow_client import AirflowClient from flowsense.domain import DAGAnalysis +from flowsense.infrastructure.airflow import AirflowClient mcp = MCPServer("FlowSense Engine") @@ -81,10 +81,11 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: @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=dag_id, - source=AirflowClient(), - ) + with AirflowClient() as source: + analysis = analyze_dag( + dag_id=dag_id, + source=source, + ) return serialize_analysis(analysis) diff --git a/tests/test_airflow_client.py b/tests/test_airflow_client.py index 371459d..66fe588 100644 --- a/tests/test_airflow_client.py +++ b/tests/test_airflow_client.py @@ -1,5 +1,6 @@ from unittest.mock import MagicMock, call, patch +import httpx import pytest from flowsense.config import AirflowConfig @@ -7,7 +8,12 @@ @pytest.fixture -def client() -> AirflowClient: +def http_client() -> MagicMock: + return MagicMock(spec=httpx.Client) + + +@pytest.fixture +def client(http_client: MagicMock) -> AirflowClient: with patch( "flowsense.infrastructure.airflow.client.get_airflow_config", return_value=AirflowConfig( @@ -16,16 +22,15 @@ def client() -> AirflowClient: password="airflow", ), ): - airflow_client = AirflowClient() + airflow_client = AirflowClient(http_client=http_client) airflow_client._token = "token" return airflow_client -@patch("flowsense.infrastructure.airflow.client.httpx.get") def test_get_dag_runs_collects_all_pages( - mock_get: MagicMock, client: AirflowClient, + http_client: MagicMock, ) -> None: first_page = MagicMock() first_page.json.return_value = { @@ -39,13 +44,13 @@ def test_get_dag_runs_collects_all_pages( "total_entries": PAGE_SIZE + 1, } - mock_get.side_effect = [first_page, second_page] + http_client.get.side_effect = [first_page, second_page] result = client.get_dag_runs("demo") assert len(result["dag_runs"]) == PAGE_SIZE + 1 assert result["total_entries"] == PAGE_SIZE + 1 - assert mock_get.call_args_list == [ + assert http_client.get.call_args_list == [ call( "http://airflow.test/api/v2/dags/demo/dagRuns", headers={ @@ -53,7 +58,6 @@ def test_get_dag_runs_collects_all_pages( "Accept": "application/json", }, params={"limit": PAGE_SIZE, "offset": 0}, - timeout=10.0, ), call( "http://airflow.test/api/v2/dags/demo/dagRuns", @@ -62,7 +66,6 @@ def test_get_dag_runs_collects_all_pages( "Accept": "application/json", }, params={"limit": PAGE_SIZE, "offset": PAGE_SIZE}, - timeout=10.0, ), ] @@ -87,31 +90,38 @@ def test_get_dag_runs_collects_all_pages( ), ], ) -@patch("flowsense.infrastructure.airflow.client.httpx.get") def test_paginated_endpoints_use_their_collection_key( - mock_get: MagicMock, method_name: str, collection_key: str, expected_url: str, args: tuple[str, ...], client: AirflowClient, + http_client: MagicMock, ) -> None: response = MagicMock() response.json.return_value = { collection_key: [{"id": "item_1"}], "total_entries": 1, } - mock_get.return_value = response + http_client.get.return_value = response result = getattr(client, method_name)(*args) assert result[collection_key] == [{"id": "item_1"}] - mock_get.assert_called_once_with( + http_client.get.assert_called_once_with( expected_url, headers={ "Authorization": "Bearer token", "Accept": "application/json", }, params={"limit": PAGE_SIZE, "offset": 0}, - timeout=10.0, ) + + +def test_does_not_close_injected_http_client( + client: AirflowClient, + http_client: MagicMock, +) -> None: + client.close() + + http_client.close.assert_not_called() From 145f3eaa3a7fcbd689e2ff92b20e9f073a36fd16 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:23:39 +0300 Subject: [PATCH 11/16] refactor: introduce domain-specific exceptions --- src/flowsense/application/analyzer.py | 11 ++++++++--- src/flowsense/domain/__init__.py | 8 ++++++++ src/flowsense/domain/exceptions.py | 21 ++++++++++++++++++++ src/flowsense/engine/drift.py | 8 ++++++-- src/flowsense/engine/timing.py | 10 +++++----- tests/test_analyzer.py | 28 ++++++++++++++++++++++++++- tests/test_drift.py | 13 ++++++++----- tests/test_timing.py | 5 +++-- 8 files changed, 86 insertions(+), 18 deletions(-) create mode 100644 src/flowsense/domain/exceptions.py diff --git a/src/flowsense/application/analyzer.py b/src/flowsense/application/analyzer.py index 48875dd..b1ae8ad 100644 --- a/src/flowsense/application/analyzer.py +++ b/src/flowsense/application/analyzer.py @@ -1,7 +1,12 @@ from __future__ import annotations from flowsense.application.ports import DAGDataSource -from flowsense.domain import AnalysisDiagnostic, DAGAnalysis, Severity +from flowsense.domain import ( + AnalysisDiagnostic, + DAGAnalysis, + InsufficientHistoryError, + Severity, +) from flowsense.domain.enums import SEVERITY_SCORE from flowsense.engine.drift import calculate_drift from flowsense.engine.history import build_duration_history @@ -30,7 +35,7 @@ def analyze_dag( task_id=task_id, durations=durations, ) - except ValueError as exc: + except InsufficientHistoryError as exc: diagnostics.append( AnalysisDiagnostic( code="INSUFFICIENT_TASK_HISTORY", @@ -67,7 +72,7 @@ def analyze_dag( downstream_task=downstream_task, handoff_delays=delays, ) - except ValueError as exc: + except InsufficientHistoryError as exc: diagnostics.append( AnalysisDiagnostic( code="INSUFFICIENT_HANDOFF_HISTORY", diff --git a/src/flowsense/domain/__init__.py b/src/flowsense/domain/__init__.py index 8f78cc1..8eccfe1 100644 --- a/src/flowsense/domain/__init__.py +++ b/src/flowsense/domain/__init__.py @@ -1,4 +1,9 @@ from flowsense.domain.enums import ImpactClassification, Severity +from flowsense.domain.exceptions import ( + FlowSenseError, + InsufficientHistoryError, + InvalidTaskTimingError, +) from flowsense.domain.models import TaskRun from flowsense.domain.results import ( AnalysisDiagnostic, @@ -13,7 +18,10 @@ "AnalysisDiagnostic", "DAGAnalysis", "DriftResult", + "FlowSenseError", "ImpactClassification", + "InsufficientHistoryError", + "InvalidTaskTimingError", "PropagationResult", "RootCauseResult", "Severity", diff --git a/src/flowsense/domain/exceptions.py b/src/flowsense/domain/exceptions.py new file mode 100644 index 0000000..3f9076c --- /dev/null +++ b/src/flowsense/domain/exceptions.py @@ -0,0 +1,21 @@ +class FlowSenseError(Exception): + """Base exception for expected FlowSense failures.""" + + +class InsufficientHistoryError(FlowSenseError, ValueError): + def __init__( + self, + subject_id: str, + required: int, + actual: int, + ) -> None: + self.subject_id = subject_id + self.required = required + self.actual = actual + super().__init__( + f"{subject_id} için drift hesaplamak için en az {required} run gerekli." + ) + + +class InvalidTaskTimingError(FlowSenseError, ValueError): + """Raised when task timestamps cannot produce a valid handoff timing.""" diff --git a/src/flowsense/engine/drift.py b/src/flowsense/engine/drift.py index d44e357..aabea8c 100644 --- a/src/flowsense/engine/drift.py +++ b/src/flowsense/engine/drift.py @@ -4,7 +4,7 @@ import numpy as np -from flowsense.domain import DriftResult, Severity +from flowsense.domain import DriftResult, InsufficientHistoryError, Severity def calculate_drift( @@ -12,7 +12,11 @@ def calculate_drift( durations: list[float], ) -> DriftResult: if len(durations) < 5: - raise ValueError(f"{task_id} için drift hesaplamak için en az 5 run gerekli.") + raise InsufficientHistoryError( + subject_id=task_id, + required=5, + actual=len(durations), + ) baseline_values = np.array( durations[:-1], diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py index efdc44c..67e00a5 100644 --- a/src/flowsense/engine/timing.py +++ b/src/flowsense/engine/timing.py @@ -2,7 +2,7 @@ from dataclasses import dataclass -from flowsense.domain import DriftResult, TaskRun +from flowsense.domain import DriftResult, InvalidTaskTimingError, TaskRun from flowsense.engine.drift import calculate_drift @@ -34,13 +34,13 @@ def calculate_handoff_delay( downstream_run: TaskRun, ) -> HandoffTiming: if upstream_run.end_date is None: - raise ValueError("Upstream task end_date is required.") + raise InvalidTaskTimingError("Upstream task end_date is required.") if downstream_run.start_date is None: - raise ValueError("Downstream task start_date is required.") + raise InvalidTaskTimingError("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.") + raise InvalidTaskTimingError("Task runs must belong to the same DAG run.") handoff_delay = (downstream_run.start_date - upstream_run.end_date).total_seconds() @@ -162,7 +162,7 @@ def build_handoff_history_with_diagnostics( upstream_run=upstream_run, downstream_run=downstream_run, ) - except ValueError as exc: + except InvalidTaskTimingError as exc: diagnostics.append( HandoffHistoryDiagnostic( code="INVALID_HANDOFF_TIMING", diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index 2edd4ff..be1ef8a 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -1,5 +1,7 @@ from datetime import UTC, datetime, timedelta -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch + +import pytest from flowsense.application import DAGDataSource, analyze_dag from flowsense.models import TaskRun @@ -309,3 +311,27 @@ def test_analyze_dag_reports_missing_handoff_timestamp() -> None: and diagnostic.subject_id == "extract->transform" for diagnostic in analysis.diagnostics ) + + +def test_analyze_dag_does_not_hide_unexpected_value_errors() -> None: + client = MagicMock(spec=DAGDataSource) + client.collect_task_runs.return_value = [ + TaskRun( + dag_id="demo", + dag_run_id=f"run_{index}", + task_id="transform", + state="success", + duration=3.0, + ) + for index in range(5) + ] + client.get_dag_dependencies.return_value = {"transform": []} + + with ( + patch( + "flowsense.application.analyzer.calculate_drift", + side_effect=ValueError("unexpected failure"), + ), + pytest.raises(ValueError, match="unexpected failure"), + ): + analyze_dag("demo", client) diff --git a/tests/test_drift.py b/tests/test_drift.py index 5495d50..d4e2b32 100644 --- a/tests/test_drift.py +++ b/tests/test_drift.py @@ -1,3 +1,6 @@ +import pytest + +from flowsense.domain import InsufficientHistoryError from flowsense.engine.drift import calculate_drift @@ -71,12 +74,12 @@ def test_calculate_drift_requires_minimum_history() -> None: 3.3, ] - try: + with pytest.raises(InsufficientHistoryError) as exc_info: calculate_drift( "transform", durations, ) - except ValueError as exc: - assert "en az 5 run" in str(exc) - else: - raise AssertionError("Expected ValueError for insufficient history.") + + assert exc_info.value.subject_id == "transform" + assert exc_info.value.required == 5 + assert exc_info.value.actual == 4 diff --git a/tests/test_timing.py b/tests/test_timing.py index b8aa808..7711bb3 100644 --- a/tests/test_timing.py +++ b/tests/test_timing.py @@ -2,6 +2,7 @@ import pytest +from flowsense.domain import InvalidTaskTimingError from flowsense.engine.timing import ( build_handoff_history, build_handoff_history_with_diagnostics, @@ -55,7 +56,7 @@ def test_calculate_handoff_delay_requires_upstream_end_date() -> None: start_date=datetime(2026, 8, 20, 10, 0, tzinfo=UTC), ) - with pytest.raises(ValueError): + with pytest.raises(InvalidTaskTimingError): calculate_handoff_delay( upstream_run=upstream_run, downstream_run=downstream_run, @@ -77,7 +78,7 @@ def test_calculate_handoff_delay_requires_same_dag_run() -> None: start_date=datetime(2026, 8, 20, 10, 0, 5, tzinfo=UTC), ) - with pytest.raises(ValueError): + with pytest.raises(InvalidTaskTimingError): calculate_handoff_delay( upstream_run=upstream_run, downstream_run=downstream_run, From 47aa2860ebe8859795ea1bcb71a493692e345a7b Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:28:21 +0300 Subject: [PATCH 12/16] feat: handle Airflow API errors safely --- src/flowsense/cli/main.py | 16 +++++--- .../infrastructure/airflow/__init__.py | 3 +- .../infrastructure/airflow/client.py | 40 +++++++++++++++--- .../infrastructure/airflow/exceptions.py | 16 ++++++++ src/flowsense/mcp/server.py | 16 +++++--- tests/test_airflow_client.py | 41 +++++++++++++++---- tests/test_cli.py | 22 ++++++++++ 7 files changed, 128 insertions(+), 26 deletions(-) create mode 100644 src/flowsense/infrastructure/airflow/exceptions.py create mode 100644 tests/test_cli.py diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index fa2164b..5ba22e8 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -5,7 +5,7 @@ from rich.table import Table from flowsense.application import analyze_dag -from flowsense.infrastructure.airflow import AirflowClient +from flowsense.infrastructure.airflow import AirflowApiError, AirflowClient app = typer.Typer( name="flowsense", @@ -28,11 +28,15 @@ def analyze( help="Airflow DAG id to analyze.", ), ) -> None: - with AirflowClient() as source: - analysis = analyze_dag( - dag_id=dag_id, - source=source, - ) + try: + with AirflowClient() as source: + analysis = analyze_dag( + dag_id=dag_id, + source=source, + ) + except AirflowApiError as exc: + console.print(f"[bold red]Airflow request failed:[/bold red] {exc}") + raise typer.Exit(code=1) from exc console.print(f"\n[bold]FlowSense Analysis — {analysis.dag_id}[/bold]\n") diff --git a/src/flowsense/infrastructure/airflow/__init__.py b/src/flowsense/infrastructure/airflow/__init__.py index f096438..1ccf884 100644 --- a/src/flowsense/infrastructure/airflow/__init__.py +++ b/src/flowsense/infrastructure/airflow/__init__.py @@ -1,3 +1,4 @@ from flowsense.infrastructure.airflow.client import AirflowClient +from flowsense.infrastructure.airflow.exceptions import AirflowApiError -__all__ = ["AirflowClient"] +__all__ = ["AirflowApiError", "AirflowClient"] diff --git a/src/flowsense/infrastructure/airflow/client.py b/src/flowsense/infrastructure/airflow/client.py index 5d84b04..04f4911 100644 --- a/src/flowsense/infrastructure/airflow/client.py +++ b/src/flowsense/infrastructure/airflow/client.py @@ -11,6 +11,7 @@ AirflowTaskDTO, AirflowTaskInstanceDTO, ) +from flowsense.infrastructure.airflow.exceptions import AirflowApiError from flowsense.infrastructure.airflow.mapper import ( map_dependencies, map_task_instance, @@ -46,18 +47,45 @@ def close(self) -> None: if self._owns_http_client: self._http_client.close() + def _request( + self, + method: str, + url: str, + **kwargs: object, + ) -> httpx.Response: + try: + response = self._http_client.request( + method=method, + url=url, + **kwargs, + ) + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise AirflowApiError( + method=method, + endpoint=exc.request.url.path, + status_code=exc.response.status_code, + ) from exc + except httpx.RequestError as exc: + raise AirflowApiError( + method=method, + endpoint=exc.request.url.path, + ) from exc + + return response + def _get_token(self) -> str: if self._token: return self._token - response = self._http_client.post( - f"{self.base_url}/auth/token", + response = self._request( + method="POST", + url=f"{self.base_url}/auth/token", json={ "username": self.username, "password": self.password, }, ) - response.raise_for_status() self._token = response.json()["access_token"] return self._token @@ -78,12 +106,12 @@ def _get_paginated( last_page: dict = {} while True: - response = self._http_client.get( - url, + response = self._request( + method="GET", + url=url, headers=self._headers(), params={"limit": PAGE_SIZE, "offset": offset}, ) - response.raise_for_status() last_page = response.json() page_items = last_page[collection_key] diff --git a/src/flowsense/infrastructure/airflow/exceptions.py b/src/flowsense/infrastructure/airflow/exceptions.py new file mode 100644 index 0000000..44be514 --- /dev/null +++ b/src/flowsense/infrastructure/airflow/exceptions.py @@ -0,0 +1,16 @@ +from flowsense.domain import FlowSenseError + + +class AirflowApiError(FlowSenseError): + def __init__( + self, + method: str, + endpoint: str, + status_code: int | None = None, + ) -> None: + self.method = method + self.endpoint = endpoint + self.status_code = status_code + + status = f" with status {status_code}" if status_code is not None else "" + super().__init__(f"Airflow API {method} {endpoint} failed{status}.") diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py index 3e4e74d..2ecf0c0 100644 --- a/src/flowsense/mcp/server.py +++ b/src/flowsense/mcp/server.py @@ -4,7 +4,7 @@ from flowsense.application import analyze_dag from flowsense.domain import DAGAnalysis -from flowsense.infrastructure.airflow import AirflowClient +from flowsense.infrastructure.airflow import AirflowApiError, AirflowClient mcp = MCPServer("FlowSense Engine") @@ -81,11 +81,15 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: @mcp.tool() def analyze_airflow_dag(dag_id: str) -> dict: """Analyze an Apache Airflow DAG for temporal drift and propagation.""" - with AirflowClient() as source: - analysis = analyze_dag( - dag_id=dag_id, - source=source, - ) + try: + with AirflowClient() as source: + analysis = analyze_dag( + dag_id=dag_id, + source=source, + ) + except AirflowApiError as exc: + raise RuntimeError(str(exc)) from exc + return serialize_analysis(analysis) diff --git a/tests/test_airflow_client.py b/tests/test_airflow_client.py index 66fe588..9aa08e4 100644 --- a/tests/test_airflow_client.py +++ b/tests/test_airflow_client.py @@ -4,6 +4,7 @@ import pytest from flowsense.config import AirflowConfig +from flowsense.infrastructure.airflow import AirflowApiError from flowsense.infrastructure.airflow.client import PAGE_SIZE, AirflowClient @@ -44,15 +45,16 @@ def test_get_dag_runs_collects_all_pages( "total_entries": PAGE_SIZE + 1, } - http_client.get.side_effect = [first_page, second_page] + http_client.request.side_effect = [first_page, second_page] result = client.get_dag_runs("demo") assert len(result["dag_runs"]) == PAGE_SIZE + 1 assert result["total_entries"] == PAGE_SIZE + 1 - assert http_client.get.call_args_list == [ + assert http_client.request.call_args_list == [ call( - "http://airflow.test/api/v2/dags/demo/dagRuns", + method="GET", + url="http://airflow.test/api/v2/dags/demo/dagRuns", headers={ "Authorization": "Bearer token", "Accept": "application/json", @@ -60,7 +62,8 @@ def test_get_dag_runs_collects_all_pages( params={"limit": PAGE_SIZE, "offset": 0}, ), call( - "http://airflow.test/api/v2/dags/demo/dagRuns", + method="GET", + url="http://airflow.test/api/v2/dags/demo/dagRuns", headers={ "Authorization": "Bearer token", "Accept": "application/json", @@ -103,13 +106,14 @@ def test_paginated_endpoints_use_their_collection_key( collection_key: [{"id": "item_1"}], "total_entries": 1, } - http_client.get.return_value = response + http_client.request.return_value = response result = getattr(client, method_name)(*args) assert result[collection_key] == [{"id": "item_1"}] - http_client.get.assert_called_once_with( - expected_url, + http_client.request.assert_called_once_with( + method="GET", + url=expected_url, headers={ "Authorization": "Bearer token", "Accept": "application/json", @@ -125,3 +129,26 @@ def test_does_not_close_injected_http_client( client.close() http_client.close.assert_not_called() + + +def test_wraps_http_status_errors_without_response_body( + client: AirflowClient, + http_client: MagicMock, +) -> None: + request = httpx.Request( + "GET", + "http://airflow.test/api/v2/dags/demo/dagRuns", + ) + http_client.request.return_value = httpx.Response( + status_code=503, + request=request, + text="internal server details", + ) + + with pytest.raises(AirflowApiError) as exc_info: + client.get_dag_runs("demo") + + assert exc_info.value.method == "GET" + assert exc_info.value.endpoint == "/api/v2/dags/demo/dagRuns" + assert exc_info.value.status_code == 503 + assert "internal server details" not in str(exc_info.value) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..b6c8204 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,22 @@ +from unittest.mock import patch + +from typer.testing import CliRunner + +from flowsense.cli.main import app +from flowsense.infrastructure.airflow import AirflowApiError + + +def test_analyze_reports_airflow_api_errors() -> None: + error = AirflowApiError( + method="GET", + endpoint="/api/v2/dags/demo/dagRuns", + status_code=503, + ) + + with patch("flowsense.cli.main.AirflowClient") as client_class: + client_class.return_value.__enter__.side_effect = error + result = CliRunner().invoke(app, ["analyze", "demo"]) + + assert result.exit_code == 1 + assert "Airflow request failed" in result.output + assert "503" in result.output From 0b9ddb82e0f9d20f46b5423442f0f851d89ac5bd Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:33:49 +0300 Subject: [PATCH 13/16] chore: add static type checking --- .github/workflows/ci.yaml | 5 +++- pyproject.toml | 11 ++++++++- src/flowsense/engine/timing.py | 17 +++++++++++-- .../infrastructure/airflow/client.py | 24 ++++++++++--------- src/flowsense/py.typed | 0 uv.lock | 24 +++++++++++++++++++ 6 files changed, 66 insertions(+), 15 deletions(-) create mode 100644 src/flowsense/py.typed diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 6a0bdca..f8a62ee 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -28,5 +28,8 @@ jobs: - name: Run Ruff run: ruff check . + - name: Run Pyright + run: pyright + - name: Run unit tests - run: python -m pytest -m "not integration" -v \ No newline at end of file + run: python -m pytest -m "not integration" -v diff --git a/pyproject.toml b/pyproject.toml index f966611..32b8aa9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ [project.optional-dependencies] dev = [ + "pyright>=1.1.400", "pytest>=8.0", "ruff>=0.6", ] @@ -38,8 +39,16 @@ package-dir = {"" = "src"} [tool.setuptools.packages.find] where = ["src"] +[tool.setuptools.package-data] +flowsense = ["py.typed"] + +[tool.pyright] +include = ["src"] +typeCheckingMode = "standard" +pythonVersion = "3.12" + [tool.pytest.ini_options] pythonpath = ["src"] markers = [ "integration: tests that require a running Airflow instance", -] \ No newline at end of file +] diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py index 67e00a5..9db3d19 100644 --- a/src/flowsense/engine/timing.py +++ b/src/flowsense/engine/timing.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from datetime import datetime from flowsense.domain import DriftResult, InvalidTaskTimingError, TaskRun from flowsense.engine.drift import calculate_drift @@ -29,6 +30,18 @@ class HandoffHistoryResult: diagnostics: list[HandoffHistoryDiagnostic] +def _required_end_date(task_run: TaskRun) -> datetime: + if task_run.end_date is None: + raise InvalidTaskTimingError("Upstream task end_date is required.") + return task_run.end_date + + +def _required_start_date(task_run: TaskRun) -> datetime: + if task_run.start_date is None: + raise InvalidTaskTimingError("Downstream task start_date is required.") + return task_run.start_date + + def calculate_handoff_delay( upstream_run: TaskRun, downstream_run: TaskRun, @@ -117,7 +130,7 @@ def build_handoff_history_with_diagnostics( upstream_run = max( upstream_runs, - key=lambda task_run: task_run.end_date, + key=_required_end_date, ) for downstream_task in downstream_tasks: @@ -154,7 +167,7 @@ def build_handoff_history_with_diagnostics( downstream_run = min( downstream_runs, - key=lambda task_run: task_run.start_date, + key=_required_start_date, ) try: diff --git a/src/flowsense/infrastructure/airflow/client.py b/src/flowsense/infrastructure/airflow/client.py index 04f4911..18a3948 100644 --- a/src/flowsense/infrastructure/airflow/client.py +++ b/src/flowsense/infrastructure/airflow/client.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Self +from typing import Any, Self import httpx @@ -53,11 +53,13 @@ def _request( url: str, **kwargs: object, ) -> httpx.Response: + request_kwargs: dict[str, Any] = dict(kwargs) + try: response = self._http_client.request( method=method, url=url, - **kwargs, + **request_kwargs, ) response.raise_for_status() except httpx.HTTPStatusError as exc: @@ -87,8 +89,9 @@ def _get_token(self) -> str: }, ) - self._token = response.json()["access_token"] - return self._token + token: str = response.json()["access_token"] + self._token = token + return token def _headers(self) -> dict[str, str]: return { @@ -159,13 +162,12 @@ def collect_task_runs(self, dag_id: str) -> list[TaskRun]: dag_runs = [ AirflowDagRunDTO.model_validate(item) for item in response["dag_runs"] ] - dag_runs.sort( - key=lambda run: ( - (run.run_after or run.queued_at).timestamp() - if run.run_after or run.queued_at - else float("-inf") - ) - ) + + def run_timestamp(run: AirflowDagRunDTO) -> float: + timestamp = run.run_after or run.queued_at + return timestamp.timestamp() if timestamp is not None else float("-inf") + + dag_runs.sort(key=run_timestamp) task_runs: list[TaskRun] = [] diff --git a/src/flowsense/py.typed b/src/flowsense/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/uv.lock b/uv.lock index 64a7413..daf7c0e 100644 --- a/uv.lock +++ b/uv.lock @@ -231,6 +231,7 @@ dependencies = [ [package.optional-dependencies] dev = [ + { name = "pyright" }, { name = "pytest" }, { name = "ruff" }, ] @@ -245,6 +246,7 @@ requires-dist = [ { name = "numpy", specifier = ">=2.0" }, { name = "pandas", specifier = ">=2.2" }, { name = "pydantic", specifier = ">=2.8" }, + { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.400" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, { name = "rich", specifier = ">=13.7" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.6" }, @@ -439,6 +441,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "nodeenv" +version = "1.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/24/bf/d1bda4f6168e0b2e9e5958945e01910052158313224ada5ce1fb2e1113b8/nodeenv-1.10.0.tar.gz", hash = "sha256:996c191ad80897d076bdfba80a41994c2b47c68e224c542b48feba42ba00f8bb", size = 55611, upload-time = "2025-12-20T14:08:54.006Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" }, +] + [[package]] name = "numpy" version = "2.5.2" @@ -710,6 +721,19 @@ crypto = [ { name = "cryptography" }, ] +[[package]] +name = "pyright" +version = "1.1.411" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nodeenv" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7e/ab/265f7dc69d28113ebba19092e57b075f41543b2ed048429c5f56e2b88eac/pyright-1.1.411.tar.gz", hash = "sha256:d885a0551f2e763b089a02702174e7f4ba77548cddabc972ab86d1f7f1b0f998", size = 4112861, upload-time = "2026-06-25T02:14:06.37Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0a/49/385be530a6a5b78d1cbcd5c2e38debc8959a2fc6bdb716f4e581002979fc/pyright-1.1.411-py3-none-any.whl", hash = "sha256:dc7c72a8e2700c55baa127554040e067041ea53ccfd50bf96308cc4291c7d5d9", size = 6181526, upload-time = "2026-06-25T02:14:04.691Z" }, +] + [[package]] name = "pytest" version = "9.1.1" From 8f3d2ae300c3358913bea47e6d7c3ac9f42ac209 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:38:16 +0300 Subject: [PATCH 14/16] build: validate wheel packaging in CI --- .github/workflows/ci.yaml | 12 +++ pyproject.toml | 3 +- uv.lock | 156 ++++++-------------------------------- 3 files changed, 38 insertions(+), 133 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index f8a62ee..1662daa 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -28,8 +28,20 @@ jobs: - name: Run Ruff run: ruff check . + - name: Check formatting + run: ruff format --check . + - name: Run Pyright run: pyright - name: Run unit tests run: python -m pytest -m "not integration" -v + + - name: Build package + run: python -m build + + - name: Smoke test wheel + run: | + uv venv /tmp/flowsense-wheel-test + uv pip install --python /tmp/flowsense-wheel-test/bin/python dist/*.whl + /tmp/flowsense-wheel-test/bin/flowsense --help diff --git a/pyproject.toml b/pyproject.toml index 32b8aa9..bf2d34e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,14 +8,13 @@ dependencies = [ "httpx>=0.27", "pydantic>=2.8", "numpy>=2.0", - "pandas>=2.2", - "scipy>=1.14", "typer>=0.12", "rich>=13.7", ] [project.optional-dependencies] dev = [ + "build>=1.2", "pyright>=1.1.400", "pytest>=8.0", "ruff>=0.6", diff --git a/uv.lock b/uv.lock index daf7c0e..155527b 100644 --- a/uv.lock +++ b/uv.lock @@ -50,6 +50,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] +[[package]] +name = "build" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "os_name == 'nt'" }, + { name = "packaging" }, + { name = "pyproject-hooks" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4d/b7/1db48a9ce2984842c8c886432ec8a2719613322e868a966ba82a28862f25/build-1.6.0.tar.gz", hash = "sha256:bd2c8afc603e7a2e0ce70e2ea85f0a6d02043bafbd307f5bada0f98669eca5af", size = 113825, upload-time = "2026-08-27T21:01:16.458Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/e5/aa1e81b21aea0ce0ba435311837a37d4cb936e7461f9fecac08580073ba9/build-1.6.0-py3-none-any.whl", hash = "sha256:f7aaf1ebbb79178a02ba248bb524f2176b256017e17e8e4bd4289c7b38cc2bad", size = 31187, upload-time = "2026-08-27T21:01:14.957Z" }, +] + [[package]] name = "certifi" version = "2026.7.22" @@ -222,15 +236,14 @@ source = { editable = "." } dependencies = [ { name = "httpx" }, { name = "numpy" }, - { name = "pandas" }, { name = "pydantic" }, { name = "rich" }, - { name = "scipy" }, { name = "typer" }, ] [package.optional-dependencies] dev = [ + { name = "build" }, { name = "pyright" }, { name = "pytest" }, { name = "ruff" }, @@ -241,16 +254,15 @@ mcp = [ [package.metadata] requires-dist = [ + { name = "build", marker = "extra == 'dev'", specifier = ">=1.2" }, { name = "httpx", specifier = ">=0.27" }, { 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" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.400" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, { name = "rich", specifier = ">=13.7" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.6" }, - { name = "scipy", specifier = ">=1.14" }, { name = "typer", specifier = ">=0.12" }, ] provides-extras = ["dev", "mcp"] @@ -544,52 +556,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, ] -[[package]] -name = "pandas" -version = "3.0.5" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, - { name = "python-dateutil" }, - { name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/be/4f/5f3422a2afec5ffc46308b79e53291365a93748b498ac2e58bead0197916/pandas-3.0.5.tar.gz", hash = "sha256:dca3734d6ab7c906e6730f0788b0a1dbb9f2467731f9711f77995c8e9d62d712", size = 4658219, upload-time = "2026-07-22T22:19:28.819Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1c/54/1dc810ea558d1320b597aa140a514f2fdf1d2ea09c38cf556f13ea712ec9/pandas-3.0.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fa290c16964d4963fbfbc358928239cf3bd755b20e988ce944877def2f44471d", size = 10411717, upload-time = "2026-07-22T22:18:08.307Z" }, - { url = "https://files.pythonhosted.org/packages/68/56/fbe81c09195924d8b7b8d4461a20458fe80a6a5ed6b24f0314da684277e1/pandas-3.0.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c2e26bb46934b8a2ca0c3de1d3d606fc5f6746584791b2db264d58cf370e08dc", size = 9957095, upload-time = "2026-07-22T22:18:10.6Z" }, - { url = "https://files.pythonhosted.org/packages/e0/51/fac252f4a913ed5eabf3c11b880a9e8d5a6c10f0b2129d0462212d238b4d/pandas-3.0.5-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:73fa87b08a7ef706f8aafda39ddaccf2a99047bea62d8c88a0361bcafb2237bc", size = 10485458, upload-time = "2026-07-22T22:18:12.834Z" }, - { url = "https://files.pythonhosted.org/packages/12/98/e976540c1addf70442be7842a18cf70884a964abbf69442504f4d2939989/pandas-3.0.5-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d373ce03ffd84010ed9839fa73672a9c8256990532e158440c0085db7d914b34", size = 10998091, upload-time = "2026-07-22T22:18:15.209Z" }, - { url = "https://files.pythonhosted.org/packages/a4/8c/1f29b5be8d3fc47dd7567eb167fabba2085879b31e0287ce7cba6d3d2ff4/pandas-3.0.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2a29c53d85ea98c5e792c59ef82ee9fbe6ca902c0d0adb6b23f45ef894cd7bf6", size = 11499501, upload-time = "2026-07-22T22:18:17.689Z" }, - { url = "https://files.pythonhosted.org/packages/9d/e2/bd9c98ad2df7b38bde002adde4cdf353519da51881634323b126c55997f9/pandas-3.0.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a5ad3b02ed6bc7d7ae9b70804b2c6aa31827489d150f8e623ce82491b82085d7", size = 12060559, upload-time = "2026-07-22T22:18:20.147Z" }, - { url = "https://files.pythonhosted.org/packages/f3/9a/ffbd852d58bd74a617fe2f8ee6a58a96982271ce41cf981eab22190b4a4b/pandas-3.0.5-cp312-cp312-pyemscripten_2024_0_wasm32.whl", hash = "sha256:b2acb4650527eec6822c3dadb2b771277b65e7dae7a267d4bccf65fd1bb3fbce", size = 7197652, upload-time = "2026-07-22T22:18:22.502Z" }, - { url = "https://files.pythonhosted.org/packages/70/b5/d2d3e9ae73362ba4229651b0ee1455cf78073a1ce585f6ff693782ce263e/pandas-3.0.5-cp312-cp312-win_amd64.whl", hash = "sha256:80a611068e8a3ac23f7398c6c14eb46dc974e5cc9997f653e2dcfd1da74edd41", size = 9831691, upload-time = "2026-07-22T22:18:24.534Z" }, - { url = "https://files.pythonhosted.org/packages/52/51/dea1e89d6a6796b9c43f85a09b484ee03edb8a4c4842e73e200a8c11301c/pandas-3.0.5-cp312-cp312-win_arm64.whl", hash = "sha256:25ff585b972a18ef1fe9ffa3ac6544d9950508aa76832e5147640b6022821e49", size = 9105796, upload-time = "2026-07-22T22:18:27.064Z" }, - { url = "https://files.pythonhosted.org/packages/bf/09/7b95c4a0025227d6f118c4039b423412ac6a982db02864166185d812fbc7/pandas-3.0.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c1c05a767fe8e5b4fe9e1c29806829c582052eaedb9120a3da83ba3f69e24a5b", size = 10385742, upload-time = "2026-07-22T22:18:29.346Z" }, - { url = "https://files.pythonhosted.org/packages/8d/0c/dc78fd8c4da477b4b5e8ad37295af352190d21ef63a9ee1bc071753074cc/pandas-3.0.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:b86765f268b56f7e665b93bce9d5df69dee7f99e595cf8fb839483ab315942a3", size = 9932067, upload-time = "2026-07-22T22:18:31.833Z" }, - { url = "https://files.pythonhosted.org/packages/3e/71/3592c055cf44df9808550f9368ceda80ff2b224d355ef73fe251dcda1802/pandas-3.0.5-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c597ecf5616b5c420372c1d4d4c00dbbfba7398bea857dcc984347e1ea48417b", size = 10466756, upload-time = "2026-07-22T22:18:34.195Z" }, - { url = "https://files.pythonhosted.org/packages/e3/70/4363150359f95b4cb4bcbb34ca23572bb5495749a621a8f3d5a1ddfd293c/pandas-3.0.5-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4b11c36e218331d0387cbe3a0a5f75162357a1d92d57b2b08a336ff94b19b2be", size = 10938525, upload-time = "2026-07-22T22:18:36.81Z" }, - { url = "https://files.pythonhosted.org/packages/f7/d0/317e7a0c67c0e69fa905a0161409397a7dc2d46ff611f6ca4803352c042b/pandas-3.0.5-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cf52e1f61d229496da17dc7ab54acdee627357e7008fd4fecba3d0ba2937fa58", size = 11489303, upload-time = "2026-07-22T22:18:39.287Z" }, - { url = "https://files.pythonhosted.org/packages/f1/8d/36dade89b49e4f9d5cbdbe863772581f98c0c6d78fc39ad4c557f6f2e17e/pandas-3.0.5-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:db172144bb56422bd157812f3b021eacc255451470b31e2c633c349490a1cfee", size = 11989004, upload-time = "2026-07-22T22:18:42.208Z" }, - { url = "https://files.pythonhosted.org/packages/9c/ba/18c4ec8a746e177da05a9e7a7963781d8ea195780724f854601b6ebd6b78/pandas-3.0.5-cp313-cp313-win_amd64.whl", hash = "sha256:0d298e951f23016ce4699951d044ae6418dbc91bf68cefca0f77666fcbb4e5c6", size = 9826896, upload-time = "2026-07-22T22:18:44.539Z" }, - { url = "https://files.pythonhosted.org/packages/de/ec/28a57266b753799a87b8bc79e7887ac6fd981b8c6d2978a0b7e7b6bd708c/pandas-3.0.5-cp313-cp313-win_arm64.whl", hash = "sha256:66266d3442a5e8b3c90274c2b8b230bee42dd1c286bc822cc2f9f2c7e12b883e", size = 9094790, upload-time = "2026-07-22T22:18:47.468Z" }, - { url = "https://files.pythonhosted.org/packages/51/2f/cf6aae281264f4463f0875bcbb15fd2bb6d291cc535187dad1732475e4a9/pandas-3.0.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2f264fc46911cc8131a7322a16199bbf8e353d27c10bb211f5bd0c814324dc36", size = 10390034, upload-time = "2026-07-22T22:18:49.818Z" }, - { url = "https://files.pythonhosted.org/packages/06/ec/5189518c7a7659c4bdcc6b1eb32c46c6f3c86b0661ffd84143d1112c7732/pandas-3.0.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:53730687fcd161883b24e10411c06d6a4c0f2275d2faf3bb2bc25deb4ba8007c", size = 9980065, upload-time = "2026-07-22T22:18:52.249Z" }, - { url = "https://files.pythonhosted.org/packages/ea/f1/598503ce8d7e3c35601e0747ba288c7864baae66380725bc12f13f884dfe/pandas-3.0.5-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:960d3ebcf249f75206899fcd2c6de53f736b7265759ced0d3e559df0b8b709b0", size = 10545532, upload-time = "2026-07-22T22:18:54.813Z" }, - { url = "https://files.pythonhosted.org/packages/fa/de/ceae2adf7034e07e9910299fe412e1819c4f0dd520700a888bcb03625448/pandas-3.0.5-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e94c2c5ca43bd3ca32bf64d32308887b65e5f9bfd8023ea52755107a999f93b", size = 10963120, upload-time = "2026-07-22T22:18:57.42Z" }, - { url = "https://files.pythonhosted.org/packages/66/25/86e0f4451874eb79e688deeebe3c451fec4557f8952005818d800ee8ac7e/pandas-3.0.5-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e819dd5f62966b481a8cb649d3299ebd886a1ea91ed5a99bf7ce77c98d18ab94", size = 11563178, upload-time = "2026-07-22T22:18:59.729Z" }, - { url = "https://files.pythonhosted.org/packages/f3/45/8643daa3b4147e433adfcccefdd0380d3aad79d86b15d8999730fe1944d5/pandas-3.0.5-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:3c5ed2e7c06e91d340dfd091d7934f9bc82e4a36b95f647f090b9d1c9ac649da", size = 12028708, upload-time = "2026-07-22T22:19:02.164Z" }, - { url = "https://files.pythonhosted.org/packages/96/58/ad979ae617615576e8aafd569c9d4b62f1191d896e38f51d66ba06f3b89a/pandas-3.0.5-cp314-cp314-win_amd64.whl", hash = "sha256:cd8f7c6dc98527058ee6264219343f5392240a6f1bfa654fc5d79023020d0c92", size = 9951806, upload-time = "2026-07-22T22:19:04.596Z" }, - { url = "https://files.pythonhosted.org/packages/69/32/7ac03886b304049a9d2625ee88f59af760d8a93bd30ed9239bce7b9869a8/pandas-3.0.5-cp314-cp314-win_arm64.whl", hash = "sha256:5183427f5a8156d480f30333777bc978be93650a49a7c01db26adffe95b31e85", size = 9238297, upload-time = "2026-07-22T22:19:06.836Z" }, - { url = "https://files.pythonhosted.org/packages/be/ed/1d1f2ee5547d5167face2376d11c8b2a4c7bfff5a416ee7a9046891fab1e/pandas-3.0.5-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:303da736987d481074ca720ada325f8bd80c64ebc2d45ed79b29df3aaa4a26ca", size = 10849690, upload-time = "2026-07-22T22:19:09.391Z" }, - { url = "https://files.pythonhosted.org/packages/57/55/17e17152e98fbb0c4b1e562bc65387a2f20a80db0f4a86bf8d3a0e4248d4/pandas-3.0.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:3b2801bbb049d0136f6c213eae02b5fca969384fc2064dd728d8620552aa49da", size = 10509945, upload-time = "2026-07-22T22:19:11.773Z" }, - { url = "https://files.pythonhosted.org/packages/88/90/817d44dbf83facf9556f33576d9af0a241981e7bb5c00606c0bcb5df8dda/pandas-3.0.5-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cce3a9d11d2b1f82c69a27ec1f4948a170e2c403c4bbfa8cca62e3fdebe2ef3a", size = 10392197, upload-time = "2026-07-22T22:19:14.024Z" }, - { url = "https://files.pythonhosted.org/packages/f1/da/889f00c0a6f5aa1545add70abbf01502dff87ab577adb855bd631c54d2f2/pandas-3.0.5-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ef01af4d8dc6cd2c8d6c7736f149574ef93fe043811eeb5e445f2647154b5040", size = 10862726, upload-time = "2026-07-22T22:19:16.351Z" }, - { url = "https://files.pythonhosted.org/packages/bc/98/f1e934fb3c98fce859c6147c6785816c7b5b9ab7821115c5d8c4de9842b9/pandas-3.0.5-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e2759e890db96dfcffdbd9b86c3c2cb6afaf58def482820317e06163ec1066cd", size = 11414864, upload-time = "2026-07-22T22:19:18.981Z" }, - { url = "https://files.pythonhosted.org/packages/fe/be/d448af7d657d82e1888dd8551f79c6d6fb161080b5b9752d84d910ec2319/pandas-3.0.5-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b58b1b39d46a5862e3fb18f50d1a201398619d16a0f9f73f57eea5583cf0e63c", size = 11925105, upload-time = "2026-07-22T22:19:21.515Z" }, - { url = "https://files.pythonhosted.org/packages/29/c1/ccb4238212c8c4f496c584f3044d94e0c030ed8e1d68999db46c91c2242f/pandas-3.0.5-cp314-cp314t-win_amd64.whl", hash = "sha256:1c10461f6eeb35d8f05b6184c65c8b9991663b66c46b1d559b682cb34ae7c6ea", size = 10387612, upload-time = "2026-07-22T22:19:24.257Z" }, - { url = "https://files.pythonhosted.org/packages/d2/cf/6a51b2c38980e04c279fd2fa908a1b0982064e860444acfca4ec2e2c8359/pandas-3.0.5-cp314-cp314t-win_arm64.whl", hash = "sha256:3c5015fd1730fbf883647e88068176c839c102cea883ba1769a6f4593bfc1f8c", size = 9509776, upload-time = "2026-07-22T22:19:26.694Z" }, -] - [[package]] name = "pluggy" version = "1.6.0" @@ -721,6 +687,15 @@ crypto = [ { name = "cryptography" }, ] +[[package]] +name = "pyproject-hooks" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/82/28175b2414effca1cdac8dc99f76d660e7a4fb0ceefa4b4ab8f5f6742925/pyproject_hooks-1.2.0.tar.gz", hash = "sha256:1e859bd5c40fae9448642dd871adf459e5e2084186e8d2c2a79a824c970da1f8", size = 19228, upload-time = "2024-09-29T09:24:13.293Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bd/24/12818598c362d7f300f18e74db45963dbcb85150324092410c8b49405e42/pyproject_hooks-1.2.0-py3-none-any.whl", hash = "sha256:9e5c6bfa8dcc30091c74b0cf803c81fdd29d94f01992a7707bc97babb1141913", size = 10216, upload-time = "2024-09-29T09:24:11.978Z" }, +] + [[package]] name = "pyright" version = "1.1.411" @@ -750,18 +725,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, ] -[[package]] -name = "python-dateutil" -version = "2.9.0.post0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "six" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } -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" @@ -947,57 +910,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/3e/4132e539aed78c148854d4997a2685b0ed4dc4e87110b59ce528564e184e/ruff-0.16.3-py3-none-win_arm64.whl", hash = "sha256:b8ca152da82c1acc1fa8d5874b15951935f0eef46f10e6954c83859011b6178a", size = 11399302, upload-time = "2026-08-13T15:17:10.908Z" }, ] -[[package]] -name = "scipy" -version = "1.18.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6a/19/ca10ead60b0acc80b2b833c2c4a4f2ff753d0f58b811f70d911c7e94a25c/scipy-1.18.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:7bd21faaf5a1a3b2eff922d02db5f191b99a6518db9078a8fb23169f6d22259a", size = 31056519, upload-time = "2026-06-19T14:59:45.203Z" }, - { url = "https://files.pythonhosted.org/packages/96/72/1e6442a00cd2924d361aa1b642ab6373ec35c6fabf311a760be9f76e0f13/scipy-1.18.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:265915e79107de9f946b855e50d7470d5893ec3f54b342e1aa6201cbdcd8bb6b", size = 28681889, upload-time = "2026-06-19T14:59:48.103Z" }, - { url = "https://files.pythonhosted.org/packages/9b/2d/11dd93d21e147a73ba22bd75c0b9208d3a2e0ec76d53170ce7d9029b1015/scipy-1.18.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:9ab7b758be6940954a713ee466e2043e9f6e2ed965c1fce5c91039f4be3d90a9", size = 20423580, upload-time = "2026-06-19T14:59:50.665Z" }, - { url = "https://files.pythonhosted.org/packages/9c/01/93552f75e0d2a7dd115a45e59209c51e8d514daff02fc887d2623be06fe1/scipy-1.18.0-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:97b6cddaaee0a779ef6b5ca83c9604b27cc16b2b8fc22c142652df8793319fb8", size = 23054441, upload-time = "2026-06-19T14:59:53.564Z" }, - { url = "https://files.pythonhosted.org/packages/3c/23/21f5e703643d66f21faa6b4c73195bfcad70c55efcb4f1ab327cd7c4101a/scipy-1.18.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:52a96e21517c7292375c0e27dd796a811f03fcea5fd4d108fdfea8145dcf17ab", size = 33968720, upload-time = "2026-06-19T14:59:56.415Z" }, - { url = "https://files.pythonhosted.org/packages/dd/aa/1b939f6c67ed68635bb538e6752d3dacc02f66535182e939a89581a44e9c/scipy-1.18.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1f55797419e16e7f30cf88ffb3113ce0467f00cfe3f70d5c281730b21769bfc2", size = 35287115, upload-time = "2026-06-19T14:59:59.411Z" }, - { url = "https://files.pythonhosted.org/packages/b6/ff/eec46be7e9234208f801062b53e1983085eddebd693f6c9bfb03b459830d/scipy-1.18.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ad033410e2e0672ffdc1042110cef20e1c46f8fd0616cee1d44d8d58fad8fc11", size = 35577989, upload-time = "2026-06-19T15:00:02.235Z" }, - { url = "https://files.pythonhosted.org/packages/84/ca/210d4759c7210bb7d269437421959b39a33434e2776b60c5cb8a763bb30a/scipy-1.18.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4a55985d54c769c872e64b7f4c8a81cc30ef700cc04296abbbf3705439c126de", size = 37421717, upload-time = "2026-06-19T15:00:05.102Z" }, - { url = "https://files.pythonhosted.org/packages/2b/54/9a9edb45345bd6744da5ddfb6628e5d5185920494c6a67ec45b6381004cb/scipy-1.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:71ccc8faa2dd16ac310233203474a8b5cb67f10dedd54a3116d34943f4b19132", size = 36597428, upload-time = "2026-06-19T15:00:08.112Z" }, - { url = "https://files.pythonhosted.org/packages/99/0e/33f32a2a58987e26aec0f7df252cbbad1e90ae77bdbc76f40dd4ed0cf0ea/scipy-1.18.0-cp312-cp312-win_arm64.whl", hash = "sha256:d88363fd9d8fbd3511bd273f1a49efb2a540773ddf92a91d57498ce7dd7f3e76", size = 24351481, upload-time = "2026-06-19T15:00:11.103Z" }, - { url = "https://files.pythonhosted.org/packages/05/52/9c0136c2de7ae0779b7b366447766cec6d9f0702c56bb8ffeb04c8fd3af4/scipy-1.18.0-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:09143f676d157d9f546d663504ef9c1becb819824f1afc018814176411942446", size = 31036107, upload-time = "2026-06-19T15:00:14.03Z" }, - { url = "https://files.pythonhosted.org/packages/02/73/0291a64843270f4efb86cdcf2ee0f2048631b65ec6b405398b2b4dbf11bf/scipy-1.18.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5efe260f69417b97ddae455bfb5a95e8359f7f66ad7fa9522a60feb66f169520", size = 28663303, upload-time = "2026-06-19T15:00:16.819Z" }, - { url = "https://files.pythonhosted.org/packages/d3/0f/10ffa0b697a572f4e0d48b92a88895d366422f019f723e7e14a84c050dac/scipy-1.18.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:68363b7eaacd8b5dd426df56d782cc156468ac79a127a1b87ca597d6e2e82197", size = 20404960, upload-time = "2026-06-19T15:00:19.635Z" }, - { url = "https://files.pythonhosted.org/packages/7e/d2/e896cea21ba8edd6c81d4c55b1ffcc717e79698dcbebf9641b4cfb4c6622/scipy-1.18.0-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:c5557d8be5da8e41353fcd4d21491fdbab83b062fc579e94dc09a7c8ab4f669b", size = 23034074, upload-time = "2026-06-19T15:00:22.107Z" }, - { url = "https://files.pythonhosted.org/packages/ea/b2/e83ea34279a52c03374477c74006256ec78df65fc877baa4617d6de1d202/scipy-1.18.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0d13bca67c096d89fb95ced0d8921807300fce0275643aef9533cc63a0773468", size = 33942038, upload-time = "2026-06-19T15:00:24.964Z" }, - { url = "https://files.pythonhosted.org/packages/f6/af/e8fe5fb136f51e2b01678b92cb4106d10d8cd68ec147ead2e7cb0ac75398/scipy-1.18.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a46f9273dbd0eb1cefba61c9b8648b4dfe3cbc14a080176f9a73e44b8336dc7f", size = 35266390, upload-time = "2026-06-19T15:00:28.059Z" }, - { url = "https://files.pythonhosted.org/packages/3a/49/2c5cbb907b56695fc67517811d1db234dfd83381a84814ec220aded2794d/scipy-1.18.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5aba46108853ddfc77906b6557aac839d2b52e900c1d72a1180adaaab58d265f", size = 35551324, upload-time = "2026-06-19T15:00:31.014Z" }, - { url = "https://files.pythonhosted.org/packages/bb/73/eda39f7a2d306ff0ffc574afd13c0bbb6d10a603d9a413998ee269487a80/scipy-1.18.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b6f758e35f12757b5d95c00bc6de2438e229c2664b7a92e96f205959d9f2dfa4", size = 37404785, upload-time = "2026-06-19T15:00:34.072Z" }, - { url = "https://files.pythonhosted.org/packages/b7/d2/ae881ee28d014f38e0ccbfd974a06a919ba9af34f1f74bf42b5301891d63/scipy-1.18.0-cp313-cp313-win_amd64.whl", hash = "sha256:1afac4a847207c7ff8efd321734a50b06d0280b3b2a2c0fc2f413101747ad7c7", size = 36554943, upload-time = "2026-06-19T15:00:36.903Z" }, - { url = "https://files.pythonhosted.org/packages/70/3a/21154e2d54eb3639c6bf4dbae2e531c68356bfe95990daa30df33b30d556/scipy-1.18.0-cp313-cp313-win_arm64.whl", hash = "sha256:c5dbddf60e58c2312316d097271a8e73d40eaf2eabfa4d95ed7d3695bbf2ce7b", size = 24350911, upload-time = "2026-06-19T15:00:40.062Z" }, - { url = "https://files.pythonhosted.org/packages/78/b5/915a19b3de2f7430062b509653563db1633ddbb6f021b06731521115d4e2/scipy-1.18.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:4c256ee70c0d1a8a2ace807e199ccd4e3f57037433842abb3fb36bc17eaa9578", size = 31036253, upload-time = "2026-06-19T15:00:43.216Z" }, - { url = "https://files.pythonhosted.org/packages/d7/88/b72def7262e150d16be13fca37a96481138d624e700340bc3362a7588929/scipy-1.18.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:2ef3abc54a4ffc53765374b0d5728532dfdd2585ed23f6b11c206a1f0b1b9af8", size = 28673758, upload-time = "2026-06-19T15:00:46.663Z" }, - { url = "https://files.pythonhosted.org/packages/91/02/2e636a61a525632c373cf6a9c24442a3ffb79e364d38e98b32042964ac32/scipy-1.18.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f2a6af57bd9e4a75d70e4117e78a1bbee84f79ae3fbb6d0111005d6ebcc4cb8d", size = 20415514, upload-time = "2026-06-19T15:00:49.399Z" }, - { url = "https://files.pythonhosted.org/packages/c9/b6/2135974442f6aba159d9d39d774a1c8cb19947016725d69fecc685df45bf/scipy-1.18.0-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:3f1ac564d3bf6c03d861d2cd87a1bea0da2887136f7fb1bf519c05a8971452d6", size = 23034398, upload-time = "2026-06-19T15:00:51.941Z" }, - { url = "https://files.pythonhosted.org/packages/f6/e6/ba89ec5abf6ee9257c0d1ec985573f3ae32742c24bc03e016388a40b1b15/scipy-1.18.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40395a5fcd1abee49a5c7aaa98c29db393eedc835138560a588c47ec16156690", size = 33998032, upload-time = "2026-06-19T15:00:54.838Z" }, - { url = "https://files.pythonhosted.org/packages/7f/c4/bc41eb19b0fd0db868f4132920879019318d80cc522ad8f2bca4611af808/scipy-1.18.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ca01e8ae69f1b18e9a58d91afead31be3cef0dd905a10249dac559ee15460a0", size = 35283333, upload-time = "2026-06-19T15:00:58.152Z" }, - { url = "https://files.pythonhosted.org/packages/53/a4/cbdeef6eb3830a8462a9d4ada814de5fc984345cc9ecf17cbec51a036f1e/scipy-1.18.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7a7f3b01647384dbc3a711e8c6778e0aabbe93959249fef5c7393396bcac0867", size = 35610216, upload-time = "2026-06-19T15:01:01.155Z" }, - { url = "https://files.pythonhosted.org/packages/80/4d/b2b82502b65f661d1b789c1665dcdf315d5f12194e06fc0b37946294ebae/scipy-1.18.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:6aa94e78ec192a30063a5e72e561c28af769dc311190b24fe91774eff1969709", size = 37418960, upload-time = "2026-06-19T15:01:04.155Z" }, - { url = "https://files.pythonhosted.org/packages/93/3e/902d836831474b0ab5a37d16404f7bc5fafd9efba632890e271ba952635f/scipy-1.18.0-cp314-cp314-win_amd64.whl", hash = "sha256:2d8bbdc6c817f5b4006a54d799d4f5bab6f910193cbb9a1ff310833d4d270f61", size = 37288845, upload-time = "2026-06-19T15:01:07.822Z" }, - { url = "https://files.pythonhosted.org/packages/b6/43/8d73b337a3bdb14daa0314f0434210747c02d79d729ce1777574a817dcf6/scipy-1.18.0-cp314-cp314-win_arm64.whl", hash = "sha256:18e9575f1569b2c54174e6159d32942e03731177f63dce7975f0a0c88d102f5b", size = 24988971, upload-time = "2026-06-19T15:01:11.076Z" }, - { url = "https://files.pythonhosted.org/packages/b4/b4/f11918b0508a2787031a0499a03fbe3546f3bb5ca05d01038c45b278c09a/scipy-1.18.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:f351e0dd702687d12a402b867a1b4146a256923e1c38317cbc472f6372b94707", size = 31399325, upload-time = "2026-06-19T15:01:13.723Z" }, - { url = "https://files.pythonhosted.org/packages/7b/d1/1f287b57c0ff0ee5185dff3946d92c8017d39b0e431f0ae79a3ff1859512/scipy-1.18.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:7c7a51b33ce387193c97f228320cf8e87361daa1bba750638677729598b3e677", size = 29092110, upload-time = "2026-06-19T15:01:16.908Z" }, - { url = "https://files.pythonhosted.org/packages/ff/1a/7b74eb6c392fdcb27d414c0e7558a6d0231eb3b6d73571f479bb81ea8794/scipy-1.18.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:84031d7b052a54fae2f8632e0ec802073d385476eb9a63079bce6e23ef9283d4", size = 20833811, upload-time = "2026-06-19T15:01:20.488Z" }, - { url = "https://files.pythonhosted.org/packages/7c/ad/f3941716320a7b9cb4d68734a903b45fe16eff5fb7da7e16f2e619304979/scipy-1.18.0-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:56abf29a7c067dde59be8b9a22d606a4ea1b2f2a4b756d9d903c62818f5dacce", size = 23396644, upload-time = "2026-06-19T15:01:23.364Z" }, - { url = "https://files.pythonhosted.org/packages/22/22/1446b62ffe07f9719b7d9b1b6a4e05a772833ae8f441fe4c22c34c9b250f/scipy-1.18.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ad44305cfa24b1ba5803cbbebf033590ccbac1aa5d612d727b785325ab408b0", size = 34079318, upload-time = "2026-06-19T15:01:26.002Z" }, - { url = "https://files.pythonhosted.org/packages/56/3b/b87da667098bb470fa30c7011b0ba351ee976dd395c78798c66e941665a3/scipy-1.18.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:945c1761b93f38d7f99ae81ae80c63e621471608c7eeead563f6df025585cd58", size = 35324320, upload-time = "2026-06-19T15:01:28.881Z" }, - { url = "https://files.pythonhosted.org/packages/f8/a1/c7932f91909759b0267f75fdea34e91309f96b895757534b76a90b6b4344/scipy-1.18.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1a4441f15d620578772a49e5ab48c0ee1f7a0220e387110283062729136b2553", size = 35699541, upload-time = "2026-06-19T15:01:31.968Z" }, - { url = "https://files.pythonhosted.org/packages/f7/86/5185061a1fcc41d18c5dc2463969b3a3964b31d9ac67b2fb05d4c7ff7670/scipy-1.18.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:9aac6192fac56bf2ca534389d24623f07b39ff83317d58287285e7fbd622ff76", size = 37472480, upload-time = "2026-06-19T15:01:35.136Z" }, - { url = "https://files.pythonhosted.org/packages/31/8e/f04c68e39919a010d34f2ee1367fd705b0a25a02f609d755f0bfbc0a15fc/scipy-1.18.0-cp314-cp314t-win_amd64.whl", hash = "sha256:e40baea28ae7f5475c779741e2d90b1247c78531207b49c7030e698ff81cee3f", size = 37365390, upload-time = "2026-06-19T15:01:38.091Z" }, - { url = "https://files.pythonhosted.org/packages/d5/19/969dc072906c84dd0a3b05dcf57ea750936087d7873549e408b35cfc3f97/scipy-1.18.0-cp314-cp314t-win_arm64.whl", hash = "sha256:368e0a705903c466aa5f08eefb39e6b1b6b2d659e7352a31fd9e2438365be0f8", size = 25279661, upload-time = "2026-06-19T15:01:40.817Z" }, -] - [[package]] name = "shellingham" version = "1.5.4" @@ -1007,15 +919,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, ] -[[package]] -name = "six" -version = "1.17.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, -] - [[package]] name = "sse-starlette" version = "3.4.8" @@ -1087,15 +990,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/67/81/4add07e5172b7ac40d8ed5ff580409a7801a4fe26d529bdd915401dabfbe/typing_inspection-0.4.4-py3-none-any.whl", hash = "sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147", size = 14750, upload-time = "2026-08-12T12:37:24.648Z" }, ] -[[package]] -name = "tzdata" -version = "2026.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/92/ff/5a28bdfd8c3ebec42564ac7d0e54ca3db65044a9314a97f9564fa7a1e926/tzdata-2026.3.tar.gz", hash = "sha256:4a1518b8993086a7982523e071643f3c0e5f213e75b21318e78bcabfff9d1415", size = 198674, upload-time = "2026-07-10T08:50:37.887Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e5/6d/b53b99a9f2766d095985947a5782f1702cabb129a34f7a802d7197af832f/tzdata-2026.3-py2.py3-none-any.whl", hash = "sha256:dc096730c87af6cab1b171c9d532be840741ff5d459015e7f6947bd7d7e54931", size = 348168, upload-time = "2026-07-10T08:50:36.46Z" }, -] - [[package]] name = "uvicorn" version = "0.52.3" From edea49b31ea75eee5badae8d580382844a23eaf9 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:40:52 +0300 Subject: [PATCH 15/16] feat: define supported public library API --- README.md | 33 +++++++++++++++++++++++++++--- src/flowsense/__init__.py | 43 +++++++++++++++++++++++++++++++++++++++ tests/test_public_api.py | 35 +++++++++++++++++++++++++++++++ 3 files changed, 108 insertions(+), 3 deletions(-) create mode 100644 tests/test_public_api.py diff --git a/README.md b/README.md index b769207..2d3ba19 100644 --- a/README.md +++ b/README.md @@ -135,6 +135,29 @@ Then run: flowsense analyze ``` +## Library API + +FlowSense can also be used as a Python library through its supported top-level +API: + +```python +from flowsense import AirflowClient, analyze_dag + +with AirflowClient() as source: + analysis = analyze_dag( + dag_id="flowsense_demo", + source=source, + ) + +print(analysis.overall_severity) +print(analysis.primary_origin) +``` + +Custom data sources can implement the `DAGDataSource` protocol and be passed to +`analyze_dag`. Names exported directly from `flowsense` form the supported public +API. Imports from internal packages such as `flowsense.engine` should be treated +as implementation details and may change before version 1.0. + ## MCP Server Start the FlowSense MCP server over stdio: @@ -183,8 +206,8 @@ ruff format . ```text src/flowsense/ -├── cli/ -├── collector/ +├── application/ +├── domain/ ├── engine/ │ ├── drift.py │ ├── history.py @@ -192,8 +215,12 @@ src/flowsense/ │ ├── propagation.py │ ├── root_cause.py │ └── timing.py +├── infrastructure/ +│ └── airflow/ +├── cli/ ├── mcp/ -└── models/ +├── collector/ # backward-compatible imports +└── models/ # backward-compatible imports ``` ## Detection Approach diff --git a/src/flowsense/__init__.py b/src/flowsense/__init__.py index e69de29..0a6dc9b 100644 --- a/src/flowsense/__init__.py +++ b/src/flowsense/__init__.py @@ -0,0 +1,43 @@ +from importlib.metadata import PackageNotFoundError, version + +from flowsense.application import DAGDataSource, analyze_dag +from flowsense.domain import ( + AnalysisDiagnostic, + DAGAnalysis, + DriftResult, + FlowSenseError, + ImpactClassification, + InsufficientHistoryError, + InvalidTaskTimingError, + PropagationResult, + RootCauseResult, + Severity, + TaskImpact, + TaskRun, +) +from flowsense.infrastructure.airflow import AirflowApiError, AirflowClient + +try: + __version__ = version("flowsense") +except PackageNotFoundError: + __version__ = "0.0.0" + +__all__ = [ + "AirflowApiError", + "AirflowClient", + "AnalysisDiagnostic", + "DAGAnalysis", + "DAGDataSource", + "DriftResult", + "FlowSenseError", + "ImpactClassification", + "InsufficientHistoryError", + "InvalidTaskTimingError", + "PropagationResult", + "RootCauseResult", + "Severity", + "TaskImpact", + "TaskRun", + "__version__", + "analyze_dag", +] diff --git a/tests/test_public_api.py b/tests/test_public_api.py new file mode 100644 index 0000000..485eefd --- /dev/null +++ b/tests/test_public_api.py @@ -0,0 +1,35 @@ +import flowsense + + +class EmptyDataSource: + def collect_task_runs(self, dag_id: str) -> list[flowsense.TaskRun]: + return [] + + def get_dag_dependencies(self, dag_id: str) -> dict[str, list[str]]: + return {} + + +def test_top_level_api_analyzes_custom_data_source() -> None: + analysis = flowsense.analyze_dag( + dag_id="demo", + source=EmptyDataSource(), + ) + + assert analysis.dag_id == "demo" + assert analysis.overall_severity is flowsense.Severity.NORMAL + + +def test_top_level_api_declares_supported_exports() -> None: + expected_exports = { + "AirflowApiError", + "AirflowClient", + "DAGAnalysis", + "DAGDataSource", + "Severity", + "TaskRun", + "__version__", + "analyze_dag", + } + + assert expected_exports <= set(flowsense.__all__) + assert flowsense.__version__ From 30e392f5adc476be2acb40170f52e6caafdd6845 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 23:47:20 +0300 Subject: [PATCH 16/16] feat: add configurable analysis policy --- README.md | 19 +++++++ src/flowsense/__init__.py | 6 ++ src/flowsense/application/analyzer.py | 11 +++- src/flowsense/cli/main.py | 25 +++++++++ src/flowsense/domain/__init__.py | 6 +- src/flowsense/domain/enums.py | 6 ++ src/flowsense/domain/policy.py | 31 +++++++++++ src/flowsense/domain/results.py | 2 + src/flowsense/engine/drift.py | 30 +++++++--- src/flowsense/engine/history.py | 19 ++++--- src/flowsense/engine/timing.py | 10 +++- src/flowsense/mcp/server.py | 32 ++++++++++- tests/test_policy.py | 79 +++++++++++++++++++++++++++ tests/test_public_api.py | 2 + 14 files changed, 258 insertions(+), 20 deletions(-) create mode 100644 src/flowsense/domain/policy.py create mode 100644 tests/test_policy.py diff --git a/README.md b/README.md index 2d3ba19..776d47b 100644 --- a/README.md +++ b/README.md @@ -153,6 +153,25 @@ print(analysis.overall_severity) print(analysis.primary_origin) ``` +Analysis behavior can be customized with an immutable policy: + +```python +from flowsense import AnalysisPolicy, MappedTaskAggregation + +policy = AnalysisPolicy( + minimum_history=10, + baseline_window=30, + medium_threshold=2.5, + high_threshold=4.0, + critical_threshold=6.0, + mapped_task_aggregation=MappedTaskAggregation.MAX, +) +``` + +`baseline_window` limits the number of historical values used before the current +run. Mapped task durations can be aggregated with `MAX`, `MEAN`, or `SUM`. The +same policy options are available through the CLI and MCP tool. + Custom data sources can implement the `DAGDataSource` protocol and be passed to `analyze_dag`. Names exported directly from `flowsense` form the supported public API. Imports from internal packages such as `flowsense.engine` should be treated diff --git a/src/flowsense/__init__.py b/src/flowsense/__init__.py index 0a6dc9b..b5654b3 100644 --- a/src/flowsense/__init__.py +++ b/src/flowsense/__init__.py @@ -2,13 +2,16 @@ from flowsense.application import DAGDataSource, analyze_dag from flowsense.domain import ( + DEFAULT_ANALYSIS_POLICY, AnalysisDiagnostic, + AnalysisPolicy, DAGAnalysis, DriftResult, FlowSenseError, ImpactClassification, InsufficientHistoryError, InvalidTaskTimingError, + MappedTaskAggregation, PropagationResult, RootCauseResult, Severity, @@ -23,9 +26,11 @@ __version__ = "0.0.0" __all__ = [ + "DEFAULT_ANALYSIS_POLICY", "AirflowApiError", "AirflowClient", "AnalysisDiagnostic", + "AnalysisPolicy", "DAGAnalysis", "DAGDataSource", "DriftResult", @@ -33,6 +38,7 @@ "ImpactClassification", "InsufficientHistoryError", "InvalidTaskTimingError", + "MappedTaskAggregation", "PropagationResult", "RootCauseResult", "Severity", diff --git a/src/flowsense/application/analyzer.py b/src/flowsense/application/analyzer.py index b1ae8ad..0793520 100644 --- a/src/flowsense/application/analyzer.py +++ b/src/flowsense/application/analyzer.py @@ -2,7 +2,9 @@ from flowsense.application.ports import DAGDataSource from flowsense.domain import ( + DEFAULT_ANALYSIS_POLICY, AnalysisDiagnostic, + AnalysisPolicy, DAGAnalysis, InsufficientHistoryError, Severity, @@ -22,9 +24,13 @@ def analyze_dag( dag_id: str, source: DAGDataSource, + policy: AnalysisPolicy = DEFAULT_ANALYSIS_POLICY, ) -> DAGAnalysis: task_runs = source.collect_task_runs(dag_id) - duration_history = build_duration_history(task_runs) + duration_history = build_duration_history( + task_runs, + aggregation=policy.mapped_task_aggregation, + ) diagnostics: list[AnalysisDiagnostic] = [] drift_results = {} @@ -34,6 +40,7 @@ def analyze_dag( drift_results[task_id] = calculate_drift( task_id=task_id, durations=durations, + policy=policy, ) except InsufficientHistoryError as exc: diagnostics.append( @@ -71,6 +78,7 @@ def analyze_dag( upstream_task=upstream_task, downstream_task=downstream_task, handoff_delays=delays, + policy=policy, ) except InsufficientHistoryError as exc: diagnostics.append( @@ -134,4 +142,5 @@ def analyze_dag( propagation_results=propagation_results, dependencies=dependencies, diagnostics=diagnostics, + policy=policy, ) diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index 5ba22e8..8a9a5aa 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -1,10 +1,13 @@ from __future__ import annotations +from typing import Annotated + import typer from rich.console import Console from rich.table import Table from flowsense.application import analyze_dag +from flowsense.domain import AnalysisPolicy, MappedTaskAggregation from flowsense.infrastructure.airflow import AirflowApiError, AirflowClient app = typer.Typer( @@ -27,12 +30,34 @@ def analyze( ..., help="Airflow DAG id to analyze.", ), + minimum_history: int = typer.Option(5, min=2), + baseline_window: int | None = typer.Option(None, min=1), + medium_threshold: float = typer.Option(2.0, min=0.0), + high_threshold: float = typer.Option(3.5, min=0.0), + critical_threshold: float = typer.Option(5.0, min=0.0), + mapped_task_aggregation: Annotated[ + MappedTaskAggregation, + typer.Option(), + ] = MappedTaskAggregation.MAX, ) -> None: + try: + policy = AnalysisPolicy( + minimum_history=minimum_history, + baseline_window=baseline_window, + medium_threshold=medium_threshold, + high_threshold=high_threshold, + critical_threshold=critical_threshold, + mapped_task_aggregation=mapped_task_aggregation, + ) + except ValueError as exc: + raise typer.BadParameter(str(exc)) from exc + try: with AirflowClient() as source: analysis = analyze_dag( dag_id=dag_id, source=source, + policy=policy, ) except AirflowApiError as exc: console.print(f"[bold red]Airflow request failed:[/bold red] {exc}") diff --git a/src/flowsense/domain/__init__.py b/src/flowsense/domain/__init__.py index 8eccfe1..a3128b2 100644 --- a/src/flowsense/domain/__init__.py +++ b/src/flowsense/domain/__init__.py @@ -1,10 +1,11 @@ -from flowsense.domain.enums import ImpactClassification, Severity +from flowsense.domain.enums import ImpactClassification, MappedTaskAggregation, Severity from flowsense.domain.exceptions import ( FlowSenseError, InsufficientHistoryError, InvalidTaskTimingError, ) from flowsense.domain.models import TaskRun +from flowsense.domain.policy import DEFAULT_ANALYSIS_POLICY, AnalysisPolicy from flowsense.domain.results import ( AnalysisDiagnostic, DAGAnalysis, @@ -15,13 +16,16 @@ ) __all__ = [ + "DEFAULT_ANALYSIS_POLICY", "AnalysisDiagnostic", + "AnalysisPolicy", "DAGAnalysis", "DriftResult", "FlowSenseError", "ImpactClassification", "InsufficientHistoryError", "InvalidTaskTimingError", + "MappedTaskAggregation", "PropagationResult", "RootCauseResult", "Severity", diff --git a/src/flowsense/domain/enums.py b/src/flowsense/domain/enums.py index c86a7ab..5b5c974 100644 --- a/src/flowsense/domain/enums.py +++ b/src/flowsense/domain/enums.py @@ -15,6 +15,12 @@ class ImpactClassification(StrEnum): COMBINED = "COMBINED" +class MappedTaskAggregation(StrEnum): + MAX = "MAX" + MEAN = "MEAN" + SUM = "SUM" + + SEVERITY_SCORE: dict[Severity, int] = { Severity.NORMAL: 0, Severity.MEDIUM: 1, diff --git a/src/flowsense/domain/policy.py b/src/flowsense/domain/policy.py new file mode 100644 index 0000000..e24d710 --- /dev/null +++ b/src/flowsense/domain/policy.py @@ -0,0 +1,31 @@ +from dataclasses import dataclass + +from flowsense.domain.enums import MappedTaskAggregation + + +@dataclass(frozen=True) +class AnalysisPolicy: + minimum_history: int = 5 + baseline_window: int | None = None + medium_threshold: float = 2.0 + high_threshold: float = 3.5 + critical_threshold: float = 5.0 + mapped_task_aggregation: MappedTaskAggregation = MappedTaskAggregation.MAX + + def __post_init__(self) -> None: + if self.minimum_history < 2: + raise ValueError("minimum_history must be at least 2.") + if self.baseline_window is not None: + if self.baseline_window < 1: + raise ValueError("baseline_window must be at least 1.") + if self.baseline_window < self.minimum_history - 1: + raise ValueError( + "baseline_window must contain enough values for minimum_history." + ) + if not ( + 0 < self.medium_threshold < self.high_threshold < self.critical_threshold + ): + raise ValueError("Severity thresholds must be positive and increasing.") + + +DEFAULT_ANALYSIS_POLICY = AnalysisPolicy() diff --git a/src/flowsense/domain/results.py b/src/flowsense/domain/results.py index df4ef7e..9a7f41e 100644 --- a/src/flowsense/domain/results.py +++ b/src/flowsense/domain/results.py @@ -3,6 +3,7 @@ from dataclasses import dataclass, field from flowsense.domain.enums import ImpactClassification, Severity +from flowsense.domain.policy import DEFAULT_ANALYSIS_POLICY, AnalysisPolicy @dataclass(frozen=True) @@ -59,3 +60,4 @@ class DAGAnalysis: propagation_results: list[PropagationResult] dependencies: dict[str, list[str]] diagnostics: list[AnalysisDiagnostic] = field(default_factory=list) + policy: AnalysisPolicy = DEFAULT_ANALYSIS_POLICY diff --git a/src/flowsense/engine/drift.py b/src/flowsense/engine/drift.py index aabea8c..92b38ca 100644 --- a/src/flowsense/engine/drift.py +++ b/src/flowsense/engine/drift.py @@ -4,22 +4,33 @@ import numpy as np -from flowsense.domain import DriftResult, InsufficientHistoryError, Severity +from flowsense.domain import ( + DEFAULT_ANALYSIS_POLICY, + AnalysisPolicy, + DriftResult, + InsufficientHistoryError, + Severity, +) def calculate_drift( task_id: str, durations: list[float], + policy: AnalysisPolicy = DEFAULT_ANALYSIS_POLICY, ) -> DriftResult: - if len(durations) < 5: + if len(durations) < policy.minimum_history: raise InsufficientHistoryError( subject_id=task_id, - required=5, + required=policy.minimum_history, actual=len(durations), ) + baseline_durations = durations[:-1] + if policy.baseline_window is not None: + baseline_durations = baseline_durations[-policy.baseline_window :] + baseline_values = np.array( - durations[:-1], + baseline_durations, dtype=float, ) @@ -35,7 +46,10 @@ def calculate_drift( if math.isclose(current, median): robust_z_score = 0.0 else: - robust_z_score = math.copysign(5.0, current - median) + robust_z_score = math.copysign( + policy.critical_threshold, + current - median, + ) else: robust_z_score = 0.6745 * (current - median) / mad @@ -46,11 +60,11 @@ def calculate_drift( absolute_z = abs(robust_z_score) - if absolute_z >= 5: + if absolute_z >= policy.critical_threshold: severity = Severity.CRITICAL - elif absolute_z >= 3.5: + elif absolute_z >= policy.high_threshold: severity = Severity.HIGH - elif absolute_z >= 2: + elif absolute_z >= policy.medium_threshold: severity = Severity.MEDIUM else: severity = Severity.NORMAL diff --git a/src/flowsense/engine/history.py b/src/flowsense/engine/history.py index 612b60e..4baa6eb 100644 --- a/src/flowsense/engine/history.py +++ b/src/flowsense/engine/history.py @@ -2,28 +2,33 @@ from collections import defaultdict -from flowsense.domain import TaskRun +from flowsense.domain import MappedTaskAggregation, TaskRun def build_duration_history( task_runs: list[TaskRun], + aggregation: MappedTaskAggregation = MappedTaskAggregation.MAX, ) -> dict[str, list[float]]: """Build logical-task history using the slowest mapped instance per DAG run.""" - durations_by_run_and_task: dict[tuple[str, str], float] = {} + grouped_durations: dict[tuple[str, str], list[float]] = defaultdict(list) for task_run in task_runs: if task_run.state != "success" or task_run.duration is None: continue key = (task_run.dag_run_id, task_run.task_id) - current_duration = durations_by_run_and_task.get(key) + grouped_durations[key].append(task_run.duration) - if current_duration is None or task_run.duration > current_duration: - durations_by_run_and_task[key] = task_run.duration + def aggregate(values: list[float]) -> float: + if aggregation is MappedTaskAggregation.SUM: + return sum(values) + if aggregation is MappedTaskAggregation.MEAN: + return sum(values) / len(values) + return max(values) history: dict[str, list[float]] = defaultdict(list) - for (_dag_run_id, task_id), duration in durations_by_run_and_task.items(): - history[task_id].append(duration) + for (_dag_run_id, task_id), durations in grouped_durations.items(): + history[task_id].append(aggregate(durations)) return dict(history) diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py index 9db3d19..bba650d 100644 --- a/src/flowsense/engine/timing.py +++ b/src/flowsense/engine/timing.py @@ -3,7 +3,13 @@ from dataclasses import dataclass from datetime import datetime -from flowsense.domain import DriftResult, InvalidTaskTimingError, TaskRun +from flowsense.domain import ( + DEFAULT_ANALYSIS_POLICY, + AnalysisPolicy, + DriftResult, + InvalidTaskTimingError, + TaskRun, +) from flowsense.engine.drift import calculate_drift @@ -204,10 +210,12 @@ def calculate_handoff_drift( upstream_task: str, downstream_task: str, handoff_delays: list[float], + policy: AnalysisPolicy = DEFAULT_ANALYSIS_POLICY, ) -> DriftResult: edge_id = f"{upstream_task}->{downstream_task}" return calculate_drift( task_id=edge_id, durations=handoff_delays, + policy=policy, ) diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py index 2ecf0c0..055aeb7 100644 --- a/src/flowsense/mcp/server.py +++ b/src/flowsense/mcp/server.py @@ -3,7 +3,11 @@ from mcp.server import MCPServer from flowsense.application import analyze_dag -from flowsense.domain import DAGAnalysis +from flowsense.domain import ( + AnalysisPolicy, + DAGAnalysis, + MappedTaskAggregation, +) from flowsense.infrastructure.airflow import AirflowApiError, AirflowClient mcp = MCPServer("FlowSense Engine") @@ -14,6 +18,14 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: "dag_id": analysis.dag_id, "runs_analyzed": analysis.runs_analyzed, "overall_severity": analysis.overall_severity, + "policy": { + "minimum_history": analysis.policy.minimum_history, + "baseline_window": analysis.policy.baseline_window, + "medium_threshold": analysis.policy.medium_threshold, + "high_threshold": analysis.policy.high_threshold, + "critical_threshold": analysis.policy.critical_threshold, + "mapped_task_aggregation": analysis.policy.mapped_task_aggregation, + }, "primary_origin": ( { "task_id": analysis.primary_origin.task_id, @@ -79,13 +91,29 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: @mcp.tool() -def analyze_airflow_dag(dag_id: str) -> dict: +def analyze_airflow_dag( + dag_id: str, + minimum_history: int = 5, + baseline_window: int | None = None, + medium_threshold: float = 2.0, + high_threshold: float = 3.5, + critical_threshold: float = 5.0, + mapped_task_aggregation: MappedTaskAggregation = MappedTaskAggregation.MAX, +) -> dict: """Analyze an Apache Airflow DAG for temporal drift and propagation.""" try: with AirflowClient() as source: analysis = analyze_dag( dag_id=dag_id, source=source, + policy=AnalysisPolicy( + minimum_history=minimum_history, + baseline_window=baseline_window, + medium_threshold=medium_threshold, + high_threshold=high_threshold, + critical_threshold=critical_threshold, + mapped_task_aggregation=mapped_task_aggregation, + ), ) except AirflowApiError as exc: raise RuntimeError(str(exc)) from exc diff --git a/tests/test_policy.py b/tests/test_policy.py new file mode 100644 index 0000000..459af2a --- /dev/null +++ b/tests/test_policy.py @@ -0,0 +1,79 @@ +import pytest + +from flowsense import AnalysisPolicy, MappedTaskAggregation +from flowsense.engine.drift import calculate_drift +from flowsense.engine.history import build_duration_history +from flowsense.models import TaskRun + + +def test_policy_validates_history_and_thresholds() -> None: + with pytest.raises(ValueError): + AnalysisPolicy(minimum_history=1) + + with pytest.raises(ValueError): + AnalysisPolicy( + medium_threshold=3.5, + high_threshold=2.0, + ) + + with pytest.raises(ValueError): + AnalysisPolicy(minimum_history=5, baseline_window=3) + + +def test_drift_uses_configured_minimum_history_and_thresholds() -> None: + policy = AnalysisPolicy( + minimum_history=3, + medium_threshold=1.0, + high_threshold=2.0, + critical_threshold=3.0, + ) + + result = calculate_drift( + "transform", + [1.0, 1.1, 2.0], + policy=policy, + ) + + assert result.severity == "CRITICAL" + + +def test_drift_uses_recent_baseline_window() -> None: + policy = AnalysisPolicy(minimum_history=3, baseline_window=2) + + result = calculate_drift( + "transform", + [100.0, 100.0, 2.0, 2.2, 2.1], + policy=policy, + ) + + assert result.baseline == pytest.approx(2.1) + assert result.severity == "NORMAL" + + +@pytest.mark.parametrize( + ("aggregation", "expected"), + [ + (MappedTaskAggregation.MAX, 5.0), + (MappedTaskAggregation.MEAN, 3.5), + (MappedTaskAggregation.SUM, 7.0), + ], +) +def test_history_uses_mapped_task_aggregation( + aggregation: MappedTaskAggregation, + expected: float, +) -> None: + runs = [ + TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="mapped", + state="success", + duration=duration, + map_index=index, + ) + for index, duration in enumerate([2.0, 5.0]) + ] + + assert build_duration_history(runs, aggregation=aggregation) == { + "mapped": [expected] + } diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 485eefd..8574874 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -23,9 +23,11 @@ def test_top_level_api_declares_supported_exports() -> None: expected_exports = { "AirflowApiError", "AirflowClient", + "AnalysisPolicy", "DAGAnalysis", "DAGDataSource", "Severity", + "MappedTaskAggregation", "TaskRun", "__version__", "analyze_dag",