Metadata-Version: 2.4
Name: airflow-diagnostics-mcp
Version: 0.2.0
Summary: AI-powered diagnostics for Apache Airflow DAGs — root-cause analysis, performance regression detection, and health audits
Author-email: Joshua Cajethan <joshuasamuelmopur@gmail.com>
License: Apache-2.0
License-File: LICENSE
Requires-Python: >=3.12
Requires-Dist: apache-airflow>=2.0.0
Requires-Dist: duckdb>=0.9.0
Requires-Dist: fastmcp>=3.0.0
Requires-Dist: httpx>=0.24.0
Requires-Dist: psycopg2-binary>=2.9.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: python-dotenv>=1.0.0
Requires-Dist: sqlalchemy>=2.0.0
Requires-Dist: structlog>=24.0.0
Provides-Extra: dev
Requires-Dist: black>=23.0.0; extra == 'dev'
Requires-Dist: mypy>=1.0.0; extra == 'dev'
Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Description-Content-Type: text/markdown

# Airflow Diagnostics MCP

**Stop asking Claude to trigger DAGs. Start asking it why they failed.**

An MCP server that brings AI-powered diagnostics to Apache Airflow. When a DAG fails, most people read logs. We analyze them.

## What is this?

Airflow Diagnostics MCP is not another REST API wrapper. Seven servers already do that. Instead, we provide the *reasoning layer* your LLM needs to actually debug DAGs:

- **Root-cause diagnosis** — Not just raw logs, but structured findings: "Your task hit OOM because it processed 50GB instead of 2GB"
- **Performance regression detection** — "This task is 60% slower than last week, here's why"
- **Flaky task surfacing** — "This sensor fails 65% of the time in a specific way"
- **Health audits** — Automated checklist for retry storms, zombie tasks, pool starvation, SLA misses
- **Optimization suggestions** — Rule-based recommendations for parallelism, sensors, and resource allocation

## Try it in 60 seconds

### 1. Start demo Airflow (docker-compose)

```bash
cd demo
docker-compose up -d
# Wait for Airflow webserver to be healthy (~30s)
# http://localhost:8081 (admin / admin)
```

### 2. Install the MCP server

```bash
# Zero-setup (recommended) — uvx runs the published package without installing globally
uvx airflow-diagnostics-mcp

# ...or install locally for development
uv pip install -e .
```

### 3. Configure Claude Desktop

Add to `~/.config/Claude/claude_desktop_config.json` (or the Cursor / Claude Code equivalent):

```json
{
  "mcpServers": {
    "airflow-diagnostics": {
      "command": "uvx",
      "args": ["airflow-diagnostics-mcp"],
      "env": {
        "AIRFLOW_REST_API_URL": "http://localhost:8081/api/v1",
        "AIRFLOW_USERNAME": "admin",
        "AIRFLOW_PASSWORD": "admin",
        "AIRFLOW_METADATA_DB_URL": "postgresql://airflow:airflow@localhost:5432/airflow"
      }
    }
  }
}
```

Restart Claude Desktop. You should now see the Airflow Diagnostics tools available.

### 4. Trigger a failure and diagnose it

In the Airflow UI (http://localhost:8081):

1. Enable the `retry_storm` DAG
2. Trigger a run
3. Wait for a task to fail

Then in Claude:

```
Diagnose the failed run of retry_storm.
```

You'll get back a structured analysis with root cause, failure chain, and recovery steps.

**Demo DAGs included** (enable any of them to exercise the matching tool):

| DAG | Exercises | Ask Claude |
|---|---|---|
| `retry_storm` | `diagnose_failed_run`, `audit_dag_health` | *"Why is retry_storm failing?"* |
| `slow_regression` | `analyze_dag_performance` | *"Is slow_regression getting slower?"* |
| `flaky_api_dag` | `find_flaky_tasks` | *"Which tasks are flaky?"* |
| `pool_starvation` | `audit_dag_health` (pool check) | *"Audit pool_starvation for health issues"* |
| `heavy_xcom` | `find_metadata_db_hotspots` | *"Find metadata DB hotspots"* |
| `healthy_etl` | baseline comparison | *"Compare these two runs"* |

## Tools

### `diagnose_failed_run(dag_id, run_id?)`

Identifies the root cause of a failed DAG run by analyzing logs, checking for known patterns, and examining the failure chain.

**Example output:**

```json
{
  "root_cause": "Task failed due to excessive retries (5 attempts)",
  "severity": "high",
  "failure_chain": ["flaky_task", "unstable_task"],
  "suggested_actions": [
    "Increase timeout or optimize task logic",
    "Check for transient dependency failures",
    "Review logs for specific error patterns"
  ],
  "evidence": {
    "first_failed_task": "flaky_task",
    "try_number": 5,
    "matched_rules": [...]
  }
}
```

### `analyze_dag_performance(dag_id, lookback_days?)`

Analyzes performance trends, detects bottlenecks, and flags regressions.

### `find_flaky_tasks(lookback_days?, min_runs?)`

Surfaces tasks with intermittent failures across all DAGs, clusters failure patterns.

### `audit_dag_health(dag_id, lookback_days?)`

Runs a 6-point health checklist: retry storms, sensor failures, zombie tasks (via `job.latest_heartbeat`), pool starvation (queued latency by pool), SLA misses, and long-tail duration variance. Produces a prioritised `HealthReport` with overall status (`healthy` / `degraded` / `unhealthy`).

### `suggest_optimizations(dag_id, lookback_days?)`

Composes the health audit and performance analysis, then feeds signals into heuristic rules to produce ranked recommendations: enable retries, parallelise short tasks, migrate long sensors to deferrable operators, move high-concurrency tasks off `default_pool`, and act on detected health issues. Suggestions are sorted by estimated impact.

### `compare_runs(dag_id, run_id_a, run_id_b)`

Task-level diff between two runs: duration deltas, regressions, state differences.

### `find_metadata_db_hotspots(top_n?)`

Postgres-only. Uses `pg_total_relation_size`, `pg_stat_user_tables`, and `OCTET_LENGTH(xcom.value)` to identify the biggest tables, the DAGs generating the most task instances, and the (DAG, task) pairs pushing large XCom payloads through the DB. Returns concrete cleanup recommendations (`airflow db clean` targets, remote XCom backend hints). Returns an empty report with a note on non-Postgres backends.

## How it's different

| Feature | astro-airflow-mcp | airflow-mcp-server | airflow-diagnostics-mcp |
|---|---|---|---|
| **List / trigger DAGs** | ✅ | ✅ | — |
| **Root-cause diagnosis** | ❌ | ❌ | ✅ |
| **Performance regression detection** | ❌ | ❌ | ✅ |
| **Flaky task surfacing** | ❌ | ❌ | ✅ |
| **Health audits** | ❌ | ❌ | ✅ |
| **Metadata DB analytics** | ❌ | ❌ | ✅ |
| **Reasoning-focused output** | ❌ | ❌ | ✅ |

We recommend installing this alongside an existing MCP server. Use `astro-airflow-mcp` or `airflow-mcp-server` for operations (trigger, pause, clear), and use this server for diagnosis and analysis.

## Architecture

```
┌─────────────────────────────────────────────────────────────┐
│                    MCP Client (Claude)                      │
└───────────────────────────┬─────────────────────────────────┘
                            │ MCP protocol (stdio/HTTP)
┌───────────────────────────▼─────────────────────────────────┐
│              airflow-diagnostics-mcp (FastMCP)              │
│                                                             │
│  ┌─────────────────────┐    ┌───────────────────────────┐   │
│  │  Diagnosis Layer    │    │   Tool Surface (MCP)      │   │
│  │  - rule engines     │◄───┤   diagnose_failed_run     │   │
│  │  - regression calc  │    │   analyze_dag_performance │   │
│  │  - flake detection  │    │   find_flaky_tasks        │   │
│  │  - cost heuristics  │    │   audit_dag_health        │   │
│  └──────────┬──────────┘    │   suggest_optimizations   │   │
│             │                │   compare_runs            │   │
│  ┌──────────▼──────────┐    │   find_metadata_db_...    │   │
│  │   Data Access       │    └───────────────────────────┘   │
│  │   - REST adapter    │                                    │
│  │   - Metadata DB     │                                    │
│  │   - DuckDB cache    │                                    │
│  └──────────┬──────────┘                                    │
└─────────────┼───────────────────────────────────────────────┘
              │
     ┌────────┴────────┐
     ▼                 ▼
┌──────────┐    ┌─────────────┐
│ Airflow  │    │ Airflow     │
│ REST API │    │ Metadata DB │
│          │    │ (read-only) │
└──────────┘    └─────────────┘
```

**Why two planes:**

- REST API alone → paginates, rate-limits, no fast cross-run aggregation
- Metadata DB alone → misses real-time scheduler state
- DuckDB cache → pull relevant slice, run analytics locally, cache with TTL

## Configuration

Copy `.env.example` to `.env` and update:

```bash
# Airflow REST API
AIRFLOW_REST_API_URL=http://localhost:8081/api/v1
AIRFLOW_USERNAME=admin
AIRFLOW_PASSWORD=admin

# Metadata DB
AIRFLOW_METADATA_DB_URL=postgresql://airflow:airflow@localhost:5432/airflow

# Server
LOG_LEVEL=INFO
```

## Supported Airflow versions

- **2.x** — Fully supported (tested with 2.8+)
- **3.x** — Coming in v1.1

## Contributing

Contributions welcome! Areas for expansion:

- v2 features: LLM-powered log analysis for unknown patterns, ML-based suggestions
- More diagnosis rules: cloud-specific errors, data warehouse errors
- Integration tests against various Airflow setups
- UI dashboard (separate project)

## License

Apache 2.0 — same as Airflow.

## Credits & thanks

This project stands on the shoulders of:

- [Astronomer](https://github.com/astronomer/astronomer-cosmos) for `astro-airflow-mcp` and showing what first-class MCP support looks like
- [Apache Airflow](https://airflow.apache.org/) community for maintaining a genuinely extensible platform
- FastMCP for the fantastic decorator-based MCP framework

## Links

- **GitHub:** [github.com/mrjoshuasamuel/airflow-diagnostics-mcp](https://github.com/mrjoshuasamuel/airflow-diagnostics-mcp)
- **PyPI:** [pypi.org/project/airflow-diagnostics-mcp](https://pypi.org/project/airflow-diagnostics-mcp/)
- **Issues:** [GitHub Issues](https://github.com/mrjoshuasamuel/airflow-diagnostics-mcp/issues)
