← Plugin catalog
Data & Analytics

astronomer-data

Astronomer v0.1.0

Astronomer's toolkit for Apache Airflow and the data it moves. Ask for a pipeline, a fix or an answer from your warehouse, and the plugin follows the steps an Airflow engineer would. Build and test DAGs. Write new DAGs, extend existing ones, and test them locally before you deploy. Covers dbt with Cosmos, DAG Factory, Blueprint, human-in-the-loop tasks, Airflow plugins, and tasks written in Go or Java. Debug failures. Find out why a DAG fails to parse or a task keeps failing. The plugin reads the logs, finds the root cause, and fixes the code. Run Astro. Set up an Astro project, run Airflow locally, deploy, and troubleshoot Astro deployments. Migrate. Move DAGs from Airflow 2 to Airflow 3, or move a Dagster project to Airflow. Query your warehouse. Answer business questions with SQL against Snowflake, BigQuery, Postgres and other SQLAlchemy databases. Profile tables, check how fresh the data is, and generate a schema reference so later questions go faster. Trace lineage. See what feeds a table and what breaks downstream if it changes. Most skills run local tools such as the Astro CLI and uv, so they work best in Codex. Warehouse skills read connection details from ~/.astro/agents/warehouse.yml.

Language: English · Automatically detected from descriptions.

Package details

Publisher declarations from the archived package. These are separate from our research and the live service's terms.

Package author
Astronomer

Package observed Sep 30, 2026.

Files & skills

File archives

Plugin package104 files · 387 KBBrowse files →
Skill instructions
airflow14.9 KB

View saved version →

---
name: airflow
description: Queries, manages, and troubleshoots Apache Airflow using the `af` CLI. Use when working with anything related to Airflow - a DAG, a DAG run, a task log, an import or parse error, a broken DAG, or any Airflow operation. Covers listing and triggering DAGs, retrying runs, reading task logs, diagnosing failures, debugging import and parse errors, checking connections, variables and pools, exploring the REST API, and monitoring health (for example "trigger a pipeline", "retry a run", "list connections", "check Airflow health", "why did my DAG fail"). This is the entrypoint that routes to sibling skills for authoring, testing, deploying, and migrating Airflow 2 to 3. Not for warehouse/SQL analytics on Airflow metadata tables (use analyzing-data); for deep root-cause reports use debugging-dags or airflow-investigation.
---

# Airflow Operations

Use `af` commands to query, manage, and troubleshoot Airflow workflows.

## Astro CLI

The [Astro CLI](https://www.astronomer.io/docs/astro/cli/overview) is the recommended way to run Airflow locally and deploy to production. It provides a containerized Airflow environment that works out of the box:

```bash
# Initialize a new project
astro dev init

# Start local Airflow (webserver at http://localhost:8080)
astro dev start

# Parse DAGs to catch errors quickly (no need to start Airflow)
astro dev parse

# Run pytest against your DAGs
astro dev pytest

# Deploy to production
astro deploy            # Full deploy (image + DAGs)
astro deploy --dags     # DAG-only deploy (fast, no image build)
```

For more details:
- **New project?** See the **setting-up-astro-project** skill
- **Local environment?** See the **managing-astro-local-env** skill
- **Deploying?** See the **deploying-airflow** skill

---

## Running the CLI

These commands assume `af` is on PATH. Run via `astro otto` to get it automatically, or install standalone with `uv tool install astro-airflow-mcp`.

## Instance Configuration

Manage multiple Airflow instances with persistent configuration:

```bash
# Add a new instance
af instance add prod --url https://airflow.example.com --token "$API_TOKEN"
af instance add staging --url https://staging.example.com --username admin --password admin

# List and switch instances
af instance list      # Shows all instances in a table
af instance use prod  # Switch to prod instance
af instance current   # Show current instance
af instance delete old-instance

# Auto-discover instances (use --dry-run to preview first)
af instance discover --dry-run        # Preview all discoverable instances
af instance discover                  # Discover from all backends (astro, local)
af instance discover astro            # Discover Astro deployments only
af instance discover astro --all-workspaces  # Include all accessible workspaces
af instance discover local            # Scan common local Airflow ports
af instance discover local --scan     # Deep scan all ports 1024-65535

# IMPORTANT: Always run with --dry-run first and ask for user consent before
# running discover without it. The non-dry-run mode creates API tokens in
# Astro Cloud, which is a sensitive action that requires explicit approval.

# Show where an instance came from (file path + scope)
af instance show prod

# Override instance for a single command via env vars
AIRFLOW_API_URL=https://staging.example.com AIRFLOW_AUTH_TOKEN=$STG af dags list

# Or switch persistently
af instance use staging
```

Config layout (mirrors `git config` system/global/local):

| Scope | File | Committed? |
|---|---|---|
| Global | `~/.astro/config.yaml` | n/a (per-user) |
| Project shared | `<root>/.astro/config.yaml` | yes |
| Project local | `<root>/.astro/config.local.yaml` | no (gitignored) |

`<root>` is found by walking up from cwd looking for `.astro/`. Default write routing inside a project: `add`/`discover` → project-shared, `use` → project-local. Override with `--global` / `--project` / `--local`. Set `AF_CONFIG=<path>` to bypass layering and use a single file.

Migrate from the legacy `~/.af/config.yaml` with `af migrate` (idempotent; renames the old file to `.bak`).

Tokens in config can reference environment variables using `${VAR}` syntax:
```yaml
instances:
- name: prod
  url: https://airflow.example.com
  auth:
    token: ${AIRFLOW_API_TOKEN}
```

Or use environment variables directly (no config file needed):

```bash
export AIRFLOW_API_URL=http://localhost:8080
export AIRFLOW_AUTH_TOKEN=your-token-here
# Or username/password:
export AIRFLOW_USERNAME=admin
export AIRFLOW_PASSWORD=admin
```

Or CLI flags: `af --airflow-url http://localhost:8080 --token "$TOKEN" <command>`

## Quick Reference

| Command | Description |
|---------|-------------|
| `af health` | System health check |
| `af dags list` | List all DAGs |
| `af dags get <dag_id>` | Get DAG details |
| `af dags explore <dag_id>` | Full DAG investigation |
| `af dags source <dag_id>` | Get DAG source code |
| `af dags pause <dag_id>` | Pause DAG scheduling |
| `af dags unpause <dag_id>` | Resume DAG scheduling |
| `af dags errors` | List import errors |
| `af dags warnings` | List DAG warnings |
| `af dags stats` | DAG run statistics |
| `af runs list` | List DAG runs |
| `af runs get <dag_id> <run_id>` | Get run details |
| `af runs trigger <dag_id>` | Trigger a DAG run |
| `af runs trigger-wait <dag_id>` | Trigger and wait for completion |
| `af runs delete <dag_id> <run_id>` | Permanently delete a DAG run |
| `af runs clear <dag_id> <run_id>` | Clear a run for re-execution |
| `af runs diagnose <dag_id> <run_id>` | Diagnose failed run |
| `af tasks list <dag_id>` | List tasks in DAG |
| `af tasks get <dag_id> <task_id>` | Get task definition |
| `af tasks instance <dag_id> <run_id> <task_id>` | Get task instance |
| `af tasks logs <dag_id> <run_id> <task_id>` | Get task logs |
| `af config version` | Airflow version |
| `af config show` | Full configuration |
| `af config connections` | List connections |
| `af config variables` | List variables |
| `af config variable <key>` | Get specific variable |
| `af config pools` | List pools |
| `af config pool <name>` | Get pool details |
| `af config plugins` | List plugins |
| `af config providers` | List providers |
| `af config assets` | List assets/datasets |
| `af api <endpoint>` | Direct REST API access |
| `af api ls` | List available API endpoints |
| `af api ls --filter X` | List endpoints matching pattern |
| `af registry providers` | List providers in the Airflow Registry |
| `af registry modules <provider>` | List operators/hooks/sensors/transfers in a provider |
| `af registry parameters <provider>` | Constructor signatures (name, type, default, required) for a provider's classes |
| `af registry connections <provider>` | Connection types a provider exposes |

## User Intent Patterns

### Getting Started
- "How do I run Airflow locally?" / "Set up Airflow" -> use the **managing-astro-local-env** skill (uses Astro CLI)
- "Create a new Airflow project" / "Initialize project" -> use the **setting-up-astro-project** skill (uses Astro CLI)
- "How do I install Airflow?" / "Get started with Airflow" -> use the **setting-up-astro-project** skill

### DAG Operations
- "What DAGs exist?" / "List all DAGs" -> `af dags list`
- "Tell me about DAG X" / "What is DAG Y?" -> `af dags explore <dag_id>`
- "What's the schedule for DAG X?" -> `af dags get <dag_id>`
- "Show me the code for DAG X" -> `af dags source <dag_id>`
- "Stop DAG X" / "Pause this workflow" -> `af dags pause <dag_id>`
- "Resume DAG X" -> `af dags unpause <dag_id>`
- "Are there any DAG errors?" -> `af dags errors`
- "Create a new DAG" / "Write a pipeline" -> use the **authoring-dags** skill

### Run Operations
- "What runs have executed?" -> `af runs list`
- "Run DAG X" / "Trigger the pipeline" -> `af runs trigger <dag_id>`
- "Run DAG X and wait" -> `af runs trigger-wait <dag_id>`
- "Why did this run fail?" -> `af runs diagnose <dag_id> <run_id>`
- "Delete this run" / "Remove stuck run" -> `af runs delete <dag_id> <run_id>`
- "Clear this run" / "Retry this run" / "Re-run this" -> `af runs clear <dag_id> <run_id>`
- "Test this DAG and fix if it fails" -> use the **testing-dags** skill

### Task Operations
- "What tasks are in DAG X?" -> `af tasks list <dag_id>`
- "Get task logs" / "Why did task fail?" -> `af tasks logs <dag_id> <run_id> <task_id>`
- "Full root cause analysis" / "Diagnose and fix" -> use the **debugging-dags** skill

### Data Operations
- "Is the data fresh?" / "When was this table last updated?" -> use the **checking-freshness** skill
- "Where does this data come from?" -> use the **tracing-upstream-lineage** skill
- "What depends on this table?" / "What breaks if I change this?" -> use the **tracing-downstream-lineage** skill

### Deployment Operations
- "Deploy my DAGs" / "Push to production" -> use the **deploying-airflow** skill
- "Set up CI/CD" / "Automate deploys" -> use the **deploying-airflow** skill
- "Deploy to Kubernetes" / "Set up Helm" -> use the **deploying-airflow** skill
- "astro deploy" / "DAG-only deploy" -> use the **deploying-airflow** skill

### System Operations
- "What version of Airflow?" -> `af config version`
- "What connections exist?" -> `af config connections`
- "Are pools full?" -> `af config pools`
- "Is Airflow healthy?" -> `af health`

### API Exploration
- "What API endpoints are available?" -> `af api ls`
- "Find variable endpoints" -> `af api ls --filter variable`
- "Access XCom values" / "Get XCom" -> `af api xcom-entries -F dag_id=X -F task_id=Y`
- "Get event logs" / "Audit trail" -> `af api event-logs -F dag_id=X`
- "Create connection via API" -> `af api connections -X POST --body '{...}'`
- "Create variable via API" -> `af api variables -X POST -F key=name -f value=val`

### Registry Discovery
- "What operators does provider X have?" -> `af registry modules <provider>`
- "What are the constructor params for operator Y?" -> `af registry parameters <provider>`
- "What providers exist?" / "Is there a provider for Z?" -> `af registry providers`
- "What connection types does provider X expose?" -> `af registry connections <provider>`
- "Writing a DAG with a specific operator" -> use registry to verify current signature before copying examples

## Common Workflows

### Validate DAGs Before Deploying

If you're using the Astro CLI, you can validate DAGs without a running Airflow instance:

```bash
# Parse DAGs to catch import errors and syntax issues
astro dev parse

# Run unit tests
astro dev pytest
```

Otherwise, validate against a running instance:

```bash
af dags errors     # Check for parse/import errors
af dags warnings   # Check for deprecation warnings
```

### Discover Operator Signatures Before Writing Code

The Airflow Registry at `airflow.apache.org/registry` is the authoritative source for provider classes and their current constructor signatures. Prefer it over memory or stale documentation when authoring DAGs — the registry reflects the live provider release.

```bash
# List all providers and pick the one you need
af registry providers | jq '.providers[] | {id, name, version}'

# List every operator / hook / sensor in a provider (e.g. standard, amazon, google)
af registry modules standard \
  | jq '.modules[] | {name, type, import_path, docs_url}'

# Get the current constructor signature for a specific class
af registry parameters standard \
  | jq '.classes["airflow.providers.standard.operators.hitl.ApprovalOperator"].parameters'

# Filter modules by substring (useful when you know the concept but not the class)
af registry modules standard \
  | jq '.modules[] | select(.import_path | test("hitl"))'
```

Results are cached locally: 1 hour for the latest version, 30 days for pinned versions (which are immutable). Add `--version X.Y.Z` to any `modules` / `parameters` / `connections` call to target a specific release.

### Investigate a Failed Run

```bash
# 1. List recent runs to find failure
af runs list --dag-id my_dag

# 2. Diagnose the specific run
af runs diagnose my_dag manual__2024-01-15T10:00:00+00:00

# 3. Get logs for failed task (from diagnose output)
af tasks logs my_dag manual__2024-01-15T10:00:00+00:00 extract_data

# 4. After fixing, clear the run to retry all tasks
af runs clear my_dag manual__2024-01-15T10:00:00+00:00
```

### Morning Health Check

```bash
# 1. Overall system health
af health

# 2. Check for broken DAGs
af dags errors

# 3. Check pool utilization
af config pools
```

### Understand a DAG

```bash
# Get comprehensive overview (metadata + tasks + source)
af dags explore my_dag
```

### Check Why DAG Isn't Running

```bash
# Check if paused
af dags get my_dag

# Check for import errors
af dags errors

# Check recent runs
af runs list --dag-id my_dag
```

### Trigger and Monitor

```bash
# Option 1: Trigger and wait (blocking)
af runs trigger-wait my_dag --timeout 1800

# Option 2: Trigger and check later
af runs trigger my_dag
af runs get my_dag <run_id>
```

## Output Format

All commands output JSON (except `instance` commands which use human-readable tables):

```bash
af dags list
# {
#   "total_dags": 5,
#   "returned_count": 5,
#   "dags": [...]
# }
```

Use `jq` for filtering:

```bash
# Find failed runs
af runs list | jq '.dag_runs[] | select(.state == "failed")'

# Get DAG IDs only
af dags list | jq '.dags[].dag_id'

# Find paused DAGs
af dags list | jq '[.dags[] | select(.is_paused == true)]'
```

## Task Logs Options

```bash
# Get logs for specific retry attempt
af tasks logs my_dag run_id task_id --try 2

# Get logs for mapped task index
af tasks logs my_dag run_id task_id --map-index 5
```

## Direct API Access with `af api`

Use `af api` for endpoints not covered by high-level commands (XCom, event-logs, backfills, etc).

```bash
# Discover available endpoints
af api ls
af api ls --filter variable

# Basic usage
af api dags
af api dags -F limit=10 -F only_active=true
af api variables -X POST -F key=my_var -f value="my value"
af api variables/old_var -X DELETE
```

**Field syntax**: `-F key=value` auto-converts types, `-f key=value` keeps as string.

**Full reference**: See [api-reference.md](api-reference.md) for all options, common endpoints (XCom, event-logs, backfills), and examples.

## Related Skills

| Skill | Use when... |
|-------|-------------|
| **authoring-dags** | Creating or editing DAG files with best practices |
| **testing-dags** | Iterative test -> debug -> fix -> retest cycles |
| **debugging-dags** | Deep root cause analysis and failure diagnosis |
| **checking-freshness** | Checking if data is up to date or stale |
| **tracing-upstream-lineage** | Finding where data comes from |
| **tracing-downstream-lineage** | Impact analysis -- what breaks if something changes |
| **deploying-airflow** | Deploying DAGs to production (Astro, Docker Compose, Kubernetes) |
| **migrating-airflow-2-to-3** | Upgrading DAGs from Airflow 2.x to 3.x |
| **managing-astro-local-env** | Starting, stopping, or troubleshooting local Airflow |
| **setting-up-astro-project** | Initializing a new Astro/Airflow project |
| **airflow-state-store** | Per-task checkpointing, watermarks, crash-safe operators (Airflow 3.3+) |
| **airflow-hitl** | Pausing a DAG for human approval or input (Airflow 3.1+) |

Referenced files: 2

airflow-hitl7.26 KB

View saved version →

---
name: airflow-hitl
description: Builds human-in-the-loop (HITL) Airflow workflows - approval gates, form input, and human-driven branching. Use when a DAG needs a human in the loop - an approval or reject step, sign-off before a task runs, a decision or approval UI, branching on a human choice, or collecting form input mid-run; also on mentions of ApprovalOperator, HITLOperator, HITLBranchOperator, HITLEntryOperator, or HITLTrigger. Requires Airflow 3.1+. Not for AI/LLM task calls (see migrating-ai-sdk-to-common-ai).
---

# Airflow Human-in-the-Loop Operators

Pause a DAG until a human responds via the Airflow UI or REST API. HITL operators are deferrable — they release their worker slot while waiting.

> **Requires Airflow 3.1+** (`af config version`).
>
> **UI location**: Browse → Required Actions. Respond from the task instance page's Required Actions tab.
>
> **Cross-references**: `migrating-ai-sdk-to-common-ai` for AI/LLM task decorators; `airflow` for registry and API discovery commands used below.

---

## Step 1 — Pick the capability you need

| Capability | Class (verify in Step 2) |
|---|---|
| Approve or reject; downstream skips on reject | `ApprovalOperator` |
| Present N options and return which were chosen | `HITLOperator` |
| Branch to one or more downstream tasks based on a choice | `HITLBranchOperator` |
| Collect a form (no approve/select step) | `HITLEntryOperator` |
| Use the HITL trigger directly (advanced / custom operators) | `HITLTrigger` |

This is the only place class names are hardcoded. The provider adds, renames, and removes params across releases — do not copy parameter lists from memory. Fetch the current signature before writing code.

---

## Step 2 — Discover the current signatures from the Airflow Registry

Before writing HITL code, run these to see the live roster and constructor params (see the `airflow` skill for the full `af registry` reference):

```bash
# Every HITL-related module in the standard provider
af registry modules standard \
  | jq '.modules[] | select(.import_path | test("\\.hitl\\.")) | {name, type, import_path, short_description, docs_url}'

# Constructor signatures: name, type, default, required, description
af registry parameters standard \
  | jq '.classes | to_entries[] | select(.key | test("\\.hitl\\.")) | {fqn: .key, parameters: .value.parameters}'

# Pin to the exact installed provider version
af config providers \
  | jq '.providers[] | select(.package_name == "apache-airflow-providers-standard") | .version'
# then: af registry parameters standard --version <VERSION>
```

If the registry shows a param that this skill does not mention, prefer the registry. If the registry shows a class that is not in Step 1, treat it as additive — the decision table above may be stale.

---

## Step 3 — Canonical example (approval gate)

Starting point for any HITL task. Adapt by swapping the class name and params per Step 2.

```python
from airflow.providers.standard.operators.hitl import ApprovalOperator
from airflow.sdk import dag, task, chain, Param
from pendulum import datetime

@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
def approval_example():
    @task
    def prepare():
        return "Review quarterly report"

    approval = ApprovalOperator(
        task_id="approve_report",
        subject="Report Approval",
        body="{{ ti.xcom_pull(task_ids='prepare') }}",
        defaults="Approve",              # Auto-selected on timeout
        params={"comments": Param("", type="string")},
    )

    @task
    def after_approval(result):
        print(f"Decision: {result['chosen_options']}")

    chain(prepare(), approval)
    after_approval(approval.output)

approval_example()
```

For the other classes in Step 1, the shape is the same (`task_id`, `subject`, plus class-specific params). Verify each constructor through Step 2 — for example, `HITLBranchOperator` requires every option either to match a downstream task id directly or to be resolved via a mapping param surfaced in the registry.

---

## Step 4 — Behavior contracts (stable across versions)

### Timeout
- With `defaults` set: task succeeds on timeout, default option(s) selected.
- Without `defaults`: task fails on timeout.

### Markdown + Jinja in `body`
`body` supports Markdown and is Jinja-templatable. Render XCom context directly:

```python
body = """**Total Budget:** {{ ti.xcom_pull(task_ids='get_budget') }}

| Category | Amount |
|----------|--------|
| Marketing | $1M |
"""
```

### Callbacks
All HITL operators accept the standard Airflow callback kwargs (`on_success_callback`, `on_failure_callback`, etc.).

### Notifiers
HITL operators accept a `notifiers` list. Inside a notifier's `notify(context)` method, build a link to the pending task with `HITLOperator.generate_link_to_ui_from_context(context, base_url=...)`.

### Restricting who can respond
The parameter name and accepted identifier format depend on the active auth manager. Do **not** hardcode — check which one is active and which kwarg the current provider exposes:

```bash
af config show | jq '.auth_manager // .core.auth_manager'
```

Then look up the current kwarg in Step 2 (at the time of writing it is `assigned_users`, accepting identifiers in whatever format the active auth manager uses — Astro uses the Astro user ID, FabAuthManager uses email, SimpleAuthManager uses username).

---

## Step 5 — Responding from external integrations

For Slack bots, custom apps, or scripts. Discover the live endpoint rather than hardcoding a path:

```bash
af api ls --filter hitl           # live endpoint list
af api spec \
  | jq '.paths | to_entries[] | select(.key | test("hitl"))'   # request/response schemas
```

The PATCH-to-respond pattern is stable; the exact path is discovered. Typical shape:

```python
import os, requests

HOST = os.environ["AIRFLOW_HOST"]
TOKEN = os.environ["AIRFLOW_API_TOKEN"]
HEADERS = {"Authorization": f"Bearer {TOKEN}"}

# List pending — use the path from `af api ls --filter hitl`
requests.get(f"{HOST}/<path>", headers=HEADERS, params={"state": "pending"})

# Respond — same discovered path family, PATCH
requests.patch(
    f"{HOST}/<path>/{dag_id}/{run_id}/{task_id}",
    headers=HEADERS,
    json={"chosen_options": ["Approve"], "params_input": {"comments": "ok"}},
)
```

---

## Step 6 — Safety checks

- [ ] Airflow version ≥ 3.1 (`af config version`).
- [ ] Constructor kwargs match the current registry output from Step 2 — no `respondents`-vs-`assigned_users` style drift.
- [ ] For branching: every option resolves to a downstream task id (directly or via the mapping kwarg from Step 2).
- [ ] Every value in `defaults` is also in `options`.
- [ ] `execution_timeout` set; `defaults` configured if timeout should succeed rather than fail.
- [ ] API token configured if external responders are part of the flow.

---

## References

The upstream docs URL is surfaced per-module by the registry — do not hardcode:

```bash
af registry modules standard \
  | jq '.modules[] | select(.import_path | test("\\.hitl\\.")) | {name, docs_url}'
```

## Related skills

- **airflow** — `af registry`, `af api`, `af config` command reference.
- **migrating-ai-sdk-to-common-ai** — AI/LLM task decorators and GenAI patterns (common-ai provider).
- **authoring-dags** — general DAG writing best practices.
- **testing-dags** — iterative test → debug → fix cycles.
airflow-plugins23.6 KB

View saved version →

---
name: airflow-plugins
description: Builds Airflow 3.1+ plugins that embed FastAPI apps, custom UI pages, React components, middleware, macros, and operator links directly into the Airflow UI. Use when building anything custom inside Airflow 3.1+ that involves Python and a browser-facing interface - creating an Airflow plugin, adding a custom UI page or nav entry, building FastAPI-backed endpoints inside Airflow, serving static assets from a plugin, embedding a React app, adding middleware to the API server, creating custom operator extra links, or calling the Airflow REST API from inside a plugin; also when AirflowPlugin, fastapi_apps, external_views, react_apps, or plugin registration come up.
---

# Airflow 3 Plugins

Airflow 3 plugins let you embed FastAPI apps, React UIs, middleware, macros, operator buttons, and custom timetables directly into the Airflow process. No sidecar, no extra server.

> **CRITICAL**: Plugin components (fastapi_apps, react_apps, external_views) require **Airflow 3.1+**. **NEVER import `flask`, `flask_appbuilder`, or use `appbuilder_views` / `flask_blueprints`** — these are Airflow 2 patterns and will not work in Airflow 3. If existing code uses them, rewrite the entire registration block using FastAPI.
>
> **Security**: FastAPI plugin endpoints are **not automatically protected** by Airflow auth. If your endpoints need to be private, implement authentication explicitly using FastAPI's security utilities.
>
> **Restart required**: Changes to Python plugin files require restarting the API server. Static file changes (HTML, JS, CSS) are picked up immediately. Set `AIRFLOW__CORE__LAZY_LOAD_PLUGINS=False` during development to load plugins at startup rather than lazily.
>
> **Relative paths always**: In `external_views`, `href` must have no leading slash. In HTML and JavaScript, use relative paths for all assets and `fetch()` calls. Absolute paths break behind reverse proxies.

### Before writing any code, verify

1. Am I using `fastapi_apps` / FastAPI — not `appbuilder_views` / Flask?
2. Are all HTML/JS asset paths and `fetch()` calls relative (no leading slash)?
3. Are all synchronous SDK or SQLAlchemy calls wrapped in `asyncio.to_thread()`?
4. Do the `static/` and `assets/` directories exist before the FastAPI app mounts them?
5. If the endpoint must be private, did I add explicit FastAPI authentication?

---

## Step 1: Choose plugin components

A single plugin class can register multiple component types at once.

| Component | What it does | Field |
|-----------|-------------|-------|
| Custom API endpoints | FastAPI app mounted in Airflow process | `fastapi_apps` |
| Nav / page link | Embeds a URL as an iframe or links out | `external_views` |
| React component | Custom React app embedded in Airflow UI | `react_apps` |
| API middleware | Intercepts all Airflow API requests/responses | `fastapi_root_middlewares` |
| Jinja macros | Reusable Python functions in DAG templates | `macros` |
| Task instance button | Extra link button in task Detail view | `operator_extra_links` / `global_operator_extra_links` |
| Custom timetable | Custom scheduling logic | `timetables` |
| Event hooks | Listener callbacks for Airflow events | `listeners` |

---

## Step 2: Plugin registration skeleton

### Project file structure

Give each plugin its own subdirectory under `plugins/` — this keeps the Python file, static assets, and templates together and makes multi-plugin projects manageable:

```
plugins/
  my-plugin/
    plugin.py       # AirflowPlugin subclass — auto-discovered by Airflow
    static/
      index.html
      app.js
    assets/
      icon.svg
```

`BASE_DIR = Path(__file__).parent` in `plugin.py` resolves to `plugins/my-plugin/` — static and asset paths will be correct relative to that. Create the subdirectory and any static/assets folders before starting Airflow, or `StaticFiles` will raise on import.

```python
from pathlib import Path
from airflow.plugins_manager import AirflowPlugin
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

BASE_DIR = Path(__file__).parent

app = FastAPI(title="My Plugin")

# Both directories must exist before Airflow starts or FastAPI raises on import
app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")
app.mount("/assets", StaticFiles(directory=BASE_DIR / "assets"), name="assets")


class MyPlugin(AirflowPlugin):
    name = "my_plugin"

    fastapi_apps = [
        {
            "app": app,
            "url_prefix": "/my-plugin",   # plugin available at {AIRFLOW_HOST}/my-plugin/
            "name": "My Plugin",
        }
    ]

    external_views = [
        {
            "name": "My Plugin",
            "href": "my-plugin/ui",              # NO leading slash — breaks on Astro and reverse proxies
            "destination": "nav",                # see locations table below
            "category": "browse",                # nav bar category (nav destination only)
            "url_route": "my-plugin",            # unique route name (required for React apps)
            "icon": "/my-plugin/static/icon.svg" # DOES use a leading slash — served by FastAPI
        }
    ]
```

### External view locations

| `destination` | Where it appears |
|--------------|-----------------|
| `"nav"` | Left navigation bar (also set `category`) |
| `"dag"` | Extra tab on every Dag page |
| `"dag_run"` | Extra tab on every Dag run page |
| `"task"` | Extra tab on every task page |
| `"task_instance"` | Extra tab on every task instance page |

### Nav bar categories (`destination: "nav"`)

Set `"category"` to place the link under a specific nav group: `"browse"`, `"admin"`, or omit for top-level.

### External URLs and minimal plugins

`href` can be a relative path to an internal endpoint (`"my-plugin/ui"`) or a full external URL. A plugin with only `external_views` and no `fastapi_apps` is valid — no backend needed for a simple link or tab:

```python
from airflow.plugins_manager import AirflowPlugin

class LearnViewPlugin(AirflowPlugin):
    name = "learn_view_plugin"

    external_views = [
        {
            "name": "Learn Airflow 3",
            "href": "https://www.astronomer.io/docs/learn",
            "destination": "dag",   # adds a tab to every Dag page
            "url_route": "learn"
        }
    ]
```

The no-leading-slash rule applies to internal paths only — full `https://` URLs are fine.

---

## Step 3: Serve the UI entry point

```python
@app.get("/ui", response_class=FileResponse)
async def serve_ui():
    return FileResponse(BASE_DIR / "static" / "index.html")
```

In HTML, always use **relative paths**. Absolute paths break when Airflow is mounted at a sub-path:

```html
<!-- correct -->
<link rel="stylesheet" href="static/app.css" />
<script src="static/app.js?v=20240315"></script>

<!-- breaks behind a reverse proxy -->
<script src="/my-plugin/static/app.js"></script>
```

Same rule in JavaScript:

```javascript
fetch('api/dags')           // correct — relative to current page
fetch('/my-plugin/api/dags') // breaks on Astro and sub-path deploys
```

---

## Step 4: Call the Airflow API from your plugin

> **Only needed if your plugin calls the Airflow REST API.** Plugins that only serve static files, register `external_views`, or use direct DB access do not need this step — skip to Step 5 or Step 6.

### Add the dependency

Only if REST API communication is being implemented: add `apache-airflow-client` to the project's dependencies. Check which file exists and act accordingly:

| File found | Action |
|------------|--------|
| `requirements.txt` | Append `apache-airflow-client` |
| `pyproject.toml` (uv / poetry) | `uv add apache-airflow-client` or `poetry add apache-airflow-client` |
| None of the above | Tell the user: "Add `apache-airflow-client` to your dependencies before running the plugin." |

Use `apache-airflow-client` to talk to Airflow's own REST API. The SDK is **synchronous** but FastAPI routes are async — never call blocking SDK methods directly inside `async def` or you will stall the event loop and freeze all concurrent requests.

### JWT token management

Cache one token per process. Refresh 5 minutes before the 1-hour expiry. Use double-checked locking so multiple concurrent requests don't all race to refresh simultaneously:

> Replace `MYPLUGIN_` with a short uppercase prefix derived from the plugin name (e.g. if the plugin is called "Trip Analyzer", use `TRIP_ANALYZER_`). If no plugin name has been given yet, ask the user before writing env var names.

```python
import asyncio
import os
import threading
import time
import airflow_client.client as airflow_sdk
import requests

AIRFLOW_HOST  = os.environ.get("MYPLUGIN_HOST",     "http://localhost:8080")
AIRFLOW_USER  = os.environ.get("MYPLUGIN_USERNAME", "admin")
AIRFLOW_PASS  = os.environ.get("MYPLUGIN_PASSWORD", "admin")
AIRFLOW_TOKEN = os.environ.get("MYPLUGIN_TOKEN")    # Astronomer Astro: Deployment API token

_cached_token: str | None = None
_token_expires_at: float  = 0.0
_token_lock = threading.Lock()


def _fetch_fresh_token() -> str:
    """Exchange username/password for a JWT via Airflow's auth endpoint."""
    response = requests.post(
        f"{AIRFLOW_HOST}/auth/token",
        json={"username": AIRFLOW_USER, "password": AIRFLOW_PASS},
        timeout=10,
    )
    response.raise_for_status()
    return response.json()["access_token"]


def _get_token() -> str:
    # Astronomer Astro production: use static Deployment API token directly
    if AIRFLOW_TOKEN:
        return AIRFLOW_TOKEN
    global _cached_token, _token_expires_at
    now = time.monotonic()
    # Fast path — no lock if still valid
    if _cached_token and now < _token_expires_at:
        return _cached_token
    # Slow path — one thread refreshes, others wait
    with _token_lock:
        if _cached_token and now < _token_expires_at:
            return _cached_token
        _cached_token = _fetch_fresh_token()
        _token_expires_at = now + 55 * 60  # refresh 5 min before 1-hour expiry
    return _cached_token


def _make_config() -> airflow_sdk.Configuration:
    config = airflow_sdk.Configuration(host=AIRFLOW_HOST)
    config.access_token = _get_token()
    return config
```

After implementing auth, tell the user:

- **Local development**: set `MYPLUGIN_USERNAME` and `MYPLUGIN_PASSWORD` in `.env` — JWT exchange happens automatically.
- **Astronomer Astro (production)**: create a Deployment API token and set it as `MYPLUGIN_TOKEN` — the JWT exchange is skipped entirely:
  1. Astro UI → open the Deployment → **Access** → **API Tokens** → **+ Deployment API Token**
  2. Copy the token value (shown only once)
  3. `astro deployment variable create MYPLUGIN_TOKEN=<token>`

  `MYPLUGIN_USERNAME` and `MYPLUGIN_PASSWORD` are not needed on Astro.

### Wrapping SDK calls with asyncio.to_thread

```python
from fastapi import HTTPException
from airflow_client.client.api import DAGApi

@app.get("/api/dags")
async def list_dags():
    try:
        def _fetch():
            with airflow_sdk.ApiClient(_make_config()) as client:
                return DAGApi(client).get_dags(limit=100).dags
        dags = await asyncio.to_thread(_fetch)
        return [{"dag_id": d.dag_id, "is_paused": d.is_paused, "timetable_summary": d.timetable_summary} for d in dags]
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
```

> **API field names**: Never guess response field names — verify against the [REST API reference](https://airflow.apache.org/docs/apache-airflow/stable/stable-rest-api-ref.html). Key `DAGResponse` fields: `dag_id`, `dag_display_name`, `description`, `is_paused`, `timetable_summary`, `timetable_description`, `fileloc`, `owners`, `tags`.

The pattern is always: define a plain inner `def _fetch()` with all SDK logic, then `await asyncio.to_thread(_fetch)`.

### Alternative: Direct database access

> **Warning — use with caution and tell the user.** The Airflow metadb is not a public interface. Direct writes or poorly-formed queries can corrupt scheduler state. Whenever you use this pattern, explicitly tell the user: "This accesses Airflow's internal database directly. The internal models are not part of the public API, can change between Airflow versions, and incorrect queries can cause issues in the metadb. Prefer `apache-airflow-client` unless the operation is not exposed via the REST API."

Since FastAPI plugin endpoints run inside the **API server process** (not in a task worker), they have direct access to Airflow's internal SQLAlchemy models — no HTTP round-trip or JWT needed. Use only for read operations not exposed via the REST API, or when the extra HTTP overhead genuinely matters. Always wrap DB calls in `asyncio.to_thread()` — SQLAlchemy queries are blocking.

```python
from airflow.models import DagBag, DagModel
from airflow.utils.db import provide_session

@app.get("/api/dags/status")
async def dag_status():
    def _fetch():
        @provide_session
        def _query(session=None):
            dagbag = DagBag()
            paused = sum(
                1 for dag_id in dagbag.dags
                if (m := session.query(DagModel).filter(DagModel.dag_id == dag_id).first())
                and m.is_paused
            )
            return {"total": len(dagbag.dags), "paused": paused}
        return _query()
    return await asyncio.to_thread(_fetch)
```

---

## Step 5: Common API endpoint patterns

> **If you need an SDK method or field not shown in the examples below**, verify it before generating code — do not guess. Either run `python3 -c "from airflow_client.client.api import <Class>; print([m for m in dir(<Class>) if not m.startswith('_')])"` in any environment where the SDK is installed, or search the [`apache/airflow-client-python`](https://github.com/apache/airflow-client-python) repo for the class definition.

```python
from airflow_client.client.api import DAGApi, DagRunApi
from airflow_client.client.models import TriggerDAGRunPostBody, DAGPatchBody


@app.post("/api/dags/{dag_id}/trigger")
async def trigger_dag(dag_id: str):
    def _run():
        with airflow_sdk.ApiClient(_make_config()) as client:
            return DagRunApi(client).trigger_dag_run(dag_id, TriggerDAGRunPostBody())
    result = await asyncio.to_thread(_run)
    return {"run_id": result.dag_run_id, "state": normalize_state(result.state)}


@app.patch("/api/dags/{dag_id}/pause")
async def toggle_pause(dag_id: str, is_paused: bool):
    def _run():
        with airflow_sdk.ApiClient(_make_config()) as client:
            DAGApi(client).patch_dag(dag_id, DAGPatchBody(is_paused=is_paused))
    await asyncio.to_thread(_run)
    return {"dag_id": dag_id, "is_paused": is_paused}


@app.delete("/api/dags/{dag_id}")
async def delete_dag(dag_id: str):
    def _run():
        with airflow_sdk.ApiClient(_make_config()) as client:
            DAGApi(client).delete_dag(dag_id)
    await asyncio.to_thread(_run)
    return {"deleted": dag_id}


def normalize_state(raw) -> str:
    """Convert SDK enum objects to plain strings before sending to the frontend."""
    if raw is None:
        return "never_run"
    return str(raw).lower()
```

### DAG runs, task instances, and logs

These are the most common calls beyond basic DAG CRUD. For anything not shown here, consult the [REST API reference](https://airflow.apache.org/docs/apache-airflow/stable/stable-rest-api-ref.html) for available endpoints and the matching Python SDK class/method names.

```python
from airflow_client.client.api import DagRunApi, TaskInstanceApi

# Latest run for a DAG
@app.get("/api/dags/{dag_id}/runs/latest")
async def latest_run(dag_id: str):
    def _fetch():
        with airflow_sdk.ApiClient(_make_config()) as client:
            runs = DagRunApi(client).get_dag_runs(dag_id, limit=1, order_by="-start_date").dag_runs
            return runs[0] if runs else None
    run = await asyncio.to_thread(_fetch)
    if not run:
        return {"state": "never_run"}
    return {"run_id": run.dag_run_id, "state": normalize_state(run.state)}


# Task instances for a specific run
@app.get("/api/dags/{dag_id}/runs/{run_id}/tasks")
async def task_instances(dag_id: str, run_id: str):
    def _fetch():
        with airflow_sdk.ApiClient(_make_config()) as client:
            return TaskInstanceApi(client).get_task_instances(dag_id, run_id).task_instances
    tasks = await asyncio.to_thread(_fetch)
    return [{"task_id": t.task_id, "state": normalize_state(t.state)} for t in tasks]


# Task log (try_number starts at 1)
@app.get("/api/dags/{dag_id}/runs/{run_id}/tasks/{task_id}/logs/{try_number}")
async def task_log(dag_id: str, run_id: str, task_id: str, try_number: int):
    def _fetch():
        with airflow_sdk.ApiClient(_make_config()) as client:
            return TaskInstanceApi(client).get_log(
                dag_id, run_id, task_id, try_number, map_index=-1
            )
    result = await asyncio.to_thread(_fetch)
    return {"log": result.content if hasattr(result, "content") else str(result)}
```

### Streaming proxy

Use `StreamingResponse` to proxy binary content from an external URL through the plugin — useful when the browser can't fetch the resource directly (CORS, auth, etc.):

```python
import requests
from starlette.responses import StreamingResponse

@app.get("/api/files/{filename}")
async def proxy_file(filename: str):
    def _stream():
        r = requests.get(f"https://files.example.com/{filename}", stream=True)
        r.raise_for_status()
        return r
    response = await asyncio.to_thread(_stream)
    return StreamingResponse(
        response.iter_content(chunk_size=8192),
        media_type="application/octet-stream",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )
```

Note that `requests.get()` is blocking — fetch in `asyncio.to_thread` so the event loop isn't stalled while waiting for the remote server.

---

## Step 6: Other plugin component types

### Macros

Macros are loaded by the **scheduler** (and DAG processor), not the API server. Restart the scheduler after changes.

```python
from airflow.plugins_manager import AirflowPlugin

def format_confidence(confidence: float) -> str:
    return f"{confidence * 100:.2f}%"

class MyPlugin(AirflowPlugin):
    name = "my_plugin"
    macros = [format_confidence]
```

Use in any templated field — including with XCom:

```
{{ macros.my_plugin.format_confidence(0.95) }}

{{ macros.my_plugin.format_confidence(ti.xcom_pull(task_ids='score_task')['confidence']) }}
```

The naming pattern is always `macros.{plugin_name}.{function_name}`.

### Middleware

Middleware applies to **all** Airflow API requests, including the built-in REST API and any FastAPI plugins. Use sparingly and filter requests explicitly if needed:

```python
from starlette.middleware.base import BaseHTTPMiddleware
from fastapi import Request, Response

class AuditMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        # runs before every request to the Airflow API server
        response = await call_next(request)
        return response

class MyPlugin(AirflowPlugin):
    name = "my_plugin"
    fastapi_root_middlewares = [
        {"middleware": AuditMiddleware, "args": [], "kwargs": {}, "name": "Audit"}
    ]
```

### Operator extra links

```python
from airflow.sdk.bases.operatorlink import BaseOperatorLink

class MyDashboardLink(BaseOperatorLink):
    name = "Open in Dashboard"

    def get_link(self, operator, *, ti_key, **context) -> str:
        return f"https://my-dashboard.example.com/tasks/{ti_key.task_id}"

class MyPlugin(AirflowPlugin):
    name = "my_plugin"
    global_operator_extra_links = [MyDashboardLink()]  # appears on every task
    # operator_extra_links = [MyDashboardLink()]       # attach to specific operator instead
```

### React apps

React apps are embedded as JavaScript bundles served via FastAPI. The bundle must expose itself as a global variable matching the plugin name:

```javascript
// In your bundle (e.g. my-app.js)
globalThis['My Plugin'] = MyComponent;   // matches plugin name
globalThis.AirflowPlugin = MyComponent;  // fallback Airflow looks for
```

```python
class MyPlugin(AirflowPlugin):
    name = "my_plugin"
    fastapi_apps = [{"app": app, "url_prefix": "/my-plugin", "name": "My Plugin"}]
    react_apps = [
        {
            "name": "My Plugin",
            "bundle_url": "/my-plugin/my-app.js",
            "destination": "nav",
            "category": "browse",
            "url_route": "my-plugin",
        }
    ]
```

The same bundle can be registered to multiple destinations by adding multiple entries — each needs a unique `url_route`:

```python
react_apps = [
    {"name": "My Widget", "bundle_url": "/my-plugin/widget.js", "destination": "nav",  "url_route": "my-widget-nav"},
    {"name": "My Widget", "bundle_url": "/my-plugin/widget.js", "destination": "dag",  "url_route": "my-widget-dag"},
]
```

> React app integration is experimental in Airflow 3.1. Interfaces may change in future releases.

---

## Step 7: Environment variables and deployment

Never hardcode credentials:

```python
AIRFLOW_HOST = os.environ.get("MYPLUGIN_HOST",     "http://localhost:8080")
AIRFLOW_USER = os.environ.get("MYPLUGIN_USERNAME", "admin")
AIRFLOW_PASS = os.environ.get("MYPLUGIN_PASSWORD", "admin")
```

**Local Astro CLI:**
```
# .env
MYPLUGIN_HOST=http://localhost:8080
MYPLUGIN_USERNAME=admin
MYPLUGIN_PASSWORD=admin
```

```bash
astro dev restart              # required after any Python plugin change

# Check logs by component (Astro CLI):
astro dev logs --api-server    # FastAPI apps, external_views — plugin import errors show here
astro dev logs --scheduler     # macros, timetables, listeners, operator links
astro dev logs --dag-processor # DAG parsing errors

# Non-Astro:
airflow plugins                # CLI — lists all loaded plugins
```

**Production Astronomer:**
```bash
astro deployment variable create --deployment-id <id> MYPLUGIN_HOST=https://airflow.example.com
```

**Auto-reload during development** (skips lazy loading):
```
AIRFLOW__CORE__LAZY_LOAD_PLUGINS=False
```

**Cache busting for static files** after deploy:
```html
<script src="static/app.js?v=20240315-1"></script>
```

**Verify the plugin loaded**: open **Admin > Plugins** in the Airflow UI.

**OpenAPI docs** are auto-generated for FastAPI plugins:
- Swagger UI: `{AIRFLOW_HOST}/{url_prefix}/docs`
- OpenAPI JSON: `{AIRFLOW_HOST}/{url_prefix}/openapi.json`

---

## Common pitfalls

| Problem | Cause | Fix |
|---------|-------|-----|
| Nav link goes to 404 | Leading `/` in `href` | `"my-plugin/ui"` not `"/my-plugin/ui"` |
| Nav icon not showing | Missing `/` in `icon` | `icon` takes an absolute path: `"/my-plugin/static/icon.svg"` |
| Event loop freezes under load | Sync SDK called directly in `async def` | Wrap with `asyncio.to_thread()` |
| 401 errors after 1 hour | JWT expires with no refresh | Use the 5-minute pre-expiry refresh pattern |
| `StaticFiles` raises on startup | Directory missing | Create `assets/` and `static/` before starting |
| Plugin not showing up | Python file changed without restart | `astro dev restart` |
| Endpoints accessible without login | FastAPI apps are not auto-authenticated | Add FastAPI security (e.g. OAuth2, API key) if endpoints must be private |
| Middleware affecting wrong routes | Middleware applies to all API traffic | Filter by `request.url.path` inside `dispatch()` |
| JS `fetch()` breaks on Astro | Absolute path in `fetch()` | Always use relative paths: `fetch('api/dags')` |

---

## References

- [Airflow plugins documentation](https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/plugins.html)
- [Airflow REST API reference](https://airflow.apache.org/docs/apache-airflow/stable/stable-rest-api-ref.html) — full endpoint list with SDK class/method names
- [Astronomer: Using Airflow plugins](https://www.astronomer.io/docs/learn/using-airflow-plugins)
airflow-state-store24 KB

View saved version →

---
name: airflow-state-store
description: Persists task and asset state across retries and DAG runs using Airflow 3.3's AIP-103 key/value stores (`task_state_store`, `asset_state_store`) and the crash-safe `ResumableJobMixin`. Use when the user asks about task state store, checkpointing in tasks, persisting state across retries, job IDs surviving worker crashes, watermarks, asset metadata, resumable tasks, crash-safe operators, or "what's new in Airflow 3.3". Also use proactively when reading a DAG that uses Variables or XCom for intra-task coordination state (flag the anti-pattern), or when reviewing any DAG that submits a job to an external system (Databricks, Snowflake, BigQuery, Redshift, Spark, dbt Cloud, EMR, AWS Batch, etc.) and waits for it to finish, whether as one submit-and-wait operator or split into submit + sensor/polling tasks; covers `wait_for_termination`, `deferrable`, `durable`, and hand-rolled polling sensors. State-persistence pieces require Airflow 3.3+; submit+poll architecture guidance applies on any version.
---

# Airflow Task State Store (AIP-103)

Airflow 3.3 ships two key/value stores and a crash-safety mixin for operators that submit external jobs.

> **`task_state_store`, `asset_state_store`, and `ResumableJobMixin`'s crash-safety guarantee require Airflow 3.3+.** Check first:
> ```bash
> af config version
> ```
> Below 3.3: `task_state_store`/`asset_state_store` are unavailable, and `durable=True` is a no-op — provider operators ship a pre-3.3 `ResumableJobMixin` shim that always submits fresh (see Section 5). Tell the user those specific features aren't available yet and link the AIP-103 tracking issue. This does **not** gate Section 6's Triggerer-vs-`mode="reschedule"` decision, or the general "green submit ≠ success" anti-pattern — those apply on any Airflow version. On a pre-3.3 DAG, give that guidance in full; only drop the "`durable=True` adds crash-safety" half of it.

---

## Section 1 — Pick the right primitive

| I need to… | Use |
|---|---|
| Persist a cursor, offset, or job ID so a retry can resume instead of restart | `task_state_store` |
| Pass small coordination state within one task across retries (not between tasks) | `task_state_store` |
| Store a watermark or last-processed timestamp per asset, surviving across DAG runs | `asset_state_store` |
| Cache asset-level metadata (manifest hash, row count, schema version) | `asset_state_store` |
| Make an existing non deferrable operator crash-safe when it submits to an external system | `task_state_store` or `ResumableJobMixin` |

**When NOT to use these:**
- Passing data *between* tasks -> use XCom
- Large payloads (model weights, dataframes) -> use XCom with an object storage backend
- Config or secrets shared across DAGs -> use Variables or Connections

---

## Section 2 — Detect anti-patterns in existing DAGs (on demand)

When the user asks to review a DAG or asks "is there a better way", scan for these patterns and flag them:

| Pattern seen in DAG | Problem | Recommend |
|---|---|---|
| `Variable.get(...)` / `Variable.set(...)` inside a `@task` body for per-run state | Variables are global and shared; no scoping to task instance or retry | `task_state_store` |
| `context["ti"].xcom_push(key="job_id", ...)` to survive retries | XCom is scoped to a DAG run, not a retry; a new ti_id is issued per retry | `task_state_store` or `ResumableJobMixin` |
| Manual `if Variable.get("job_id"): reconnect else: submit` retry-resume logic | Reimplements what `ResumableJobMixin` already provides, without the crash-safety guarantee | `ResumableJobMixin` |
| `Variable.set("last_processed_at", ...)` for watermarks | Global; any DAG or task can overwrite it; no scoping to asset | `asset_state_store` |
| Separate `submit` task (`wait_for_termination=False` / fire-and-forget) + a second sensor/polling task waiting on the same external job (Databricks, Snowflake, BigQuery, Redshift, Spark, etc.) | A green submit task only means the job was *accepted*, not that it *succeeded* — only the sensor task's outcome reflects reality. | See **Section 6, "Submit-and-poll DAGs: one task or two?"** — the right call depends on Triggerer availability and job duration, not a single fixed answer. |

Show a before/after snippet when flagging. Use the canonical examples in Steps 3–5 as the "after".

**The submit+sensor split is worth a comment even when the sensor code is bug-free** — reviewing the sensor's code quality (correct `mode=`, correct terminal-state handling, cached hook) is a separate question from whether the two-task split is the right architecture. Follow **Section 6, "Submit-and-poll DAGs: one task or two?"** for that decision; don't re-derive it here.

---

## Section 3 — `task_state_store`: per-task coordination state

`task_state_store` is a key/value store scoped to a single task instance identity (dag_id + run_id + task_id + map_index). It survives retries — a new retry on the same task reads the same store.

```python
from airflow.sdk import dag, task
from pendulum import datetime

@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
def etl_with_checkpoint():

    @task(retries=3)
    def process_records(**context):
        task_state_store = context["task_state_store"]  # injected by Airflow, no setup needed
        cursor = task_state_store.get("last_cursor", default=0)
        records = fetch_records_after(cursor)
        for record in records:
            process(record)
            cursor = record["id"]
            task_state_store.set("last_cursor", cursor)   # checkpoint after each record

    process_records()

etl_with_checkpoint()
```

**API:**
```python
from airflow.sdk import NEVER_EXPIRE

task_state_store.get(key, default=None)                        # returns a JsonValue or default
task_state_store.set(key, value)                               # uses default_retention_days
task_state_store.set(key, value, retention=timedelta(days=7))  # per-key TTL override
task_state_store.set(key, value, retention=NEVER_EXPIRE)       # never expires regardless of config
task_state_store.delete(key)                                   # no-op if key does not exist
task_state_store.clear()                                       # delete all keys for this task instance
```

**Key rules:**
- Values must be JSON-serializable (`str`, `int`, `float`, `bool`, `list`, `dict` — `None` values are rejected).
- Default expiry is controlled by `[state_store] default_retention_days` (0 = never expire).
- Use `NEVER_EXPIRE` for keys that must outlive the default retention window (e.g. a job ID for a multi-day Spark job).
- Max value size defaults to 64 KB; configurable via `[state_store] max_value_storage_bytes` (0 = no limit). For larger payloads, configure a custom `[state_store] backend` or a worker side backend configured via: `[workers] state_store_backend`.

**Mapped tasks — each index has its own namespace:**

When a task is dynamically mapped (`task.expand(...)`), each map index gets an isolated `task_state_store` scoped to its own `map_index`. Indices do not share state.

```python
@task(retries=2)
def process_partition(partition_id, **context):
    task_state_store = context["task_state_store"]
    # Scoped to THIS index only — other indices have their own copy
    cursor = task_state_store.get("cursor", default=0)
    task_state_store.set("cursor", new_cursor)

process_partition.expand(partition_id=[0, 1, 2, 3])
```

`clear()` clears only the current index. To wipe state across all map indices of a task group, use the CLI or core API.

**Before (anti-pattern):**
```python
@task
def process(**context):
    cursor = Variable.get("etl_cursor", default_var=0)
    # ... process ...
    Variable.set("etl_cursor", new_cursor)  # global, any task can overwrite
```

**After:**
```python
@task(retries=3)
def process(**context):
    task_state_store = context["task_state_store"]
    cursor = task_state_store.get("cursor", default=0)
    # ... process ...
    task_state_store.set("cursor", new_cursor)    # scoped to this task instance
```

---

## Section 4 — `asset_state_store`: per-asset metadata across DAG runs

`asset_state_store` is scoped to an asset, not a task instance. It persists across DAG runs — the same key on the same asset is readable and writable by any task that produces or consumes it.

```python
from airflow.sdk import DAG, Asset, task
from datetime import datetime, timezone

ORDERS = Asset(name="orders/daily", uri="s3://warehouse/orders/daily")

with DAG(dag_id="producer", schedule=None, start_date=datetime(2026, 1, 1), catchup=False):

    @task(inlets=[ORDERS], outlets=[ORDERS])
    def load(asset_state_store=None):        # asset_state_store injected by Airflow — declare as a kwarg
        asset_state_store = asset_state_store[ORDERS]

        watermark = asset_state_store.get("watermark", default="2026-01-01T00:00:00+00:00")
        records = fetch_records_since(watermark)

        now = datetime.now(tz=timezone.utc).isoformat()
        asset_state_store.set("watermark", now)
        asset_state_store.set("last_run_summary", {"rows_loaded": len(records), "completed_at": now})

    load()
```

**Reading the store from a consumer DAG:**
```python
with DAG(dag_id="consumer", schedule=[ORDERS], start_date=datetime(2026, 1, 1), catchup=False):

    @task(inlets=[ORDERS])
    def consume(asset_state_store=None):
        asset_state_store = asset_state_store[ORDERS]
        summary = asset_state_store.get("last_run_summary") or {}
        print(f"Processing {summary.get('rows_loaded')} rows up to {asset_state_store.get('watermark')}")

    consume()
```

**Key rules:**
- `asset_state_store` is injected by Airflow as a named kwarg — declare it as `def my_task(asset_state_store=None)`. Do NOT combine with `**context`; Airflow injects it separately.
- Use `datetime.now(tz=timezone.utc).isoformat()` for timestamps — never `datetime.utcnow()` (not timezone-aware).
- Same JSON-serializable value constraint as `task_state_store`.
- No per-key expiry — asset state store entries have no TTL (the asset outlives any single run).
- Readable by any DAG that declares the asset as an inlet or outlet.

**Mapped tasks — last writer wins:**

`asset_state_store` is scoped to the asset, not the map index. If multiple mapped indices write the same key concurrently, the last write wins. Use distinct keys per index or ensure only one index writes to a given key.

```python
@task(outlets=[my_asset])
def load_partition(partition_id, asset_state_store=None):
    asset_state_store = asset_state_store[my_asset]
    # Distinct key per index — no race condition
    asset_state_store.set(f"offset_{partition_id}", new_offset)
```

**Before (anti-pattern):**
```python
Variable.set(f"watermark_{asset_name}", new_offset)   # global, not scoped to asset
```

**After:**
```python
@task(inlets=[my_asset], outlets=[my_asset])
def load(asset_state_store=None):
    asset_state_store = asset_state_store[my_asset]
    asset_state_store.set("watermark", new_offset)
```

---

## Section 5 — `ResumableJobMixin`: crash-safe external job submission

Use whenever a task submits a job to an external system (Spark, Databricks, dbt Cloud, AWS Batch, etc.) and could resubmit a duplicate on retry — whether that same task also polls for completion inside one `execute()` call, or hands the job id to a separate downstream poll/sensor task. Without this mixin (or a provider operator that already builds it in, like `DatabricksSubmitRunOperator`'s `durable=True` default), a worker crash after submission means the next retry of the **submit** task resubmits a duplicate job — that risk exists regardless of whether polling happens in that same task or a separate one.

**Scope check before recommending it:** if submit and poll are already two *separate* tasks (e.g. a `submit` task handing a job id to a downstream sensor/poll task via XCom), the **poll/sensor task** doesn't need this mixin — it never submits anything, so retrying it is already safe. The **submit task** still does; splitting off the poll step doesn't make the submit task's own duplicate-submission risk go away. See the table row below.

**When NOT to use `ResumableJobMixin`:**

| Situation | Use instead | Why |
|---|---|---|
| A Triggerer is deployed and a deferrable operator exists (or can be written) | Deferrable operator | Frees the worker slot during polling; more resource-efficient |
| The task fans out many concurrent I/O operations within a single execution | `async def` task / `BaseAsyncOperator` | Async is for high-throughput I/O, not crash recovery |
| `retries=0` | — | Crash recovery has nothing to reconnect to |
| The external system has no trackable job ID (`submit_job` returns `None`) | Plain operator | The mixin's crash-safety guarantee is silently disabled; adds no value |
| Submit and poll are already split into two separate tasks (submit task → XCom → sensor/poll task) | Nothing for the poll/sensor task — but the submit task itself still wants `durable=True` | The poll/sensor task never submits anything, so retrying it is already idempotent — no checkpoint needed there. The submit task can still resubmit on its own retry unless it's itself crash-safe: e.g. `DatabricksSubmitRunOperator.execute()` routes through `execute_resumable()` (and checkpoints the run id under `durable=True`, the default) even with `wait_for_termination=False` — keep that default on. A hand-rolled submit task (a plain `@task` calling the hook directly) gets none of this for free and should implement `ResumableJobMixin` itself. |

`ResumableJobMixin` holds the worker slot for the full polling duration — the same as a standard synchronous operator. The benefit is crash safety and job continuity, not resource efficiency.

**Opting out of crash recovery:**

The mixin ships with `durable=True` by default. Set `durable=False` to skip all `task_state_store` interaction and run a plain submit/poll/result cycle — useful in test environments or when the external system has its own dedup:

```python
MyBatchOperator(task_id="job", durable=False)

# Or via default_args to disable for all tasks in a DAG:
with DAG("my_dag", default_args={"durable": False}):
    ...
```

### Implementing the mixin

```python
from airflow.sdk import BaseOperator, ResumableJobMixin
from pydantic import JsonValue


class MyBatchOperator(BaseOperator, ResumableJobMixin):

    external_id_key = "batch_job_id"   # key used in task_state_store; set once, never rename

    def execute(self, context):
        return self.execute_resumable(context)  # never call self.execute() — call this

    def submit_job(self, context) -> JsonValue:
        # Submit and return the job identifier. This value is persisted to task_state_store
        # before polling starts. Return None only if the system has no trackable ID
        # (in that case crash-safety is disabled and the job resubmits on every retry).
        return self.hook.submit_batch(...)

    def get_job_status(self, external_id: JsonValue, context) -> str:
        # Query the external system. Return a raw status string.
        return self.hook.get_status(external_id)

    def is_job_active(self, status: str) -> bool:
        # Return True if the job is still running and should be reconnected to.
        return status in ("RUNNING", "PENDING", "QUEUED")

    def is_job_succeeded(self, status: str) -> bool:
        return status == "SUCCEEDED"

    def poll_until_complete(self, external_id: JsonValue, context) -> None:
        # Block until the job reaches a terminal state. Raise on failure.
        self.hook.wait(external_id)

    def get_job_result(self, external_id: JsonValue, context):
        # Return the job result after success. Return None if not applicable.
        return None
```

### What happens on retry

| Job state on retry | Mixin behaviour |
|---|---|
| Still running | Reconnects — calls `poll_until_complete` without resubmitting |
| Already succeeded | Returns `get_job_result` immediately |
| Failed / unknown | Submits a fresh job |

### `external_id_key` warning

> **Never rename `external_id_key` on an operator that is already deployed with in-flight task instances.** The old key is stored in `task_state_store` under the previous name. A rename makes the mixin treat every active retry as a fresh submission, defeating the crash-safety guarantee.

### Before (anti-pattern):
```python
def execute(self, context):
    job_id = Variable.get("spark_job_id", default_var=None)
    if job_id and self._is_running(job_id):
        self._wait(job_id)
    else:
        job_id = self.hook.submit(...)
        Variable.set("spark_job_id", job_id)   # global, race-prone
        self._wait(job_id)
```

**After:**
```python
class MySparkOperator(BaseOperator, ResumableJobMixin):
    external_id_key = "spark_job_id"
    def execute(self, context): return self.execute_resumable(context)
    def submit_job(self, context): return self.hook.submit(...)
    # ... implement the 5 other methods ...
```

For the "submit + separate sensor task" DAG shape specifically — whether to collapse it to one task — see **Section 6, "Submit-and-poll DAGs: one task or two?"**. That's an architecture call driven by Triggerer availability and job duration. It's a different question from whether the submit task itself needs `durable=True` for crash-safety (usually yes, whether or not the split stays) — see the table row above.

---

## Section 6 — Submit-and-poll DAGs: one task or two?

A common DAG shape for Databricks/Snowflake/BigQuery/Redshift/Spark jobs: a `submit` task that fires the job with `wait_for_termination=False` (or equivalent) and returns immediately, followed by a hand-rolled sensor task that polls the same job to completion. The submit task going green the instant the external system *accepts* the job (not when it finishes) is always worth pointing out — a retry, alert, or SLA on the submit task alone tells you nothing about real job success, only the sensor task's outcome does. But whether the split itself should be removed depends on Triggerer availability:

**If a Triggerer is deployed** — collapse to one task with `deferrable=True` **and set `wait_for_termination=True` explicitly**. `deferrable=True` does not by itself imply `wait_for_termination=True` — on `DatabricksSubmitRunOperator`, the deferral helper only defers if `wait_for_termination` is also `True`; carry `wait_for_termination=False` over from an old submit task and it silently skips deferral and polling entirely, reproducing the exact fire-and-forget gap this section exists to remove.

This still frees the worker during polling and the task's own outcome reflects the real job result — but it isn't a strictly-better upgrade with zero tradeoff: the deferred path submits the job directly and never goes through the mixin's checkpointing step, so `durable=True` protects nothing here (unlike the synchronous path below). A worker crash between submission and the trigger taking over still resubmits on retry.

```python
run_job = DatabricksSubmitRunOperator(
    task_id="run_job",
    tasks=[{"task_key": "job", "notebook_task": {"notebook_path": NOTEBOOK_PATH}}],
    deferrable=True,
    wait_for_termination=True,   # required — deferrable=True alone does not imply this
)
```

**If no Triggerer is deployed** — this is a genuine tradeoff, not an automatic call either way. It splits into two cases depending on whether holding a worker for the job's duration is acceptable:

- **Job is OK to run synchronously (short/moderate runtime, or worker capacity isn't scarce) and the operator supports durable execution** (inherits `ResumableJobMixin` from Section 5, exposes `durable=True`) — recommend replacing the two-task pattern with **one task** in the operator's durable mode: `wait_for_termination=True`, `durable=True` (both already the default). This is not just fewer moving parts — the operator's built-in poll loop is covered by `ResumableJobMixin` end-to-end (submit *and* poll), so a worker crash mid-poll reconnects instead of resubmitting. A hand-rolled sensor task is a separate, non-mixin code path with its own retry semantics — often weaker (check whether it has its own `retries` set; many don't). Collapsing here removes an entire hand-maintained file, not just a task:

  ```python
  run_job = DatabricksSubmitRunOperator(
      task_id="run_job",
      tasks=[{"task_key": "job", "notebook_task": {"notebook_path": NOTEBOOK_PATH}}],
      wait_for_termination=True,   # synchronous — holds the worker for the run, acceptable here
      durable=True,                # default; ResumableJobMixin checkpoints the run id end-to-end
  )
  ```

- **Job is genuinely long-running and worker slots are scarce enough that holding one for the full duration is the real cost to avoid** — keep the `submit` + `mode="reschedule"` sensor split. That is the legitimate, worker-efficient design in this case, not an anti-pattern to remove on sight.

Either way:
- Verify downstream tasks (and any alerting/SLA) depend on the **sensor** task, not `submit`, if the split stays. If something downstream keys off `submit` succeeding, that's the real bug — fix the dependency, not the architecture.
- Name the tradeoff explicitly in the review (worker-slot cost and split outcome vs. single-task correctness and built-in crash-safety) rather than asserting one side is simply "the right pattern."

**Crash-safety and the collapse decision are related, but they're not the same call.** Whether to remove the split is decided by worker-slot cost (above) — that's independent of `durable=True`, which the submit task should generally keep on regardless of whether you collapse (see the "When NOT to use `ResumableJobMixin`" table in Section 5: the poll/sensor task needs nothing, but the submit task's own duplicate-submission risk doesn't disappear just because it's split from the poll task). When you do collapse, prefer the operator's own `durable=True` mode over keeping the hand-rolled sensor bolted on — the built-in path already carries `ResumableJobMixin`'s crash-safety for both submit and poll in one place, which a hand-rolled sensor doesn't replicate.

---

## Section 7 — Configuration reference

```ini
[state_store]
# Full dotted path to the storage backend. Default writes to the Airflow metadata DB.
backend = airflow.state.metastore.MetastoreStateStoreBackend

# Days to retain task state store entries after their last update. 0 = disable time-based cleanup.
# Does NOT affect asset_state_store rows — asset state store has no TTL.
default_retention_days = 30

# Rows deleted per batch during cleanup. 0 = no batching (single unbounded delete).
# Tune on large deployments to reduce lock contention.
state_cleanup_batch_size = 0

# Auto-delete all task state store keys when a task succeeds. Default: False.
# Does NOT affect asset_state_store — asset state store persists across runs and must be cleared explicitly.
clear_on_success = False
```

**Worker-side backend** (optional, `[workers]` section) — routes task state store writes through a local backend before they reach the API server. Useful when large payloads or credentialed storage should stay on the worker:

```ini
[workers]
state_store_backend = mypackage.store.WorkerSideBackend
```

---

## Section 8 — Safety checklist

- [ ] Airflow version ≥ 3.3 (`af config version`)
- [ ] Values are JSON-serializable (`str`, `int`, `float`, `bool`, `list`, `dict` — no `datetime`, no custom objects)
- [ ] `task_state_store` keys are short, descriptive strings (avoid dots and slashes)
- [ ] Mapped tasks writing to `asset_state_store`: use distinct keys per index or accept last-writer-wins semantics
- [ ] Mapped tasks: fleet-wide state clear uses CLI/core API from a downstream task, not `clear()` inside the task body
- [ ] `ResumableJobMixin`: `external_id_key` is set and will not be renamed after deployment
- [ ] `ResumableJobMixin`: `execute()` calls `self.execute_resumable(context)`, not custom logic
- [ ] `ResumableJobMixin`: `durable=False` is intentional if crash recovery is disabled
- [ ] Large payloads (> configured `max_value_storage_bytes`) use a custom `[state_store] backend` or a worker side backend configured via: `[workers] state_store_backend`

---

## Related skills

- **authoring-dags** — general DAG writing patterns and conventions.
- **airflow-hitl** — pausing a DAG for human approval (Airflow 3.1+).
- **airflow** — `af config`, `af registry`, and general Airflow CLI reference.
analyzing-data4.96 KB

View saved version →

---
name: analyzing-data
description: Queries the data warehouse with SQL and answers business questions about data. Use when answering anything that needs warehouse data - counts, metrics, trends, aggregations, joins across tables, data lookups, or ad-hoc SQL analysis (for example "who uses X", "how many Y", "show me Z", "find customers", "what is the count").
---

# Data Analysis

Answer business questions by querying the data warehouse. The kernel auto-starts on first `exec` call.

**All CLI commands below are relative to this skill's directory.** Before running any `scripts/cli.py` command, `cd` to the directory containing this file.

## Workflow

1. **Pattern lookup** — Check for a cached query strategy:
   ```bash
   uv run scripts/cli.py pattern lookup "<user's question>"
   ```
   If a pattern exists, follow its strategy. Record the outcome after executing:
   ```bash
   uv run scripts/cli.py pattern record <name> --success  # or --failure
   ```

2. **Concept lookup** — Find known table mappings:
   ```bash
   uv run scripts/cli.py concept lookup <concept>
   ```

3. **Table discovery** — If cache misses, search the codebase (`Grep pattern="<concept>" glob="**/*.sql"`) or query `INFORMATION_SCHEMA`. See [reference/discovery-warehouse.md](reference/discovery-warehouse.md).

4. **Execute query**:
   ```bash
   uv run scripts/cli.py exec "df = run_sql('SELECT ...')"
   uv run scripts/cli.py exec "print(df)"
   ```

5. **Cache learnings** — Always cache before presenting results:
   ```bash
   # Cache concept → table mapping
   uv run scripts/cli.py concept learn <concept> <TABLE> -k <KEY_COL>
   # Cache query strategy (if discovery was needed)
   uv run scripts/cli.py pattern learn <name> -q "question" -s "step" -t "TABLE" -g "gotcha"
   ```

6. **Present findings** to user.

## Kernel Functions

| Function | Returns |
|----------|---------|
| `run_sql(query, limit=100)` | Polars DataFrame |
| `run_sql_pandas(query, limit=100)` | Pandas DataFrame |
| `run_sql_many(queries, limit=100)` | List of Polars DataFrames (one per query) |

`pl` (Polars) and `pd` (Pandas) are pre-imported.

**Run independent queries together** with `run_sql_many` — they execute concurrently (Snowflake async / connection-pool fan-out) instead of one at a time:

```bash
uv run scripts/cli.py exec "dfs = run_sql_many(['SELECT ...', 'SELECT ...']); print(dfs[0])"
```

`run_sql_many` is **fail-fast**: if any query errors, the call raises and the results of the queries that succeeded are discarded. Use separate `run_sql` calls if you need partial results.

**Timeouts:** `exec` waits up to 120s by default, then interrupts the query and returns a "client stopped waiting" message (the query may still finish server-side). Raise it for known long-running queries: `uv run scripts/cli.py exec "..." -t 600`.

**Idle kernel:** the kernel self-terminates after 2h idle (preserving state until then). Override with `ASTRO_KERNEL_IDLE_TIMEOUT` (seconds; `0` disables).

## CLI Reference

### Kernel

```bash
uv run scripts/cli.py warehouse list      # List warehouses
uv run scripts/cli.py start [-w name]     # Start kernel (with optional warehouse)
uv run scripts/cli.py exec "..."          # Execute Python code
uv run scripts/cli.py status              # Kernel status
uv run scripts/cli.py restart             # Restart kernel
uv run scripts/cli.py stop                # Stop kernel
uv run scripts/cli.py install <pkg>       # Install package
```

### Concept Cache

```bash
uv run scripts/cli.py concept lookup <name>                     # Look up
uv run scripts/cli.py concept learn <name> <TABLE> -k <KEY_COL> # Learn
uv run scripts/cli.py concept list                               # List all
uv run scripts/cli.py concept import -p /path/to/warehouse.md   # Bulk import
```

### Pattern Cache

```bash
uv run scripts/cli.py pattern lookup "question"                                      # Look up
uv run scripts/cli.py pattern learn <name> -q "..." -s "..." -t "TABLE" -g "gotcha"  # Learn
uv run scripts/cli.py pattern record <name> --success                                # Record outcome
uv run scripts/cli.py pattern list                                                   # List all
uv run scripts/cli.py pattern delete <name>                                          # Delete
```

### Table Schema Cache

```bash
uv run scripts/cli.py table lookup <TABLE>            # Look up schema
uv run scripts/cli.py table cache <TABLE> -c '[...]'  # Cache schema
uv run scripts/cli.py table list                       # List cached
uv run scripts/cli.py table delete <TABLE>             # Delete
```

### Cache Management

```bash
uv run scripts/cli.py cache status                # Stats
uv run scripts/cli.py cache clear [--stale-only]  # Clear
```

## References

- [reference/discovery-warehouse.md](reference/discovery-warehouse.md) — Large table handling, warehouse exploration, INFORMATION_SCHEMA queries
- [reference/common-patterns.md](reference/common-patterns.md) — SQL templates for trends, comparisons, top-N, distributions, cohorts

Referenced files: 26

annotating-task-lineage11.1 KB

View saved version →

---
name: annotating-task-lineage
description: Annotate Airflow tasks with data lineage using inlets and outlets. Use when the user wants to add lineage metadata to tasks, specify input/output datasets, or enable lineage tracking for operators without built-in OpenLineage extraction.
---

# Annotating Task Lineage with Inlets & Outlets

This skill guides you through adding manual lineage annotations to Airflow tasks using `inlets` and `outlets`.

> **Reference:** See the [OpenLineage provider developer guide](https://airflow.apache.org/docs/apache-airflow-providers-openlineage/stable/guides/developer.html) for the latest supported operators and patterns.

### On Astro

Lineage annotations defined with inlets and outlets are visualized in Astro's enhanced **Lineage tab**, which provides cross-DAG and cross-deployment lineage views. This means your annotations are immediately visible in the Astro UI, giving you a unified view of data flow across your entire Astro organization.

## When to Use This Approach

| Scenario | Use Inlets/Outlets? |
|----------|---------------------|
| Operator has OpenLineage methods (`get_openlineage_facets_on_*`) | ❌ Modify the OL method directly |
| Operator has no built-in OpenLineage extractor | ✅ Yes |
| Simple table-level lineage is sufficient | ✅ Yes |
| Quick lineage setup without custom code | ✅ Yes |
| Need column-level lineage | ❌ Use OpenLineage methods or custom extractor |
| Complex extraction logic needed | ❌ Use OpenLineage methods or custom extractor |

> **Note:** Inlets/outlets are the lowest-priority fallback. If an OpenLineage extractor or method exists for the operator, it takes precedence. Use this approach for operators without extractors.

---

## Supported Types for Inlets/Outlets

You can use **OpenLineage Dataset** objects or **Airflow Assets** for inlets and outlets:

### OpenLineage Datasets (Recommended)

```python
from openlineage.client.event_v2 import Dataset

# Database tables
source_table = Dataset(
    namespace="postgres://mydb:5432",
    name="public.orders",
)
target_table = Dataset(
    namespace="snowflake://account.snowflakecomputing.com",
    name="staging.orders_clean",
)

# Files
input_file = Dataset(
    namespace="s3://my-bucket",
    name="raw/events/2024-01-01.json",
)
```

### Airflow Assets (Airflow 3+)

```python
from airflow.sdk import Asset

# Using Airflow's native Asset type
orders_asset = Asset(uri="s3://my-bucket/data/orders")
```

### Airflow Datasets (Airflow 2.4+)

```python
from airflow.datasets import Dataset

# Using Airflow's Dataset type (Airflow 2.4-2.x)
orders_dataset = Dataset(uri="s3://my-bucket/data/orders")
```

---

## Basic Usage

### Setting Inlets and Outlets on Operators

```python
from airflow import DAG
from airflow.operators.bash import BashOperator
from openlineage.client.event_v2 import Dataset
import pendulum

# Define your lineage datasets
source_table = Dataset(
    namespace="snowflake://account.snowflakecomputing.com",
    name="raw.orders",
)
target_table = Dataset(
    namespace="snowflake://account.snowflakecomputing.com",
    name="staging.orders_clean",
)
output_file = Dataset(
    namespace="s3://my-bucket",
    name="exports/orders.parquet",
)

with DAG(
    dag_id="etl_with_lineage",
    start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
    schedule="@daily",
) as dag:

    transform = BashOperator(
        task_id="transform_orders",
        bash_command="echo 'transforming...'",
        inlets=[source_table],           # What this task reads
        outlets=[target_table],          # What this task writes
    )

    export = BashOperator(
        task_id="export_to_s3",
        bash_command="echo 'exporting...'",
        inlets=[target_table],           # Reads from previous output
        outlets=[output_file],           # Writes to S3
    )

    transform >> export
```

### Multiple Inputs and Outputs

Tasks often read from multiple sources and write to multiple destinations:

```python
from openlineage.client.event_v2 import Dataset

# Multiple source tables
customers = Dataset(namespace="postgres://crm:5432", name="public.customers")
orders = Dataset(namespace="postgres://sales:5432", name="public.orders")
products = Dataset(namespace="postgres://inventory:5432", name="public.products")

# Multiple output tables
daily_summary = Dataset(namespace="snowflake://account", name="analytics.daily_summary")
customer_metrics = Dataset(namespace="snowflake://account", name="analytics.customer_metrics")

aggregate_task = PythonOperator(
    task_id="build_daily_aggregates",
    python_callable=build_aggregates,
    inlets=[customers, orders, products],      # All inputs
    outlets=[daily_summary, customer_metrics], # All outputs
)
```

---

## Setting Lineage in Custom Operators

When building custom operators, you have two options:

### Option 1: Implement OpenLineage Methods (Recommended)

This is the preferred approach as it gives you full control over lineage extraction:

```python
from airflow.models import BaseOperator


class MyCustomOperator(BaseOperator):
    def __init__(self, source_table: str, target_table: str, **kwargs):
        super().__init__(**kwargs)
        self.source_table = source_table
        self.target_table = target_table

    def execute(self, context):
        # ... perform the actual work ...
        self.log.info(f"Processing {self.source_table} -> {self.target_table}")

    def get_openlineage_facets_on_complete(self, task_instance):
        """Return lineage after successful execution."""
        from openlineage.client.event_v2 import Dataset
        from airflow.providers.openlineage.extractors import OperatorLineage

        return OperatorLineage(
            inputs=[Dataset(namespace="warehouse://db", name=self.source_table)],
            outputs=[Dataset(namespace="warehouse://db", name=self.target_table)],
        )
```

### Option 2: Set Inlets/Outlets Dynamically

For simpler cases, set lineage within the `execute` method (non-deferrable operators only):

```python
from airflow.models import BaseOperator
from openlineage.client.event_v2 import Dataset


class MyCustomOperator(BaseOperator):
    def __init__(self, source_table: str, target_table: str, **kwargs):
        super().__init__(**kwargs)
        self.source_table = source_table
        self.target_table = target_table

    def execute(self, context):
        # Set lineage dynamically based on operator parameters
        self.inlets = [
            Dataset(namespace="warehouse://db", name=self.source_table)
        ]
        self.outlets = [
            Dataset(namespace="warehouse://db", name=self.target_table)
        ]

        # ... perform the actual work ...
        self.log.info(f"Processing {self.source_table} -> {self.target_table}")
```

---

## Dataset Naming Helpers

Use the [OpenLineage dataset naming helpers](https://openlineage.io/docs/client/python/best-practices#dataset-naming-helpers) to ensure consistent naming across platforms:

```python
from openlineage.client.event_v2 import Dataset

# Snowflake
from openlineage.client.naming.snowflake import SnowflakeDatasetNaming

naming = SnowflakeDatasetNaming(
    account_identifier="myorg-myaccount",
    database="mydb",
    schema="myschema",
    table="mytable",
)
dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())
# -> namespace: "snowflake://myorg-myaccount", name: "mydb.myschema.mytable"

# BigQuery
from openlineage.client.naming.bigquery import BigQueryDatasetNaming

naming = BigQueryDatasetNaming(
    project="my-project",
    dataset="my_dataset",
    table="my_table",
)
dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())
# -> namespace: "bigquery", name: "my-project.my_dataset.my_table"

# S3
from openlineage.client.naming.s3 import S3DatasetNaming

naming = S3DatasetNaming(bucket="my-bucket", key="path/to/file.parquet")
dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())
# -> namespace: "s3://my-bucket", name: "path/to/file.parquet"

# PostgreSQL
from openlineage.client.naming.postgres import PostgresDatasetNaming

naming = PostgresDatasetNaming(
    host="localhost",
    port=5432,
    database="mydb",
    schema="public",
    table="users",
)
dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())
# -> namespace: "postgres://localhost:5432", name: "mydb.public.users"
```

> **Note:** Always use the naming helpers instead of constructing namespaces manually. If a helper is missing for your platform, check the [OpenLineage repo](https://github.com/OpenLineage/OpenLineage) or request it.

---

## Precedence Rules

OpenLineage uses this precedence for lineage extraction:

1. **Custom Extractors** (highest) - User-registered extractors
2. **OpenLineage Methods** - `get_openlineage_facets_on_*` in operator
3. **Hook-Level Lineage** - Lineage collected from hooks via `HookLineageCollector`
4. **Inlets/Outlets** (lowest) - Falls back to these if nothing else extracts lineage

> **Note:** If an extractor or method exists but returns no datasets, OpenLineage will check hook-level lineage, then fall back to inlets/outlets.

---

## Best Practices

### Use the Naming Helpers

Always use OpenLineage naming helpers for consistent dataset creation:

```python
from openlineage.client.event_v2 import Dataset
from openlineage.client.naming.snowflake import SnowflakeDatasetNaming


def snowflake_dataset(schema: str, table: str) -> Dataset:
    """Create a Snowflake Dataset using the naming helper."""
    naming = SnowflakeDatasetNaming(
        account_identifier="mycompany",
        database="analytics",
        schema=schema,
        table=table,
    )
    return Dataset(namespace=naming.get_namespace(), name=naming.get_name())


# Usage
source = snowflake_dataset("raw", "orders")
target = snowflake_dataset("staging", "orders_clean")
```

### Document Your Lineage

Add comments explaining the data flow:

```python
transform = SqlOperator(
    task_id="transform_orders",
    sql="...",
    # Lineage: Reads raw orders, joins with customers, writes to staging
    inlets=[
        snowflake_dataset("raw", "orders"),
        snowflake_dataset("raw", "customers"),
    ],
    outlets=[
        snowflake_dataset("staging", "order_details"),
    ],
)
```

### Keep Lineage Accurate

- Update inlets/outlets when SQL queries change
- Include all tables referenced in JOINs as inlets
- Include all tables written to (including temp tables if relevant)
- **Outlet-only and inlet-only annotations are valid.** One-sided annotations are encouraged for lineage visibility even without a corresponding inlet or outlet in another DAG.

---

## Limitations

| Limitation | Workaround |
|------------|------------|
| Table-level only (no column lineage) | Use OpenLineage methods or custom extractor |
| Overridden by extractors/methods | Only use for operators without extractors |
| Static at DAG parse time | Set dynamically in `execute()` or use OL methods |
| Deferrable operators lose dynamic lineage | Use OL methods instead; attributes set in `execute()` are lost when deferring |

---

## Related Skills

- **creating-openlineage-extractors**: For column-level lineage or complex extraction
- **tracing-upstream-lineage**: Investigate where data comes from
- **tracing-downstream-lineage**: Investigate what depends on data
authoring-dags6.91 KB

View saved version →

---
name: authoring-dags
description: Workflow and best practices for writing Apache Airflow DAGs. Use when creating a new DAG, write pipeline code, handling questions about DAG patterns and conventions or extending an existing DAG with a follow-up/downstream task. ANY request shaped like 'add a DAG named X', 'write a pipeline', 'add a task that runs after Y', or 'extend the DAG'. For testing and debugging DAGs, see the testing-dags skill.
hooks:
  Stop:
    - hooks:
        - type: command
          command: "echo 'Remember to test your DAG with the testing-dags skill'"
---

# DAG Authoring Skill

This skill guides you through creating and validating Airflow DAGs using best practices and `af` CLI commands.

> **For testing and debugging DAGs**, see the **testing-dags** skill which covers the full test -> debug -> fix -> retest workflow.

---

## Running the CLI

These commands assume `af` is on PATH. Run via `astro otto` to get it automatically, or install standalone with `uv tool install astro-airflow-mcp`.

---

## Workflow Overview

```
+-----------------------------------------+
| 1. DISCOVER                             |
|    Understand codebase & environment    |
+-----------------------------------------+
                 |
+-----------------------------------------+
| 2. PLAN                                 |
|    Propose structure, get approval      |
+-----------------------------------------+
                 |
+-----------------------------------------+
| 3. IMPLEMENT                            |
|    Write DAG following patterns         |
+-----------------------------------------+
                 |
+-----------------------------------------+
| 4. VALIDATE                             |
|    Check import errors, warnings        |
+-----------------------------------------+
                 |
+-----------------------------------------+
| 5. TEST (with user consent)             |
|    Trigger, monitor, check logs         |
+-----------------------------------------+
                 |
+-----------------------------------------+
| 6. ITERATE                              |
|    Fix issues, re-validate              |
+-----------------------------------------+
```

---

## Phase 1: Discover

Before writing code, understand the context.

### Explore the Codebase

Use file tools to find existing patterns:
- `Glob` for `**/dags/**/*.py` to find existing DAGs
- `Read` similar DAGs to understand conventions
- Check `requirements.txt` for available packages

### Query the Airflow Environment

Use `af` CLI commands to understand what's available:

| Command | Purpose |
|---------|---------|
| `af config connections` | What external systems are configured |
| `af config variables` | What configuration values exist |
| `af config providers` | What operator packages are installed |
| `af config version` | Version constraints and features |
| `af dags list` | Existing DAGs and naming conventions |
| `af config pools` | Resource pools for concurrency |

**Example discovery questions:**
- "Is there a Snowflake connection?" -> `af config connections`
- "What Airflow version?" -> `af config version`
- "Are S3 operators available?" -> `af config providers`

---

## Phase 2: Plan

Based on discovery, propose:

1. **DAG structure** - Tasks, dependencies, schedule
2. **Operators to use** - Based on available providers
3. **Connections needed** - Existing or to be created
4. **Variables needed** - Existing or to be created
5. **Packages needed** - Additions to requirements.txt

**Get user approval before implementing.**

---

## Phase 3: Implement

Write the DAG following best practices (see below). Key steps:

1. Create DAG file in appropriate location
2. Update `requirements.txt` if needed
3. Save the file

---

## Phase 4: Validate

**Use `af` CLI as a feedback loop to validate your DAG.**

### Step 1: Check Import Errors

After saving, check for parse errors (Airflow will have already parsed the file):

```bash
af dags errors
```

- If your file appears -> **fix and retry**
- If no errors -> **continue**

Common causes: missing imports, syntax errors, missing packages.

### Step 2: Verify DAG Exists

```bash
af dags get <dag_id>
```

Check: DAG exists, schedule correct, tags set, paused status.

### Step 3: Check Warnings

```bash
af dags warnings
```

Look for deprecation warnings or configuration issues.

### Step 4: Explore DAG Structure

```bash
af dags explore <dag_id>
```

Returns in one call: metadata, tasks, dependencies, source code.

### On Astro

If you're running on Astro, you can also validate locally before deploying:

- **Parse check**: Run `astro dev parse` to catch import errors and DAG-level issues without starting a full Airflow environment
- **DAG-only deploy**: Once validated, use `astro deploy --dags` for fast DAG-only deploys that skip the Docker image build — ideal for iterating on DAG code

---

## Phase 5: Test

> See the **testing-dags** skill for comprehensive testing guidance.

Once validation passes, test the DAG using the workflow in the **testing-dags** skill:

1. **Get user consent** -- Always ask before triggering
2. **Trigger and wait** -- `af runs trigger-wait <dag_id> --timeout 300`
3. **Analyze results** -- Check success/failure status
4. **Debug if needed** -- `af runs diagnose <dag_id> <run_id>` and `af tasks logs <dag_id> <run_id> <task_id>`

### Quick Test (Minimal)

```bash
# Ask user first, then:
af runs trigger-wait <dag_id> --timeout 300
```

For the full test -> debug -> fix -> retest loop, see **testing-dags**.

---

## Phase 6: Iterate

If issues found:
1. Fix the code
2. Check for import errors: `af dags errors`
3. Re-validate (Phase 4)
4. Re-test using the **testing-dags** skill workflow (Phase 5)

---

## CLI Quick Reference

| Phase | Command | Purpose |
|-------|---------|---------|
| Discover | `af config connections` | Available connections |
| Discover | `af config variables` | Configuration values |
| Discover | `af config providers` | Installed operators |
| Discover | `af config version` | Version info |
| Validate | `af dags errors` | Parse errors (check first!) |
| Validate | `af dags get <dag_id>` | Verify DAG config |
| Validate | `af dags warnings` | Configuration warnings |
| Validate | `af dags explore <dag_id>` | Full DAG inspection |

> **Testing commands** -- See the **testing-dags** skill for `af runs trigger-wait`, `af runs diagnose`, `af tasks logs`, etc.

---

## Best Practices & Anti-Patterns

For code patterns and anti-patterns, see **[reference/best-practices.md](reference/best-practices.md)**.

**Read this reference when writing new DAGs or reviewing existing ones.** It covers what patterns are correct (including Airflow 3-specific behavior) and what to avoid.

---

## Related Skills

- **testing-dags**: For testing DAGs, debugging failures, and the test -> fix -> retest loop
- **debugging-dags**: For troubleshooting failed DAGs
- **deploying-airflow**: For deploying DAGs to production (Astro or open-source)
- **migrating-airflow-2-to-3**: For migrating DAGs to Airflow 3

Referenced files: 1

authoring-go-sdk-tasks8.35 KB

View saved version →

---
name: authoring-go-sdk-tasks
description: Writes Airflow task logic in Go using the Airflow Go SDK. Use when the user wants to implement Airflow tasks in Go, asks about `BundleProvider`/`RegisterDags`, the `bundlev1` Registry/Dag interfaces, registering Go tasks (`AddTask`/`AddTaskWithName`), dependency injection by parameter type (`context.Context`, `sdk.TIRunContext`, `*slog.Logger`, `sdk.Client`), or reading connections/variables/XComs from Go. This skill covers the Go-specific native API; the shared Python-stub pattern and conceptual model live in authoring-language-sdk-tasks. For building/packing/shipping the bundle see deploying-go-sdk-bundles; for coordinator config see configuring-airflow-language-sdks.
---

# Authoring Go SDK Tasks

The Airflow Go SDK implements the language-SDK model for Go: your DAG stays in Python, and each task is a compiled Go function registered inside a **bundle** (a single native executable). This skill covers the **Go-specific** native API. The shared model (the Python `@task.stub` pattern, ID matching, the XCom-as-JSON contract) lives in **authoring-language-sdk-tasks**; read that first if you are new to language SDKs.

> **Experimental.** The Go SDK is under active development and not production-ready. Module path `github.com/apache/airflow/go-sdk` (Go 1.24+). APIs may change.

> **Related skills:** **authoring-language-sdk-tasks** (shared Python stub + concepts), **deploying-go-sdk-bundles** (build, pack, and ship the bundle), **configuring-airflow-language-sdks** (route the queue to the Go coordinator).

---

## Recap: the Python side

A Go task is paired with a Python stub that carries no logic; it declares the task, its queue, and the dependency graph. IDs must match the Go registration exactly, and `queue=` routes the task to the Go runtime. Full rules are in **authoring-language-sdk-tasks**; the minimal shape:

```python
from airflow.sdk import dag, task


@task.stub(queue="golang")
def extract(): ...


@task.stub(queue="golang")
def transform(): ...


@dag()
def simple_dag():
    extract() >> transform()


simple_dag()
```

The `queue` value (`"golang"` here) is an arbitrary label that must match the queue routed to the Go coordinator (`queue_to_coordinator`). See **configuring-airflow-language-sdks**.

---

## The bundle entry point

A bundle implements `bundlev1.BundleProvider`: report its version and register your DAGs and tasks. `main` is one line; `bundlev1server.Serve` wires the bundle to the Airflow runtime for you.

```go
package main

import (
	"log"

	v1 "github.com/apache/airflow/go-sdk/bundle/bundlev1"
	"github.com/apache/airflow/go-sdk/bundle/bundlev1/bundlev1server"
)

type myBundle struct{}

var _ v1.BundleProvider = (*myBundle)(nil)

func (m *myBundle) GetBundleVersion() v1.BundleInfo {
	return v1.BundleInfo{Name: bundleName, Version: &bundleVersion}
}

func (m *myBundle) RegisterDags(dagbag v1.Registry) error {
	simpleDag := dagbag.AddDag("simple_dag")      // dag_id must match the Python @dag name
	simpleDag.AddTask(extract)                    // task_id is the function name; must match the stub
	simpleDag.AddTaskWithName("transform", transform) // or set the task_id explicitly
	return nil
}

func main() {
	if err := bundlev1server.Serve(&myBundle{}); err != nil {
		log.Fatal(err)
	}
}
```

`AddTask(fn)` derives the `task_id` from the Go function's name; use `AddTaskWithName("<task_id>", fn)` when that name can't match the Python stub (an unexported, renamed, or reused function). `RegisterDags` is the single source of truth for task identity: the bundle's manifest (used by the packer and by the coordinator) is generated by running it, never hand-written.

---

## Task functions: dependency injection by parameter type

A task is an ordinary Go function. The runtime inspects its signature and injects arguments **by type**; declare only what you need.

| Parameter type | Injected value |
|----------------|----------------|
| `context.Context` | Task context for cancellation. Always available. |
| `sdk.TIRunContext` | Richer context (embeds `context.Context`) exposing `TaskInstance()` and `DagRun()`. See [Runtime context](#runtime-context). |
| `*slog.Logger` | Logger wired to the Airflow task log. |
| `sdk.Client` | Full Airflow model access: Variables, Connections, XComs. |
| `sdk.VariableClient` / `sdk.ConnectionClient` / `sdk.XComClient` | A narrower slice of `sdk.Client`. Prefer the narrowest you need; it documents intent and is trivial to fake in tests. |

The optional return signature is `(result, error)`: a non-nil `result` is pushed as the task's `return_value` XCom; a non-nil `error` fails the task (which triggers the stub's retry policy). Returning only `error`, or nothing, is also valid.

```go
func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) {
	conn, err := client.GetConnection(ctx, "test_http")
	if err != nil {
		return nil, err
	}
	log.Info("connected", "host", conn.Host)
	return map[string]any{"go_version": runtime.Version()}, nil
}

func transform(ctx sdk.TIRunContext, client sdk.VariableClient) error {
	val, err := client.GetVariable(ctx, "my_variable")
	if err != nil {
		return err // VariableNotFound (a sentinel error) if absent
	}
	_ = val
	return nil
}
```

---

## The `sdk.Client` surface

| Call | Returns | Notes |
|------|---------|-------|
| `GetVariable(ctx, key)` | `(string, error)` | `VariableNotFound` if absent. |
| `UnmarshalJSONVariable(ctx, key, &ptr)` | `error` | Decode a JSON variable into a struct/pointer. |
| `GetConnection(ctx, connID)` | `(Connection, error)` | `ConnectionNotFound` if absent. |
| `GetXCom(ctx, dagID, runID, taskID, mapIndex, key, value)` | `(any, error)` | `XComNotFound` only if the key is absent; a stored null returns `(nil, nil)`. |
| `PushXCom(ctx, ti, key, value)` | `error` | Rarely needed; a returned value is pushed for you. |

`Connection` exposes `ID`, `Type`, `Host`, `Port` (`int`), `Login *string`, `Password *string` (nil when unset, distinct from empty), `Path` (schema), `Extra map[string]any`, plus `GetURI()`. Not-found cases return the sentinels `sdk.VariableNotFound`, `sdk.ConnectionNotFound`, `sdk.XComNotFound`.

To read an upstream task's result, call `GetXCom` explicitly, taking the `dag_id`/`run_id`/`task_id` you need from the runtime context (below).

---

## Runtime context

Declare an `sdk.TIRunContext` parameter to read metadata about the task instance and its DAG run. It is an interface that embeds `context.Context`, so it is usable anywhere a `context.Context` is expected.

```go
func extract(ctx sdk.TIRunContext, log *slog.Logger) error {
	ti, dagRun := ctx.TaskInstance(), ctx.DagRun()
	log.Info("running",
		"task_id", ti.TaskID,
		"run_id", dagRun.RunID,
		"logical_date", dagRun.LogicalDate)
	return nil
}
```

- `TaskInstance()`: `DagID`, `RunID`, `TaskID`, `MapIndex *int` (nil when unmapped), `TryNumber`.
- `DagRun()`: `DagID`, `RunID`, and the `*time.Time` timestamps `LogicalDate`, `DataIntervalStart`, `DataIntervalEnd` (nil when not sent).

The accessors are populated from the task's startup details before the body runs. Because `TIRunContext` embeds `context.Context`, pass it straight to client calls and cancellation checks (`ctx.Done()`); declare it as your context parameter by default. In tests, build the argument with `sdk.NewTIRunContext(ctx, ti, dagRun)` (it panics on a nil `ctx`).

---

## Go-specific pitfalls

- **IDs must match the Python stub** (`dag_id` from `AddDag`, `task_id` from the registered function name), and the stub's `queue=` must route to the Go coordinator, or the task is never delivered.
- **`RegisterDags` is authoritative.** Do not hand-write the manifest; the packer generates it by running `RegisterDags`.
- **Ask for the narrowest client interface** you need (`sdk.VariableClient` over `sdk.Client`) for clearer intent and easier fakes.
- **A non-nil `error` return fails the task** and applies the stub's retries; a recovered panic is also a failure.
- See **authoring-language-sdk-tasks** for the language-agnostic pitfalls (one process per task instance, set queue and retries on the stub).

---

## Related Skills

- **authoring-language-sdk-tasks**: Shared Python-stub pattern and concepts (read first).
- **deploying-go-sdk-bundles**: Build and pack the bundle with `go tool airflow-go-pack`, then deploy it for the coordinator.
- **configuring-airflow-language-sdks**: Route the queue to the Go coordinator (`ExecutableCoordinator`).
- **authoring-dags**: General Airflow DAG authoring.
authoring-java-sdk-tasks10.2 KB

View saved version →

---
name: authoring-java-sdk-tasks
description: Writes Airflow task logic in Java, Kotlin, or any JVM language using the Airflow Java SDK. Use when the user wants to implement Airflow tasks in Java/JVM, asks about `@Builder.Dag`/`@Builder.Task`/`@Builder.XCom`, the `Task`/`BundleBuilder` interfaces, reading connections/variables/XComs from Java, the JSON-to-Java type mapping, or logging from Java tasks. This skill covers the Java-specific native API; the shared Python-stub pattern and conceptual model live in authoring-language-sdk-tasks. For building/shipping the bundle see deploying-java-sdk-bundles; for coordinator config see configuring-airflow-language-sdks.
---

# Authoring Java SDK Tasks

The Airflow Java SDK implements the language-SDK model for the JVM: your DAG stays in Python, and each task instance runs in a short-lived JVM subprocess. This skill covers the **Java-specific** native API. The shared model — the Python `@task.stub` pattern, ID matching, and the XCom-as-JSON contract — lives in **authoring-language-sdk-tasks**; read that first if you're new to language SDKs.

> **Experimental.** The Java SDK is in preview. Artifact coordinates and APIs may change.

> **Related skills:** **authoring-language-sdk-tasks** (shared Python stub + concepts), **configuring-airflow-language-sdks** (route the queue to `JavaCoordinator`), **deploying-java-sdk-bundles** (compile and ship the JAR).

---

## Recap: the Python side

Java tasks are paired with Python stubs that carry no logic — they declare the task, queue, dependency graph, and retries. IDs must match the Java annotations exactly, and an upstream argument on a stub only declares the dependency (the value is fetched in Java). Full rules are in **authoring-language-sdk-tasks**; the minimal shape:

```python
from airflow.sdk import dag, task


@dag
def sales_pipeline():                     # dag_id "sales_pipeline" -> @Builder.Dag(id="sales_pipeline")
    @task.stub(queue="java")
    def extract(): ...                    # task_id "extract" -> @Builder.Task(id="extract")

    @task.stub(queue="java")
    def transform(extracted): ...

    transform(extract())


sales_pipeline()
```

---

## Java side: two APIs

Both APIs produce identical runtime behavior; pick by style, and you can mix them in one bundle.

### Annotation-based API (recommended)

Annotate a plain class; an annotation processor generates the wiring (`<ClassName>Builder`) at compile time.

```java
import static java.lang.System.Logger.Level.INFO;
import org.apache.airflow.sdk.*;

@Builder.Dag(id = "sales_pipeline")          // must match the Python dag_id
public class SalesPipeline {
  private static final System.Logger log = System.getLogger(SalesPipeline.class.getName());

  @Builder.Task(id = "extract")              // must match the Python @task.stub name
  public long extract(Client client) {
    var conn = client.getConnection("sales_db");
    log.log(INFO, "connected to {0}", conn.host);
    return 42L;                              // return value is pushed as the return_value XCom
  }

  @Builder.Task(id = "transform")
  public long transform(
      Client client,
      @Builder.XCom(task = "extract") long recordCount) {  // pulls extract's return_value
    var threshold = (String) client.getVariable("transform_threshold");
    return recordCount * 2;
  }

  @Builder.Task   // id omitted -> the method name "load" is used
  public void load(Context context, @Builder.XCom(task = "transform") long transformed) {
    log.log(INFO, "attempt {0}, value {1}", context.ti.tryNumber, transformed);
  }
}
```

Annotation reference:

| Annotation | Purpose |
|------------|---------|
| `@Builder.Dag(id = "...")` | Marks the class as a task container. `id` must match the Python `dag_id`; if omitted, the class name is used. Optional `to = "..."` renames the generated builder (default `<ClassName>Builder`). |
| `@Builder.Task(id = "...")` | Marks a method as a task. `id` must match the Python `@task.stub` function name; if omitted, the method name is used. |
| `@Builder.XCom(task = "...", key = "...")` | Injects an upstream task's XCom as a parameter. `task` defaults to the parameter name; `key` defaults to the producing task's `return_value`. The parameter type must be compatible with the stored JSON value. |

A task method's return value is automatically pushed as that task's `return_value` XCom. A method may declare `throws Exception`; any uncaught exception fails the task instance (which triggers retries if the stub configured them).

### Interface-based API

Implement `Task` directly when you want full control over registration and XCom handling.

```java
import org.apache.airflow.sdk.*;

public class ExtractTask implements Task {
  @Override
  public void execute(Context context, Client client) throws Exception {
    var conn = client.getConnection("sales_db");
    // ... do work ...
    client.setXCom(42L);   // push return_value explicitly
  }
}
```

Register tasks manually in a `Dag` and expose it through a `BundleBuilder`:

```java
public class MyBundle implements BundleBuilder {
  @Override
  public Iterable<Dag> getDags() {
    var dag = new Dag("sales_pipeline");      // DAG ID matches Python
    dag.addTask("extract", ExtractTask.class);
    dag.addTask("transform", TransformTask.class);
    return java.util.List.of(dag);
  }
}
```

Each `Task` class needs a public no-arg constructor. Task IDs must be unique within a DAG, and DAG IDs unique within a bundle.

---

## The entry point

Every bundle has a `main` that hands your DAGs to the SDK server. The server connects to the coordinator, runs one task instance, and exits.

```java
import java.util.List;
import org.apache.airflow.sdk.*;

public class Main implements BundleBuilder {
  @Override
  public Iterable<Dag> getDags() {
    // With the annotation API, the *Builder classes are generated at compile time.
    return List.of(SalesPipelineBuilder.build());
  }

  public static void main(String[] args) {
    Server.create(args).serve(new Main().build());
  }
}
```

`Server.create(args)` parses the connection details Airflow passes on the command line — don't construct them by hand. Record this `main` class as the bundle's main class when you build it (see **deploying-java-sdk-bundles**).

---

## Talking to Airflow from a task: `Client`

A `Client` is passed into every task and is scoped to the current DAG run and task instance.

| Call | Returns | Notes |
|------|---------|-------|
| `client.getConnection(id)` | `Connection` | Fields: `id`, `type`, `host`, `schema`, `login`, `password`, `port`, `extra`. Any unset field is `null`. Throws if the connection doesn't exist. |
| `client.getVariable(key)` | `Object` (or `null`) | Cast to the type you expect, e.g. `(String) client.getVariable("threshold")`. |
| `client.getXCom(taskId)` | `Object` (or `null`) | Reads another task's `return_value` by default. Overloads accept `key`, `dagId`, `runId`, `mapIndex`, and `includePriorDates` for cross-DAG/run reads and mapped tasks. |
| `client.setXCom(value)` | — | Pushes the `return_value` XCom (interface API). Value must be JSON-serializable. With the annotation API, returning a value does this for you. |

### `Context`

The `Context` parameter exposes run metadata: `context.dagRun` (`dagId`, `runId`) and `context.ti` (`dagId`, `runId`, `taskId`, `mapIndex`, `tryNumber`). `tryNumber` is useful for retry-aware logic.

---

## XCom: Java types

XComs cross the boundary as JSON (the shared contract is in **authoring-language-sdk-tasks**). When you read one back in Java you get:

| Python type | JSON | Java type from `getXCom` |
|-------------|------|--------------------------|
| `int` | integer | `Long` (or `BigInteger` if too large) |
| `float` | decimal | `Double` |
| `str` | string | `String` |
| `bool` | boolean | `Boolean` |
| `None` | null | `null` |
| `list` | array | `List<Object>` |
| `dict` | object | `Map<String, Object>` |

Declare `@Builder.XCom` parameter types to match. A mismatch (e.g. declaring `int` when the value is a `String`) fails the task.

---

## Logging

Declare a logger as a static field named after the class — the conventional pattern regardless of framework:

```java
private static final System.Logger log = System.getLogger(SalesPipeline.class.getName());
```

For records to reach Airflow's task log store (and show in the UI), the bundle must include one of the SDK logging integration artifacts (`airflow-sdk-jpl`, `airflow-sdk-slf4j`, `airflow-sdk-log4j2`, or `airflow-sdk-jul`). The dependencies and per-framework setup are in the logging integration section of **deploying-java-sdk-bundles**. `System.Logger` (JPL) with `airflow-sdk-jpl` is the lightest option and needs no configuration.

---

## A complete worked example ships with the SDK

The SDK repository includes a runnable example under `java-sdk/example/`:

- `src/resources/dags/java_examples.py` — Python DAGs pairing Python tasks with Java stubs, including a `load` stub with `retries=1`.
- `src/java/.../AnnotationExample.java` — annotation API, including a task that fails on `tryNumber == 1` and succeeds on retry.
- `src/java/.../InterfaceExampleBuilder.java` — the same tasks via the `Task` interface and `Dag.addTask(...)`.
- `src/java/.../ExampleBundleBuilder.java` — a `BundleBuilder` returning both DAGs plus the `main` entry point.

Point users there for an end-to-end reference.

---

## Java-specific pitfalls

- **Cast `Object` returns deliberately.** `getVariable` and `getXCom` return `Object`; match the cast to the JSON type (see the table above).
- **`@Builder.XCom` parameter types must match the stored JSON type**, or the task fails at runtime.
- **The annotation processor must be on the build** for the annotation API (generates `<ClassName>Builder`); it is not needed for the interface API. See **deploying-java-sdk-bundles**.
- See **authoring-language-sdk-tasks** for the language-agnostic pitfalls (ID matching, one JVM per task instance, queue/retries on the stub).

---

## Related Skills

- **authoring-language-sdk-tasks**: Shared Python-stub pattern and concepts (read first).
- **configuring-airflow-language-sdks**: Route the `java` queue to `JavaCoordinator` and set JRE/coordinator options.
- **deploying-java-sdk-bundles**: Build the bundle (Gradle/Maven) and place the JAR where Airflow can find it.
- **authoring-dags**: General Airflow DAG authoring.
authoring-language-sdk-tasks7.16 KB

View saved version →

---
name: authoring-language-sdk-tasks
description: The language-neutral foundation for Airflow language SDKs — implement task logic in a non-Python language while the DAG stays in Python. Use when the user wants to run an Airflow task in another language (Java, Kotlin, Go, or other JVM/native languages), asks how the Python `@task.stub` pairs with native task code, how task/DAG IDs must match across the two sides, how data passes via XCom as JSON, or which language SDKs exist. This skill owns the shared Python-stub pattern and conceptual model; for a specific language's native API, build, and runtime, use that language's skill (e.g. authoring-java-sdk-tasks, authoring-go-sdk-tasks).
---

# Authoring Language SDK Tasks (Shared Foundation)

Airflow language SDKs let you implement task logic in a language other than Python while the DAG and its scheduling stay in Python. This skill describes the parts that are identical across every language SDK. Each language has its own companion skill for the native API, build tooling, and runtime — see [Per-language skills](#per-language-skills).

> **Experimental.** The language SDKs are in preview. APIs and artifact coordinates may change.

---

## The model

A DAG is authored in Python as usual. Tasks that should run in another language are declared as **stubs** routed to a dedicated queue. At runtime, Airflow hands a stub task to a **coordinator** that launches a short-lived **native subprocess** for that one task instance, runs your compiled/native code, and shuts the subprocess down.

Consequences that hold for every language SDK:

- **One subprocess per task instance** — there is no shared in-process state between task instances. Pass data via XCom or an external store.
- **The DAG, schedule, retries, and queue routing live in Python.** The native side only implements task logic.
- **Data crossing the boundary is JSON.** See [The XCom-as-JSON contract](#the-xcom-as-json-contract).

---

## The two-sided model

Every task has two halves that must agree:

1. A **Python stub** in a normal DAG file — no logic; it declares the task, its queue, the dependency graph, and retry policy.
2. A **native implementation** (Java, Go, etc.) whose IDs match the Python side and where the work happens.

### Python side (scheduling)

The example below uses the Go SDK to be concrete, but the Python side is **identical for every language SDK**. The queue name (`"golang"` here) is an arbitrary label you choose — it just has to match a key in `queue_to_coordinator` (see **configuring-airflow-language-sdks**). Pick whatever name fits the SDK you're routing to.

```python
from datetime import timedelta
from airflow.sdk import dag, task


@dag
def sales_pipeline():
    @task.stub(queue="golang")          # queue selects the coordinator (see configuring-airflow-language-sdks)
    def extract(): ...

    @task.stub(queue="golang")
    def transform(extracted): ...        # arg only declares the dependency

    @task.stub(queue="golang", retries=1, retry_delay=timedelta(seconds=5))
    def load(transformed): ...

    @task()                              # an ordinary Python task can sit downstream
    def report(loaded):
        print(f"done: {loaded}")

    report(load(transform(extract())))


sales_pipeline()
```

Rules that apply regardless of language:

- The **stub function name is the task ID** and the `@dag` name (or `dag_id=`) is the DAG ID. The native side must use these exact IDs.
- An upstream argument on a stub (e.g. `transform(extracted)`) exists **only to declare the dependency** in Python. The value itself is fetched on the native side via XCom — passing it in Python does not hand it to the native code.
- **Queue, retries, and other task arguments are set on the stub**, not in the native code. A native task that fails is reported back to Airflow, which then applies the stub's retry policy.
- The `queue` value is what routes the task to a coordinator; the same string must appear in `queue_to_coordinator` (see **configuring-airflow-language-sdks**).

---

## The XCom-as-JSON contract

XCom values are stored as JSON in Airflow's metadata database, so the boundary between Python and any native language is JSON. The Python/JSON side is the same for every SDK:

| Python type | JSON |
|-------------|------|
| `int` | number (integer) |
| `float` | number (decimal) |
| `str` | string |
| `bool` | boolean |
| `None` | null |
| `list` | array |
| `dict` | object |

Each language SDK maps these JSON types onto its own native types (e.g. a JSON integer becomes a Java `Long`). The native-type mapping lives in that language's skill. The key portability rule: a value pushed by one task is read by another **as JSON**, so the consuming side must expect a type compatible with what was stored.

---

## What is language-specific (and lives elsewhere)

This skill deliberately stops at the shared concepts. The following differ per language and are documented in each language's companion skills:

- **Native task API** — how you declare tasks, read connections/variables/XComs, and push results (annotations, interfaces, function registration, etc.).
- **Native type mapping** — the native column of the JSON table above.
- **Build and packaging** — how the artifact is compiled and bundled.
- **Runtime prerequisite** — what must be present on the worker (a language runtime for some SDKs, e.g. a JRE for the Java SDK; none for the Go SDK's self-contained bundles).

The Airflow-side wiring (which coordinator runs which queue) is shared in structure but has per-coordinator options; it lives in **configuring-airflow-language-sdks**.

---

## Language-agnostic pitfalls

- **IDs must match exactly** across the Python stub function name and the native task ID, and across `@dag`/`dag_id` and the native DAG ID. Mismatches surface as "no DAGs" or missing-XCom errors.
- **Both sides need the upstream reference.** Python declares the dependency by passing the upstream call; the native code retrieves the value via XCom.
- **Set queue and retries on the stub**, never in the native code.
- **Stub bodies must be empty.** An AST check enforces it — only `pass`, `...`, or a docstring is allowed in the body; any real logic is rejected.
- **`retry_policy` is rejected on stubs** (`@task.stub` raises `ValueError`). Use `retries`/`retry_delay` instead — a retry-policy callable runs Python in-process and would never fire for a task executing in a native subprocess.
- Assets, deferral, and some other Airflow features have limited or no support in the language SDKs today.

---

## Per-language skills

- **authoring-java-sdk-tasks**: Java/Kotlin/JVM native API, type mapping, and logging.
- **authoring-go-sdk-tasks**: Go native API — task registration, dependency injection by parameter type, and client access.
- *(Future language SDKs each add their own `authoring-<lang>-sdk-tasks` skill that builds on this one.)*

## Related Skills

- **configuring-airflow-language-sdks**: Route a queue to a coordinator and set runtime options.
- **authoring-dags**: General Airflow DAG authoring (the Python side lives here too).
- **deploying-java-sdk-bundles**: Build and ship the Java artifact.
- **deploying-go-sdk-bundles**: Build, pack, and ship the Go bundle (per-language deploy skills follow the same shape).
blueprint32.9 KB

View saved version →

---
name: blueprint
description: Define reusable Airflow task group templates with Pydantic validation and compose DAGs from YAML. Use when creating blueprint templates, composing DAGs from YAML, declaring shared variables or per-environment profiles, validating configurations, sharing templates as an installable package, or enabling no-code DAG authoring for non-engineers.
---

# Blueprint Implementation

You are helping a user work with Blueprint, a system for composing Airflow DAGs from YAML using reusable Python templates. Execute steps in order and prefer the simplest configuration that meets the user's needs.

> **Package**: `airflow-blueprint` on PyPI — this skill documents **0.5.0**
> **Repo**: https://github.com/astronomer/blueprint
> **Requires**: Python 3.10+, Airflow 2.5+
> **Cross-references**: the `airflow` skill for Astro CLI, registry, and REST API discovery commands; `authoring-dags` or `dag-factory` when the user needs full Airflow flexibility instead of validated templates.

---

## Determine What the User Needs

| User Request | Action |
|--------------|--------|
| "Create a blueprint" / "Define a template" | Go to **Creating Blueprints** |
| "Build a template from other templates" | Go to **Composing Templates** |
| "Create a DAG from YAML" / "Compose steps" | Go to **Composing DAGs in YAML** |
| "Reuse a value across steps or DAGs" / "Different value per environment" | Go to **Variables and Profiles** |
| "Use a blueprint in an existing Python DAG" / "Generate DAGs in a loop" | Go to **Blueprints in Python DAGs** |
| "Customize DAG args" / "Add tags to DAG" / "Different DAG defaults per folder" | Go to **Customizing DAG-Level Configuration** |
| "Share templates across repos" / "Install blueprints from a package" | Go to **Sharing Blueprints as a Package** |
| "Override config at runtime" / "Trigger with params" | Go to **Runtime Parameter Overrides** |
| "Post-process DAGs" / "Add callback" / "Don't let one bad file break everything" | Go to **Loader Options** |
| "Validate my YAML" / "Lint blueprint" | Go to **Validation Commands** |
| "Set up blueprint in my project" | Go to **Project Setup** |
| "Version my blueprint" | Go to **Versioning** |
| "Generate schema" / "Astro IDE setup" | Go to **Schema Generation** |
| Blueprint errors / troubleshooting | Go to **Troubleshooting** |

---

## Project Setup

If the user is starting fresh, guide them through setup:

### 1. Install the Package

Add `airflow-blueprint>=0.5.0` to `requirements.txt`.

### 2. Create the Loader

Create `dags/loader.py`:

```python
from blueprint import build_all_airflow_dags

build_all_airflow_dags()
```

> **The function name matters.** Airflow's safe-mode DAG file processor only parses files containing both `airflow` and `dag`, so the import line itself is what makes the loader discoverable. `build_all` and `build_all_dags` still work as deprecated aliases that emit `DeprecationWarning`; migrate existing loaders to `build_all_airflow_dags`.

DAG-level configuration (schedule, description, tags, default_args, etc.) is handled via YAML fields and `BlueprintDagArgs` templates — see **Customizing DAG-Level Configuration**.

### 3. Verify Installation

Run `blueprint list` from the project root. If no blueprints are found, the user needs to create blueprint classes first.

---

## Creating Blueprints

### Canonical Example

Config model, generic base class, and a `render()` returning a task or group keyed on `self.step_id`. Adapt this rather than inventing a different structure:

```python
# dags/templates/my_blueprints.py
from airflow.operators.bash import BashOperator
from airflow.utils.task_group import TaskGroup
from blueprint import Blueprint, BaseModel, Field

class MyConfig(BaseModel):
    source_table: str = Field(description="Source table name")
    batch_size: int = Field(default=1000, ge=1)

class MyBlueprint(Blueprint[MyConfig]):
    """Docstring becomes blueprint description."""

    def render(self, config: MyConfig) -> TaskGroup:
        with TaskGroup(group_id=self.step_id) as group:
            BashOperator(
                task_id="my_task",
                bash_command=f"echo '{config.source_table}'"
            )
        return group
```

### Key Rules

| Element | Requirement |
|---------|-------------|
| Config class | Must inherit from `BaseModel` |
| Blueprint class | Must inherit from `Blueprint[ConfigClass]` |
| `render()` method | Must return `TaskGroup` or `BaseOperator` |
| Task IDs | Use `self.step_id` for the group/task ID |
| Field types | Must be single-typed and YAML-compatible (see below) |

### Config Field Types Must Be YAML-Compatible

Config fields must be single-typed. Multi-type unions like `str | int` or `Union[A, B]` are **rejected at class-definition time** (raises `TypeError`) because they produce ambiguous YAML parsing and `anyOf` schemas. The check recurses through nested models, list items, and dict values.

- **Allowed**: scalars (`str`, `int`, `float`, `bool`), `Literal[...]`, `list[X]`, `dict[str, V]`, nested `BaseModel`, and `Optional[X]` / `X | None` (the nullable pattern).
- **Rejected**: `str | int`, `Union[A, B]`, or any union with more than one non-`None` arm. Bare `Any` and `dict[str, Any]` are rejected for the same reason — use an explicit single type for the value.

### Internal Fields Not Settable from YAML

Use `Field(default=..., init=False)` for fields used inside `render()` that should not be overridable from YAML. They are excluded from the constructor and omitted from JSON Schema output:

```python
class ExtractConfig(BaseModel):
    source_table: str
    _internal_batch_multiplier: int = Field(default=4, init=False)
```

### Recommend Strict Validation for Step Configs

A **step** config model inherits Pydantic's default `extra="ignore"`, so a misspelled field in a step's YAML is silently dropped rather than reported. Suggest `model_config = ConfigDict(extra="forbid")` to turn those typos into errors:

```python
class MyConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")
    source_table: str
```

DAG args config models are the opposite and need no such setting — Blueprint makes them strict for you (see **Customizing DAG-Level Configuration**).

---

## Composing Templates

A blueprint can instantiate and render **other blueprints** inside its `render()` method, letting you build higher-level templates from lower-level building blocks while exposing a single, flat config to YAML authors.

Inside `render()`, instantiate each child blueprint, set its `step_id`, call `render(...)` with a config you construct, and wire the results together inside a parent `TaskGroup`:

```python
class QualityGate(Blueprint[QualityGateConfig]):
    """Run checks then send a report — composed from Validate and Report."""

    def render(self, config: QualityGateConfig) -> TaskGroup:
        with TaskGroup(group_id=self.step_id) as group:
            validate = Validate()
            validate.step_id = "validate"
            validate_group = validate.render(ValidateConfig(checks=config.checks))

            report = Report()
            report.step_id = "report"
            report_task = report.render(ReportConfig(channel=config.report_channel))

            validate_group >> report_task
        return group
```

YAML authors then see a single step with a flat config, and the composed children stay invisible to them.

---

## Composing DAGs in YAML

### YAML Structure

```yaml
# dags/my_pipeline.dag.yaml
dag_id: my_pipeline
schedule: "@daily"
description: "My data pipeline"

steps:
  step_one:
    blueprint: my_blueprint
    source_table: raw.customers
    batch_size: 500

  step_two:
    blueprint: another_blueprint
    depends_on: [step_one]
    target: analytics.output
```

By default, only `schedule` and `description` are supported as DAG-level fields (via the built-in `DefaultDagArgs`). For other fields like `tags`, `default_args`, `catchup`, etc., see **Customizing DAG-Level Configuration**.

### Reserved Keys in Steps

| Key | Purpose |
|-----|---------|
| `blueprint` | Template name (required) |
| `depends_on` | List of upstream step names |
| `version` | Pin to specific blueprint version |
| `trigger_rule` | Airflow trigger rule for the step; validated against the installed Airflow version |

Everything else passes to the blueprint's config.

### Trigger Rules

Use `trigger_rule` to control when a step runs relative to its upstream dependencies — for example, to run a notification step even if an upstream step failed:

```yaml
steps:
  notify:
    blueprint: notify
    depends_on: [analyze]
    trigger_rule: all_done   # run regardless of whether analyze succeeded
```

Values are validated dynamically against the installed Airflow's `TriggerRule` enum, so the accepted set follows your Airflow version rather than this skill. When the step's blueprint renders a `TaskGroup`, the rule applies only to the group's **root** tasks (those with no internal upstream), preserving the blueprint author's internal wiring.

### Jinja2 Support

YAML supports Jinja2 templating with access to environment variables, Airflow variables/connections, and runtime context:

```yaml
dag_id: "{{ env.get('ENV', 'dev') }}_pipeline"
schedule: "{{ var.value.schedule | default('@daily') }}"

steps:
  extract:
    blueprint: extract
    output_path: "/data/{{ context.ds_nodash }}/output.csv"
```

Available template variables:

- `env` — environment variables
- `var` — Airflow Variables
- `conn` — Airflow Connections
- `context` — proxy that generates Airflow template expressions for runtime macros (e.g. `context.ds_nodash`, `context.dag_run.conf`, `context.task_instance.xcom_pull(...)`)
- `profile` — the active variable profile name, or nothing when none is selected. Useful for deriving a value from the profile rather than enumerating it per profile: `dag_id: "pipeline_{{ profile }}"`

For values that are fixed at parse time and shared across steps or DAGs, prefer **Variables and Profiles** over a Jinja `{% set %}` block — variables are scoped, shareable, and visible to `blueprint lint`.

---

## Variables and Profiles

DAG YAML can declare variables and reference them as `${name}`. Use this to stop repeating a value across steps and DAGs.

Declare them in a `blueprint.vars.yaml` shared by every DAG beneath it, in a DAG's own `vars:` block, or both — nearer declarations override further ones:

```yaml
# dags/blueprint.vars.yaml — shared by every DAG beneath it
vars:
  landing_dataset: raw_events
  warehouse_db: analytics
```

```yaml
# dags/customer_etl.dag.yaml
vars:
  stream: customer_events
  retention_days: 90

steps:
  load:
    blueprint: load
    target_table: ${warehouse_db}.${landing_dataset}.${stream}
    expiration_days: ${retention_days}
```

Substitution runs after YAML parsing, so `expiration_days` stays an `int` rather than becoming the string `"90"`. Values are scalars or lists, and variables may compose (`base: ${db}.${schema}`).

Variable names match `^[A-Za-z_][A-Za-z0-9_-]*$` — hyphens are allowed, and periods are reserved so dotted namespaces can be added later without ambiguity.

> **`${...}` is always a variable reference.** Anything else that uses that syntax — most often a shell variable in a `bash_command` — must be escaped as `$${...}`, or Blueprint tries to resolve it as a variable. Only `$$` immediately before `{` is treated as an escape, so a bare `$$` (a shell PID, an awk field) needs no change. `blueprint lint` reports each unescaped occurrence and names the escape in the error, so lint the project after adopting variables.

### Profiles

A variable can carry a different value per named profile, selected at build time. Environments are the obvious use, but the mechanism is just named selection:

```yaml
profiles: [prod, dev]
vars:
  warehouse_db:
    prod: analytics
    dev: sandbox
```

```python
build_all_airflow_dags(profile="prod" if is_production else "dev")
```

Every profile a DAG declares must give the variable a value; a partial mapping is an error rather than a silent fallback.

### Inspecting Variables

`blueprint vars <path>` shows the resolved value of each variable and where it came from, and flags variables a DAG never references. `blueprint lint` validates every declared profile unless `--profile` narrows it to one. Pass `--root` to match the path the loader builds from, or resolution differs between lint and runtime.

---

## Blueprints in Python DAGs

Blueprints aren't tied to the YAML composition flow. Two patterns let you use them from Python — useful for incremental adoption or data-driven DAG generation.

### Inside a Hand-Written DAG

Instantiate the Blueprint class, set its `step_id`, call `render()`, and wire it in with `>>`:

```python
# dags/hybrid_dag.py
from datetime import datetime

from airflow import DAG
from airflow.operators.bash import BashOperator

from dags.etl_blueprints import Extract, ExtractConfig

with DAG(dag_id="hybrid_python_dag", start_date=datetime(2024, 1, 1), schedule=None, catchup=False) as dag:
    setup = BashOperator(task_id="setup", bash_command="echo 'setup'")

    extract = Extract()
    extract.step_id = "extract"
    extract_group = extract.render(ExtractConfig(source_table="raw.events"))

    setup >> extract_group
```

The `step_id` you set determines the `task_id` / `group_id` the blueprint renders under.

### Programmatic Building with `Builder` / `DAGConfig`

For data-driven DAG generation (one DAG per region, tenant, etc.), build DAGs in a loop and register each in `globals()` so Airflow discovers them:

```python
from blueprint import Builder, DAGConfig

builder = Builder()

for region in ["us", "eu", "apac"]:
    config = DAGConfig(
        dag_id=f"pipeline_{region}",
        schedule="@hourly",
        steps={
            "extract": {"blueprint": "extract", "source_table": f"raw.{region}"},
        },
    )
    dag = builder.build(config, source_path=__file__)
    globals()[dag.dag_id] = dag
```

`DAGConfig` accepts the same fields you would write in YAML. Pass `source_path=__file__` so the DAG args template is resolved from this file's directory the same way a YAML file's would be — without it, resolution falls back to the project-wide default (see **Customizing DAG-Level Configuration**).

---

## Customizing DAG-Level Configuration

By default, Blueprint supports `schedule` and `description` as DAG-level YAML fields. To use other DAG constructor arguments (tags, default_args, catchup, etc.), define a `BlueprintDagArgs` subclass. Its `render()` returns a dict of kwargs passed to the Airflow `DAG()` constructor, so the accepted keys are whatever your Airflow version's `DAG` accepts.

```python
# dags/dag_args.py
from typing import Any

from pydantic import BaseModel
from blueprint import BlueprintDagArgs

class ProjectDagArgsConfig(BaseModel):
    schedule: str | None = None
    tags: list[str] = []
    owner: str = "data-team"

class ProjectDagArgs(BlueprintDagArgs[ProjectDagArgsConfig]):
    def render(self, config: ProjectDagArgsConfig) -> dict[str, Any]:
        return {
            "schedule": config.schedule,
            "tags": config.tags,
            "default_args": {"owner": config.owner},
        }
```

The declared fields then become valid DAG-level YAML keys, validated by the config model.

### Several Templates per Project

A project may define more than one template. Each DAG uses the template defined **closest above it**: resolution starts in the DAG file's own directory and walks up parent directories, so a subdirectory overrides its parents.

```
dags/
  dag_args.py             ProjectDagArgs
  customer.dag.yaml       -> ProjectDagArgs
  sandbox/
    dag_args.py           SandboxDagArgs
    probe.dag.yaml        -> SandboxDagArgs
```

A DAG with no template above it falls back to the one declared `default=True`, then to the sole registered template, then to the built-in `DefaultDagArgs`. A template is scoped to the directory holding the `.py` file that defines it, so **moving that file rescopes it** — the most common surprise in this feature.

Nothing in the DAG YAML changes: a DAG never names its template. Run `blueprint list` to see which template applies to which path, which one is the fallback, and where each is defined; `blueprint lint` names the resolved template per DAG.

A template registers under the snake_case form of its class name — `ProjectDagArgs` becomes `project_dag_args` — which is the name `blueprint schema --dag-args <name>` expects. Setting `name = "..."` overrides it, and must itself be snake_case.

### Undeclared Fields Are Rejected

A DAG args config model defines the DAG YAML's top-level surface, so Blueprint applies `extra="forbid"` to it automatically — an undeclared top-level key is an error rather than a silently ignored one. **This is the opposite default from step configs**, which ignore unknown keys unless you opt in.

This shows up in the generated schema as `additionalProperties: false`, so editors and the Astro IDE reject unknown top-level keys too.

Setting `extra` yourself on the model leaves your choice intact. To defer to the model's own policy instead, pass `allow_extra=True` on the class:

```python
class LooseDagArgs(BlueprintDagArgs[LooseConfig], allow_extra=True):
    ...
```

### Rules

- Two templates in the same directory is an error, as is two sharing a name (`name = "..."` on the class renames one) or more than one declaring `default=True`.
- If no subclass exists anywhere, the built-in `DefaultDagArgs` is used (`schedule` and `description` only).

---

## Sharing Blueprints as a Package

Blueprints can be shared across repositories as an installable package instead of copied files. The package advertises itself under the `airflow_blueprint.blueprints` entry-point group, and Blueprint discovers it once installed, with no per-repo configuration:

```toml
# pyproject.toml of the shared package
[project.entry-points."airflow_blueprint.blueprints"]
company_blueprints = "company_blueprints"
```

The value must be a plain dotted module or package path. The advertised module — and every submodule, if it is a package — is scanned exactly like a locally discovered file, so any `Blueprint` or `BlueprintDagArgs` subclass defined in it is registered.

Consumers install the package and the blueprints appear in `blueprint list` alongside local ones, with the source column distinguishing them. A package that fails to import raises `EntryPointLoadError` rather than silently vanishing from the registry.

To turn discovery off, pass `discover_entry_points=False` to the loader, or the corresponding `--no-entry-points` flag to the CLI (`--help` confirms which commands accept it).

---

## Runtime Parameter Overrides

Blueprint config fields can be overridden at DAG trigger time using Airflow params, letting users customize behavior when manually triggering DAGs.

### Opt In with `supports_params = True`

A blueprint must set the class attribute `supports_params = True` for its config fields to register as Airflow params (namespaced as `{step}__{field}`). **Without it, `self.param()` / `self.resolve_config()` do nothing and no fields appear in the trigger form.** Only opt in for blueprints that actually use those methods — otherwise dead params clutter the form with no effect.

### Canonical Example

Use `self.param()` in operator template fields, where Airflow renders the value at execution time; use `self.resolve_config()` in Python callables, where you need a validated config object. Both can appear in one blueprint:

```python
class Extract(Blueprint[ExtractConfig]):
    supports_params = True

    def render(self, config: ExtractConfig) -> TaskGroup:
        bp = self  # capture reference for the closure

        @task(task_id="run_query")
        def run_query(**context):
            resolved = bp.resolve_config(config, context)
            execute(resolved.query, resolved.batch_size)

        with TaskGroup(group_id=self.step_id) as group:
            BashOperator(
                task_id="shell_step",
                bash_command=f"run-etl --query {self.param('query')}",
            )
            run_query()
        return group
```

### How It Works

- Params are **auto-generated** from Pydantic config models and namespaced per step (e.g. `step_name__field`)
- YAML values become param defaults; Pydantic metadata (description, constraints, enum values) flows through to the Airflow trigger form
- Invalid overrides raise `ValidationError` at execution time
- Override them from the trigger form, or by posting `conf` with the namespaced names to the DAG run endpoint (`af api ls --filter dagRun` finds the current path — see the `airflow` skill)

### Trigger Form Customization

Pydantic field schema flows through to Airflow's trigger form; `json_schema_extra` controls how each field renders (`format` values such as multiline and date pickers, `examples`, `values_display`, `description_md`). The Airflow version determines which are honoured, so check against the form rather than assuming.

**Validation nuance:** only `Field` constraints that map to JSON Schema (`ge`, `le`, `pattern`, `min_length`, `max_length`, `Literal` enums) are enforced in the trigger form. Custom `@field_validator` / `@model_validator` logic does **not** map to JSON Schema, so it runs only at build time and inside `resolve_config()`. If custom validators enforce important constraints, call `self.resolve_config()` in your `@task` function so they run on overridden values.

---

## Loader Options

`build_all_airflow_dags()` takes the options that govern a whole project. The ones that change behaviour materially:

| Option | Effect |
|---|---|
| `profile=` | Selects which profile's values the `${...}` variables resolve to (see **Variables and Profiles**) |
| `skip_invalid_dags=True` | Renders the valid YAML files and skips faulty ones instead of failing the import |
| `discover_entry_points=False` | Turns off discovery of blueprints installed as packages |
| `on_dag_built=` | Callback to post-process each DAG after construction |

`discover_entry_points` is ignored when `bp_registry` is supplied directly, since that registry has already run discovery.

The rest of the signature is plumbing that rarely needs changing: `search_path` and `pattern` control YAML discovery, `register_globals` overrides the caller's `globals()`, `render_templates` and `template_context` govern Jinja, and `bp_registry` supplies a pre-built registry.

### Excluding Files with `.airflowignore`

YAML discovery honours `.airflowignore`, using Airflow's own ignore-file walker — so the syntax, the `core.dag_ignore_file_syntax` setting, and nested ignore files behave exactly as they do for the DAG processor. `blueprint lint` honours it too when scanning a directory, so a draft excluded from Airflow is also excluded from lint; passing that file explicitly still lints it, which is how you check a draft on purpose.

Patterns are matched against the tail of each path, so name patterns like `*.dag.yaml` behave as with `rglob`. `**` is not supported before Python 3.13.

### Post-Build Callbacks

Use `on_dag_built` to post-process DAGs after construction — adding tags, access controls, or audit metadata:

```python
from pathlib import Path
from blueprint import build_all_airflow_dags

def add_audit_tags(dag, yaml_path: Path) -> None:
    dag.tags.append("managed-by-blueprint")
    dag.tags.append(f"source:{yaml_path.name}")

build_all_airflow_dags(on_dag_built=add_audit_tags)
```

The callback receives the constructed Airflow `DAG` (mutable) and the `Path` of the YAML file that defined it.

### Skipping Invalid Files

`skip_invalid_dags=True` stops one bad YAML file from taking down every other DAG in the folder. Explain both costs before recommending it:

- Errors no longer surface as Airflow import errors, because the loader itself parses cleanly. They go to the DAG processor log instead, which is a much less visible place to look.
- Duplicate DAG ids stop being an error — every such file parses, and the first DAG wins.

Pair it with `blueprint lint` in CI, so invalid files are caught somewhere visible.

---

## Validation Commands

Run CLI commands with uvx:

```bash
uvx --from airflow-blueprint blueprint <command>
```

| Command | When to Use |
|---------|-------------|
| `list` | Show available blueprints, versions, sources, and DAG args templates |
| `describe <name>` | Show config schema for a blueprint |
| `lint` | Validate DAG YAML — bare to scan recursively, or pass one file |
| `vars <path>` | Show resolved variables for a DAG and where each came from |
| `schema` | Generate JSON Schema for a blueprint or for DAG-level fields |
| `new` | Interactive DAG YAML creation. `--output-dir` picks where the file lands, which also selects the DAG args template it is validated against |

Every command takes `--help`, and `-h` / `-v` work as shorthands for `--help` / `--version`.

Run them from the **project root**, not from inside `dags/` — a bare invocation resolves `dags/` relative to the working directory, so running from within it finds no blueprints. Use `--template-dir` for any other layout.

> **Provider operators in the CLI.** The `uvx --from airflow-blueprint` environment is isolated and does **not** include the Airflow provider packages your Astro Runtime project has. If templates import provider operators, add `--with <provider-package>` so the CLI can import them — otherwise `list`/`lint`/`schema` fail with `ModuleNotFoundError`:
>
> ```bash
> uvx --from airflow-blueprint --with apache-airflow-providers-google blueprint list --template-dir dags/templates
> ```

---

## Versioning

### Version Naming Convention

Versions are separate classes with a `V{N}` suffix: `Extract` is v1, `ExtractV2` is v2, and each carries its own config model. A blueprint's discovered versions must form a contiguous `1..N` sequence.

```python
class Extract(Blueprint[ExtractConfig]):        # v1
    def render(self, config): ...

class ExtractV2(Blueprint[ExtractV2Config]):    # v2, breaking changes
    def render(self, config): ...
```

### Explicit Name and Version

When the class name doesn't follow the convention, set them directly:

```python
class MyCustomExtractor(Blueprint[ExtractV3Config]):
    name = "extract"
    version = 3

    def render(self, config): ...
```

An explicit `name` must be snake_case (`^[a-z][a-z0-9_]*$`) or the class raises `ValueError` at definition time. Without one, the name is the snake_case form of the class name.

### Using Versions in YAML

Omit `version` to get the latest; pin it to hold a step on an older one:

```yaml
steps:
  legacy_extract:
    blueprint: extract
    version: 1
    source_table: raw.data
```

`blueprint list` shows the discovered versions of each blueprint.

---

## Schema Generation

Generate JSON schemas for editor autocompletion or external tooling. `blueprint schema <name>` emits a step template's config; `blueprint schema --dag-args` emits the DAG-level fields (`dag_id`, `steps`, and whatever your `BlueprintDagArgs` exposes). With multiple DAG args templates, `--dag-args` takes an optional template name.

Each emitted schema includes a top-level `templateType` field — `"blueprint"` for a step template, `"dag_args"` for DAG-level fields — so consumers can tell them apart. The command emits raw JSON when piped or written with `-o/--output`, and pretty, highlighted JSON when run interactively.

> **Write with `-o/--output`, not `>`.** Importing a template can print warnings to stdout — an Airflow deprecation warning from an operator import is the common case — and those interleave with the JSON, leaving redirected output unparseable. `-o` writes the schema alone.

> **Optional fields emit a plain type.** An optional config field is published as `{"type": "string"}`, not an `anyOf` with a null branch — optionality is carried by the schema's `required` array alone. This keeps generated clients and form renderers from producing a union wrapper type for every optional field. Airflow params deliberately differ and keep a nullable type, because an unset optional param is an explicit null rather than an absent key; do not "fix" one to match the other.

### Astro Project Auto-Detection

After creating or modifying a blueprint, **automatically check** whether the project is an Astro project by looking for a `.astro/` directory (created by `astro dev init`).

If it is, **automatically regenerate schemas** without prompting, writing one file per blueprint from `blueprint list` plus the DAG-level args schema, into `blueprint/generated-schemas/`. The Astro IDE reads that directory to render configuration forms, so keeping it in sync ensures the visual builder reflects the latest configs.

If you cannot determine whether the project is an Astro project, ask the user once and remember for the rest of the session.

---

## Troubleshooting

Error messages carry their own remediation hints; read the message before applying anything here.

### "Blueprint not found"

**Cause**: Blueprint class not in Python path.

**Fix**: Point the CLI at the right directory with `--template-dir`, and check `blueprint list` for what is actually discovered. If the blueprint is meant to come from an installed package, confirm entry-point discovery is on.

### "Extra inputs are not permitted"

**Cause**: YAML field name typo with `extra="forbid"` enabled.

**Fix**: Run `blueprint describe <name>` to see valid field names.

### DAG not appearing in Airflow

**Cause**: Missing or broken loader — including a loader that imports a deprecated alias, which Airflow safe-mode may skip.

**Fix**: Ensure `dags/loader.py` calls `build_all_airflow_dags()`. If `skip_invalid_dags=True` is set, the file parses even when a DAG is broken, so check the DAG processor log rather than the import errors view.

### "ModuleNotFoundError: No module named 'airflow.providers.X'" from the CLI

**Cause**: The standalone `uvx --from airflow-blueprint` environment doesn't include the Airflow provider packages your project has, so a template importing provider operators can't be imported. This is the CLI's isolated environment, not your project.

**Fix**: Add `--with apache-airflow-providers-X` to the uvx invocation.

### Unresolved or unexpected `${...}`

**Cause**: A `${...}` that is not a declared variable — commonly a shell variable in a `bash_command`, or a variable declared in a `blueprint.vars.yaml` outside the search root.

**Fix**: Escape non-variable occurrences as `$${...}`. For genuinely missing variables, run `blueprint vars <path>` to see what resolves and `blueprint lint` for the full list; check that `--root` matches the path the loader uses.

### "CyclicVariableError" / "CompositionDepthError"

**Cause**: Variables that reference each other in a loop, or a composition chain deeper than the resolver's limit.

**Fix**: The error names the cycle or the chain. Break it by inlining one value; `blueprint vars <path>` shows what each variable resolves to.

### "MultipleDagArgsError" / "DuplicateDagArgsError" / "MultipleDefaultDagArgsError"

**Cause**: Not that several templates exist — that is supported. These fire when resolution is ambiguous: two templates in one directory, two sharing a name, or more than one declaring `default=True`.

**Fix**: Move one template to the directory whose DAGs should use it, rename one with `name = "..."`, or leave only one `default=True`. `blueprint list` shows which template applies where.

### "DagArgsNotFoundError"

**Cause**: A named DAG args template was requested that isn't registered.

**Fix**: Check the name against `blueprint list`.

### "EntryPointLoadError"

**Cause**: An installed package advertising blueprints failed to import.

**Fix**: Import the module directly to see the real traceback, and confirm the package and its dependencies are installed in the same environment as Airflow.

### "NonContiguousVersionError" / "InvalidVersionError"

**Cause**: A blueprint's versions don't form a contiguous `1..N` sequence, or YAML pins a version that doesn't exist.

**Fix**: Ensure versions increment by one with no gaps; run `blueprint list` to see available versions.

### "non-YAML-compatible fields" (TypeError at import)

**Cause**: A config field uses a type Blueprint rejects — a multi-type union (e.g. `str | int`), bare `Any`, or `dict[str, Any]`.

**Fix**: Use a single, explicit type. `Optional[X]` / `X | None` is still allowed. See **Creating Blueprints → Config Field Types Must Be YAML-Compatible**.

### "Cyclic dependency detected"

**Cause**: Circular `depends_on` references.

**Fix**: Review step dependencies and remove cycles.

### Debugging in Airflow UI

Every Blueprint task has extra fields in **Rendered Template**:

- `blueprint_step_config` — resolved YAML config
- `blueprint_step_code` — Python source of the blueprint

---

## Verification Checklist

Before finishing, verify with the user:

- [ ] `blueprint list` shows their templates, and the expected DAG args template applies to each path
- [ ] `blueprint lint` passes (bare to scan all `*.dag.yaml` recursively, or pass a specific file — passing a directory path fails with `Is a directory`)
- [ ] `blueprint vars` resolves as expected, if variables or profiles are in use
- [ ] `dags/loader.py` exists and calls `build_all_airflow_dags()`
- [ ] DAG appears in Airflow UI without parse errors

---

## Reference

- GitHub: https://github.com/astronomer/blueprint
- PyPI: https://pypi.org/project/airflow-blueprint/

### Astro IDE

- Astro IDE Blueprint docs: https://docs.astronomer.io/astro/ide-blueprint
checking-freshness3.13 KB

View saved version →

---
name: checking-freshness
description: Quick data freshness check. Use when the user asks if data is up to date, when a table was last updated, if data is stale, or needs to verify data currency before using it.
---

# Data Freshness Check

Quickly determine if data is fresh enough to use.

## Freshness Check Process

For each table to check:

### 1. Find the Timestamp Column

Look for columns that indicate when data was loaded or updated:
- `_loaded_at`, `_updated_at`, `_created_at` (common ETL patterns)
- `updated_at`, `created_at`, `modified_at` (application timestamps)
- `load_date`, `etl_timestamp`, `ingestion_time`
- `date`, `event_date`, `transaction_date` (business dates)

Query INFORMATION_SCHEMA.COLUMNS if you need to see column names.

### 2. Query Last Update Time

```sql
SELECT
    MAX(<timestamp_column>) as last_update,
    CURRENT_TIMESTAMP() as current_time,
    TIMESTAMPDIFF('hour', MAX(<timestamp_column>), CURRENT_TIMESTAMP()) as hours_ago,
    TIMESTAMPDIFF('minute', MAX(<timestamp_column>), CURRENT_TIMESTAMP()) as minutes_ago
FROM <table>
```

### 3. Check Row Counts by Time

For tables with regular updates, check recent activity:

```sql
SELECT
    DATE_TRUNC('day', <timestamp_column>) as day,
    COUNT(*) as row_count
FROM <table>
WHERE <timestamp_column> >= DATEADD('day', -7, CURRENT_DATE())
GROUP BY 1
ORDER BY 1 DESC
```

## Freshness Status

Report status using this scale:

| Status | Age | Meaning |
|--------|-----|---------|
| **Fresh** | < 4 hours | Data is current |
| **Stale** | 4-24 hours | May be outdated, check if expected |
| **Very Stale** | > 24 hours | Likely a problem unless batch job |
| **Unknown** | No timestamp | Can't determine freshness |

## If Data is Stale

Check Airflow for the source pipeline:

1. **Find the DAG**: Which DAG populates this table? Use `af dags list` and look for matching names.

2. **Check DAG status**:
   - Is the DAG paused? Use `af dags get <dag_id>`
   - Did the last run fail? Use `af dags stats`
   - Is a run currently in progress?

3. **Diagnose if needed**: If the DAG failed, use the **debugging-dags** skill to investigate.

### On Astro

If you're running on Astro, you can also:

- **DAG history in the Astro UI**: Check the deployment's DAG run history for a visual timeline of recent runs and their outcomes
- **Astro alerts for SLA monitoring**: Configure alerts to get notified when DAGs miss their expected completion windows, catching staleness before users report it

### On OSS Airflow

- **Airflow UI**: Use the DAGs view and task logs to verify last successful runs and SLA misses

## Output Format

Provide a clear, scannable report:

```
FRESHNESS REPORT
================

TABLE: database.schema.table_name
Last Update: 2024-01-15 14:32:00 UTC
Age: 2 hours 15 minutes
Status: Fresh

TABLE: database.schema.other_table
Last Update: 2024-01-14 03:00:00 UTC
Age: 37 hours
Status: Very Stale
Source DAG: daily_etl_pipeline (FAILED)
Action: Investigate with **debugging-dags** skill
```

## Quick Checks

If user just wants a yes/no answer:
- "Is X fresh?" -> Check and respond with status + one line
- "Can I use X for my 9am meeting?" -> Check and give clear yes/no with context
configuring-airflow-language-sdks9.19 KB

View saved version →

---
name: configuring-airflow-language-sdks
description: Configures Airflow to run language SDK tasks (Java, Go, and future native SDKs) — register a coordinator, map a queue to it, ensure the runtime/artifact on workers, and tune coordinator options. Use when the user wants Airflow to route a queue to a native-language coordinator, asks about the `[sdk]` `coordinators`/`queue_to_coordinator` settings, `AIRFLOW__SDK__COORDINATORS`, `jars_root`, `executables_root` or other coordinator `kwargs`, `task_startup_timeout`, or why their native tasks aren't being picked up. Covers the shared routing mechanism plus per-coordinator options (e.g. JavaCoordinator, ExecutableCoordinator).
---

# Configuring Airflow for Language SDKs

To run language SDK tasks, Airflow needs to know two things: which **coordinator** launches the native subprocess, and which **queue** routes to that coordinator. The mechanism is identical across every language SDK — only each coordinator's `classpath` and `kwargs` differ. This skill documents the shared wiring once, then the per-coordinator options. It is platform-neutral: the same settings apply on open-source Airflow and on managed platforms like Astro.

> **Experimental.** The language SDKs are in preview; configuration keys may change.

> For the task code, see **authoring-language-sdk-tasks** (and the per-language authoring skill, e.g. **authoring-java-sdk-tasks**, **authoring-go-sdk-tasks**). For building and shipping the artifact, see the per-language deploy skill (e.g. **deploying-java-sdk-bundles**, **deploying-go-sdk-bundles**).

---

## Prerequisites on the worker

- The **runtime or artifact the SDK needs** must be present on the worker nodes, because the coordinator spawns a native subprocess per task instance. The exact requirement is per-SDK — see [Per-coordinator options](#per-coordinator-options) (the Java SDK needs a **JRE 17+**; the Go SDK needs no language runtime — the bundle is a self-contained native executable, but it must be built for the worker's OS/arch).
- The compiled/native **artifact(s)** must be reachable on the worker. See the per-language deploy skill.
- The coordinators ship with the Airflow Task SDK (`apache-airflow-task-sdk`, installed with Airflow). **No extra Python package is required.**

---

## The two settings

Both live in the `[sdk]` configuration section and apply to every language SDK:

1. **`coordinators`** — a JSON object mapping a coordinator *name* you choose to its implementation (`classpath`) and constructor `kwargs`.
2. **`queue_to_coordinator`** — a JSON object mapping a task *queue* to a coordinator name.

A task whose stub sets `queue="..."` is handed to the named coordinator, which launches the native subprocess. The coordinator name is arbitrary — it just has to be the same string in both settings. The queue name must match the `queue=` set on the Python `@task.stub`.

### Option A: `airflow.cfg`

```ini
[sdk]
coordinators = {
  "java-jdk17": {
    "classpath": "airflow.sdk.coordinators.java.JavaCoordinator",
    "kwargs": {"jars_root": ["/opt/airflow/jars"]}
  },
  "go": {
    "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator",
    "kwargs": {"executables_root": ["/opt/airflow/executable-bundles"]}
  }
}
queue_to_coordinator = {"java": "java-jdk17", "golang": "go"}
```

### Option B: environment variables

Each value must be **valid one-line JSON**. This form is convenient for containers, `.env` files, Docker Compose, and Helm.

```bash
export AIRFLOW__SDK__COORDINATORS='{"java-jdk17": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": {"jars_root": ["/opt/airflow/jars"]}}, "go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"executables_root": ["/opt/airflow/executable-bundles"]}}}'
export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"java": "java-jdk17", "golang": "go"}'
```

The examples above register **multiple coordinators** at once (one per language) and map a different queue to each — register only the ones you use.

---

## Per-coordinator options

The `classpath` and `kwargs` are specific to each coordinator. Add a subsection here as new language SDKs land.

### JavaCoordinator

- **`classpath`**: `airflow.sdk.coordinators.java.JavaCoordinator`
- **Worker runtime**: JRE 17+ (`java` on `PATH`, or set `java_executable`).

| Parameter | Default | Description |
|-----------|---------|-------------|
| `jars_root` | *(required)* | One or more directories scanned **recursively** for `.jar` files. Accepts a string or a list of strings/paths. The classpath is assembled automatically. |
| `java_executable` | `"java"` | Path to the `java` binary. Defaults to `java` on `$PATH`. |
| `jvm_args` | `[]` | Extra JVM arguments, e.g. `["-Xmx1g", "-Dsome.property=value"]`. |
| `main_class` | *(auto-detect)* | Explicit entry-point class. If omitted, the coordinator scans `jars_root` for a JAR whose manifest declares `Main-Class`. **Set this explicitly if multiple executable JARs are present** — otherwise the choice is non-deterministic. |
| `task_startup_timeout` | `10.0` | Seconds to wait for the subprocess to connect after launch. Increase it if JVM startup is slow (constrained hardware, large classpath, first cold start). |

**Java logging via `java.util.logging`.** Of the SDK logging integrations, only JPL and SLF4J are zero-config build dependencies; Log4j 2 and JUL need extra setup — see the logging integration section in **deploying-java-sdk-bundles**. JUL's documented alternative to calling `AirflowJulHandler.setup()` in `main()` is a `logging.properties` file, wired through `jvm_args`:

```ini
[sdk]
coordinators = {
  "java-jdk17": {
    "classpath": "airflow.sdk.coordinators.java.JavaCoordinator",
    "kwargs": {
      "jars_root": ["/opt/airflow/jars"],
      "jvm_args": ["-Djava.util.logging.config.file=/opt/airflow/logging.properties"]
    }
  }
}
```

### ExecutableCoordinator (Go and other self-contained-executable SDKs)

- **`classpath`**: `airflow.sdk.coordinators.executable.ExecutableCoordinator`
- **Worker runtime**: none beyond the bundle itself. The bundle is a self-contained native executable (AFBNDL01), so it needs no language runtime, but it must be built for the worker's OS/arch (a mismatch fails with `exec format error`).

| Parameter | Default | Description |
|-----------|---------|-------------|
| `executables_root` | *(required)* | One or more directories scanned **recursively** for executable bundles (AFBNDL01-trailered native binaries). Accepts a string or a list of strings/paths. Bundles are identified by the trailer magic, not by filename. The coordinator matches an incoming `dag_id` against each bundle's embedded manifest and verifies its integrity hash before launching. |
| `task_startup_timeout` | `10.0` | Seconds to wait for the subprocess to connect after launch. Increase it if bundle startup is slow (constrained hardware, first cold start). |

*(Future coordinators — for other languages — will list their own `classpath`, runtime, and `kwargs` here.)*

---

## Verifying the configuration

1. Confirm the runtime/artifact is usable where workers run — for the Java SDK, `java -version` via `astro dev bash` or `docker compose exec ...`; for the Go SDK, the packed bundle exists and matches the worker's OS/arch.
2. Confirm the artifact directory referenced in `kwargs` (e.g. `jars_root`, `executables_root`) actually contains your artifact on the worker filesystem.
3. Trigger the DAG and open the native task's logs — you should see the subprocess start and your task output.

### Troubleshooting

| Symptom | Likely cause / fix |
|---------|--------------------|
| Task fails immediately mentioning coordinator or queue | `coordinators` / `queue_to_coordinator` not valid one-line JSON, or the queue name doesn't match the stub's `queue=`. Fix the JSON and restart. |
| Runtime not found (e.g. `java: command not found`) | The language runtime isn't on the worker, or the executable path kwarg is wrong. Install the runtime and verify its version. |
| "No artifact found" / "no DAGs" / "no bundle contains dag_id" | The artifact-directory kwarg points at the wrong place, the artifact isn't there yet, or its `dag_id` doesn't match the stub. Confirm the path and the IDs. |
| Wrong/ambiguous entry point (Java) | Multiple executable JARs under `jars_root`. Set `main_class` explicitly. |
| Go bundle is skipped silently | Not a valid AFBNDL01 bundle, or its integrity hash failed (re-pack after any strip/sign/rebuild). |
| `exec format error` on the Go bundle | Built for a different OS/arch than the worker. Cross-compile with `--goos`/`--goarch` (see **deploying-go-sdk-bundles**). |
| DAG run hangs at the native task | Raise `task_startup_timeout` (e.g. `30.0`); first-run subprocess startup can be slow. |

---

## Related Skills

- **authoring-language-sdk-tasks**: The shared Python-stub pattern and conceptual model.
- **authoring-java-sdk-tasks**: Java task code and matching Python stubs.
- **deploying-java-sdk-bundles**: Build the bundle and put the artifact where the coordinator scans.
- **authoring-go-sdk-tasks**: Go task code and matching Python stubs.
- **deploying-go-sdk-bundles**: Build/pack the Go bundle and place it where the coordinator scans.
- **deploying-airflow**: General deployment of Airflow on Astro, Docker Compose, or Kubernetes.
cosmos-dbt-core14.3 KB

View saved version →

---
name: cosmos-dbt-core
description: Turns a dbt Core project into an Airflow DAG/TaskGroup using Astronomer Cosmos. Use turning a dbt Core project into an Airflow DAG or TaskGroup with Astronomer Cosmos. Before implementing, verify dbt engine, warehouse, Airflow version, execution environment, DAG vs TaskGroup, and manifest availability.
---

# Cosmos + dbt Core: Implementation Checklist

Execute steps in order. Prefer the simplest configuration that meets the user's constraints.

> **Version note**: This skill targets Cosmos 1.11+ and Airflow 3.x. If the user is on Airflow 2.x, adjust imports accordingly (see Appendix A).
>
> **Reference**: Latest stable: https://pypi.org/project/astronomer-cosmos/

> **Before starting**, confirm: (1) dbt engine = Core (not Fusion → use **cosmos-dbt-fusion**), (2) warehouse type, (3) Airflow version, (4) execution environment (Airflow env / venv / container), (5) DbtDag vs DbtTaskGroup vs individual operators, (6) manifest availability.

---

## 1. Configure Project (ProjectConfig)

| Approach | When to use | Required param |
|----------|-------------|----------------|
| Project path | Files available locally | `dbt_project_path` |
| Manifest only | `dbt_manifest` load | `manifest_path` + `project_name` |

```python
from cosmos import ProjectConfig

_project_config = ProjectConfig(
    dbt_project_path="/path/to/dbt/project",
    # manifest_path="/path/to/manifest.json",  # for dbt_manifest load mode
    # project_name="my_project",  # if using manifest_path without dbt_project_path
    # install_dbt_deps=False,  # if deps precomputed in CI
)
```

## 2. Choose Parsing Strategy (RenderConfig)

Pick ONE load mode based on constraints:

| Load mode | When to use | Required inputs | Constraints |
|-----------|-------------|-----------------|-------------|
| `dbt_manifest` | Large projects; containerized execution; fastest | `ProjectConfig.manifest_path` | Remote manifest needs `manifest_conn_id` |
| `dbt_ls` | Complex selectors; need dbt-native selection | dbt installed OR `dbt_executable_path` | Can also be used with containerized execution |
| `dbt_ls_file` | dbt_ls selection without running dbt_ls every parse | `RenderConfig.dbt_ls_path` | `select`/`exclude` won't work |
| `automatic` (default) | Simple setups; let Cosmos pick | (none) | Falls back: manifest → dbt_ls → custom |

> **CRITICAL**: Containerized execution (`DOCKER`/`KUBERNETES`/etc.)

```python
from cosmos import RenderConfig, LoadMode

_render_config = RenderConfig(
    load_method=LoadMode.DBT_MANIFEST,  # or DBT_LS, DBT_LS_FILE, AUTOMATIC
)
```

---

## 3. Choose Execution Mode (ExecutionConfig)

> **Reference**: See **[reference/cosmos-config.md](reference/cosmos-config.md#execution-modes-executionconfig)** for detailed configuration examples per mode.

Pick ONE execution mode:

| Execution mode | When to use | Speed | Required setup |
|----------------|-------------|-------|----------------|
| `WATCHER` | Fastest; single `dbt build` visibility | Fastest | dbt adapter in env OR `dbt_executable_path` or dbt Fusion |
| `WATCHER_KUBERNETES` | Fastest isolated method; single `dbt build` visibility | Fast | dbt installed in container |
| `LOCAL` + `DBT_RUNNER` | dbt + adapter in the same Python installation as Airflow | Fast | dbt 1.5+ in `requirements.txt` |
| `LOCAL` + `SUBPROCESS` | dbt + adapter available in the Airflow deployment, in an isolated Python installation | Medium | `dbt_executable_path` |
| `AIRFLOW_ASYNC` | BigQuery + long-running transforms | Fast | Airflow ≥2.8; provider deps |
| `KUBERNETES` | Isolation between Airflow and dbt | Medium | Airflow ≥2.8; provider deps |
| `VIRTUALENV` | Can't modify image; runtime venv | Slower | `py_requirements` in operator_args |
| Other containerized approaches | Support Airflow and dbt isolation | Medium | container config |

```python
from cosmos import ExecutionConfig, ExecutionMode

_execution_config = ExecutionConfig(
    execution_mode=ExecutionMode.WATCHER,  # or LOCAL, VIRTUALENV, AIRFLOW_ASYNC, KUBERNETES, etc.
)
```

---

## 4. Configure Warehouse Connection (ProfileConfig)

> **Reference**: See **[reference/cosmos-config.md](reference/cosmos-config.md#profileconfig-warehouse-connection)** for detailed ProfileConfig options and all ProfileMapping classes.

### Option A: Airflow Connection + ProfileMapping (Recommended)

```python
from cosmos import ProfileConfig
from cosmos.profiles import SnowflakeUserPasswordProfileMapping

_profile_config = ProfileConfig(
    profile_name="default",
    target_name="dev",
    profile_mapping=SnowflakeUserPasswordProfileMapping(
        conn_id="snowflake_default",
        profile_args={"schema": "my_schema"},
    ),
)
```

### Option B: Existing profiles.yml

> **CRITICAL**: Do not hardcode secrets; use environment variables.

```python
from cosmos import ProfileConfig

_profile_config = ProfileConfig(
    profile_name="my_profile",
    target_name="dev",
    profiles_yml_filepath="/path/to/profiles.yml",
)
```

---

## 5. Configure Testing Behavior (RenderConfig)

> **Reference**: See **[reference/cosmos-config.md](reference/cosmos-config.md#testing-behavior-renderconfig)** for detailed testing options.

| TestBehavior | Behavior |
|--------------|----------|
| `AFTER_EACH` (default) | Tests run immediately after each model (default) |
| `BUILD` | Combine run + test into single `dbt build` |
| `AFTER_ALL` | All tests after all models complete |
| `NONE` | Skip tests |

```python
from cosmos import RenderConfig, TestBehavior

_render_config = RenderConfig(
    test_behavior=TestBehavior.AFTER_EACH,
)
```

---

## 6. Configure operator_args

> **Reference**: See **[reference/cosmos-config.md](reference/cosmos-config.md#operator_args-configuration)** for detailed operator_args options.

```python
_operator_args = {
    # BaseOperator params
    "retries": 3,

    # Cosmos-specific params
    "install_deps": False,
    "full_refresh": False,
    "quiet": True,

    # Runtime dbt vars (XCom / params)
    "vars": '{"my_var": "{{ ti.xcom_pull(task_ids=\'pre_dbt\') }}"}',
}
```

---

## 7. Assemble DAG / TaskGroup

### Option A: DbtDag (Standalone)

```python
from cosmos import DbtDag, ProjectConfig, ProfileConfig, ExecutionConfig, RenderConfig
from cosmos.profiles import SnowflakeUserPasswordProfileMapping
from pendulum import datetime

_project_config = ProjectConfig(
    dbt_project_path="/usr/local/airflow/dbt/my_project",
)

_profile_config = ProfileConfig(
    profile_name="default",
    target_name="dev",
    profile_mapping=SnowflakeUserPasswordProfileMapping(
        conn_id="snowflake_default",
    ),
)

_execution_config = ExecutionConfig()
_render_config = RenderConfig()

my_cosmos_dag = DbtDag(
    dag_id="my_cosmos_dag",
    project_config=_project_config,
    profile_config=_profile_config,
    execution_config=_execution_config,
    render_config=_render_config,
    operator_args={},
    start_date=datetime(2025, 1, 1),
    schedule="@daily",
)
```

### Option B: DbtTaskGroup (Inside Existing DAG)

```python
from airflow.sdk import dag, task  # Airflow 3.x
# from airflow.decorators import dag, task  # Airflow 2.x
from airflow.models.baseoperator import chain
from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig, ExecutionConfig, RenderConfig
from pendulum import datetime

_project_config = ProjectConfig(dbt_project_path="/usr/local/airflow/dbt/my_project")
_profile_config = ProfileConfig(profile_name="default", target_name="dev")
_execution_config = ExecutionConfig()
_render_config = RenderConfig()

@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
def my_dag():
    @task
    def pre_dbt():
        return "some_value"

    dbt = DbtTaskGroup(
        group_id="dbt_project",
        project_config=_project_config,
        profile_config=_profile_config,
        execution_config=_execution_config,
        render_config=_render_config,
    )

    @task
    def post_dbt():
        pass

    chain(pre_dbt(), dbt, post_dbt())

my_dag()
```

### Option C: Use Cosmos operators directly

```python
import os
from datetime import datetime
from pathlib import Path
from typing import Any

from airflow import DAG

try:
    from airflow.providers.standard.operators.python import PythonOperator
except ImportError:
    from airflow.operators.python import PythonOperator

from cosmos import DbtCloneLocalOperator, DbtRunLocalOperator, DbtSeedLocalOperator, ProfileConfig
from cosmos.io import upload_to_aws_s3

DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt"
DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH))
DBT_PROJ_DIR = DBT_ROOT_PATH / "jaffle_shop"
DBT_PROFILE_PATH = DBT_PROJ_DIR / "profiles.yml"
DBT_ARTIFACT = DBT_PROJ_DIR / "target"

profile_config = ProfileConfig(
    profile_name="default",
    target_name="dev",
    profiles_yml_filepath=DBT_PROFILE_PATH,
)


def check_s3_file(bucket_name: str, file_key: str, aws_conn_id: str = "aws_default", **context: Any) -> bool:
    """Check if a file exists in the given S3 bucket."""
    from airflow.providers.amazon.aws.hooks.s3 import S3Hook

    s3_key = f"{context['dag'].dag_id}/{context['run_id']}/seed/0/{file_key}"
    print(f"Checking if file {s3_key} exists in S3 bucket...")
    hook = S3Hook(aws_conn_id=aws_conn_id)
    return hook.check_for_key(key=s3_key, bucket_name=bucket_name)


with DAG("example_operators", start_date=datetime(2024, 1, 1), catchup=False) as dag:
    seed_operator = DbtSeedLocalOperator(
        profile_config=profile_config,
        project_dir=DBT_PROJ_DIR,
        task_id="seed",
        dbt_cmd_flags=["--select", "raw_customers"],
        install_deps=True,
        append_env=True,
    )

    check_file_uploaded_task = PythonOperator(
        task_id="check_file_uploaded_task",
        python_callable=check_s3_file,
        op_kwargs={
            "aws_conn_id": "aws_s3_conn",
            "bucket_name": "cosmos-artifacts-upload",
            "file_key": "target/run_results.json",
        },
    )

    run_operator = DbtRunLocalOperator(
        profile_config=profile_config,
        project_dir=DBT_PROJ_DIR,
        task_id="run",
        dbt_cmd_flags=["--models", "stg_customers"],
        install_deps=True,
        append_env=True,
    )

    clone_operator = DbtCloneLocalOperator(
        profile_config=profile_config,
        project_dir=DBT_PROJ_DIR,
        task_id="clone",
        dbt_cmd_flags=["--models", "stg_customers", "--state", DBT_ARTIFACT],
        install_deps=True,
        append_env=True,
    )

    seed_operator >> run_operator >> clone_operator
    seed_operator >> check_file_uploaded_task
```

### Setting Dependencies on Individual Cosmos Tasks

```python
from cosmos import DbtDag, DbtResourceType
from airflow.sdk import task, chain

with DbtDag(...) as dag:
    @task
    def upstream_task():
        pass

    _upstream = upstream_task()

    for unique_id, dbt_node in dag.dbt_graph.filtered_nodes.items():
        if dbt_node.resource_type == DbtResourceType.SEED:
            my_dbt_task = dag.tasks_map[unique_id]
            chain(_upstream, my_dbt_task)
```

---

## 8. Safety Checks

Before finalizing, verify:

- [ ] Execution mode matches constraints (AIRFLOW_ASYNC → BigQuery only)
- [ ] Warehouse adapter installed for chosen execution mode
- [ ] Secrets via Airflow connections or env vars, NOT plaintext
- [ ] Load mode matches execution (complex selectors → dbt_ls)
- [ ] Airflow 3 asset URIs if downstream DAGs scheduled on Cosmos assets (see Appendix A)

---

## Appendix A: Airflow 3 Compatibility

### Import Differences

| Airflow 3.x | Airflow 2.x |
|-------------|-------------|
| `from airflow.sdk import dag, task` | `from airflow.decorators import dag, task` |
| `from airflow.sdk import chain` | `from airflow.models.baseoperator import chain` |

### Asset/Dataset URI Format Change

Cosmos ≤1.9 (Airflow 2 Datasets):
```
postgres://0.0.0.0:5434/postgres.public.orders
```

Cosmos ≥1.10 (Airflow 3 Assets):
```
postgres://0.0.0.0:5434/postgres/public/orders
```

> **CRITICAL**: Update asset URIs when upgrading to Airflow 3.

---

## Appendix B: Operational Extras

### Caching

Cosmos caches artifacts to speed up parsing. Enabled by default.

Reference: https://astronomer.github.io/astronomer-cosmos/configuration/caching.html

### Memory-Optimized Imports

```bash
AIRFLOW__COSMOS__ENABLE_MEMORY_OPTIMISED_IMPORTS=True
```

When enabled:
```python
from cosmos.airflow.dag import DbtDag  # instead of: from cosmos import DbtDag
```

### Artifact Upload to Object Storage

```bash
AIRFLOW__COSMOS__REMOTE_TARGET_PATH=s3://bucket/target_dir/
AIRFLOW__COSMOS__REMOTE_TARGET_PATH_CONN_ID=aws_default
```

```python
from cosmos.io import upload_to_cloud_storage

my_dag = DbtDag(
    # ...
    operator_args={"callback": upload_to_cloud_storage},
)
```

### dbt Docs Hosting

Cosmos serves dbt docs in the Airflow UI. The config depends on your Airflow major
version (each uses a different UI plugin system) — it is not a free single-vs-multi choice:

| Airflow         | Config                                                            | Scope                | Since          |
|-----------------|------------------------------------------------------------------|----------------------|----------------|
| 2 (FAB plugin)  | `DBT_DOCS_DIR` (+ `DBT_DOCS_CONN_ID`, `DBT_DOCS_INDEX_FILE_NAME`) | Single project       | Cosmos 1.4.0+  |
| 3.1+ (FastAPI)  | `DBT_DOCS_PROJECTS` (JSON)                                        | One or more projects | Cosmos 1.11.0+ |

Airflow 2:

```bash
AIRFLOW__COSMOS__DBT_DOCS_DIR="path/to/docs"                   # local path or S3/GCS/Azure/HTTP URI; defaults to the dbt target/ folder
AIRFLOW__COSMOS__DBT_DOCS_CONN_ID="my_conn_id"                 # optional; for cloud storage
AIRFLOW__COSMOS__DBT_DOCS_INDEX_FILE_NAME="static_index.html"  # optional; only if docs built with --static
```

Airflow 3.1+:

```bash
AIRFLOW__COSMOS__DBT_DOCS_PROJECTS='{
    "my_project": {
        "dir": "s3://bucket/docs/",
        "index": "index.html",
        "conn_id": "aws_default",
        "name": "My Project"
    }
}'
```

Pick by Airflow version, not project count. The single-project settings are the Airflow 2
path; Cosmos publishes no deprecation notice for them — do not describe them as "legacy"
or "deprecated."

Reference: https://astronomer.github.io/astronomer-cosmos/configuration/hosting-docs.html

---

## Related Skills

- **cosmos-dbt-fusion**: For dbt Fusion projects (not dbt Core)
- **authoring-dags**: General DAG authoring patterns
- **testing-dags**: Testing DAGs after creation

Referenced files: 1

cosmos-dbt-fusion6.98 KB

View saved version →

---
name: cosmos-dbt-fusion
description: Run a dbt Fusion project with Astronomer Cosmos. Use when running a dbt Fusion project with Astronomer Cosmos (Cosmos 1.11+, ExecutionMode.LOCAL on Snowflake/Databricks). Before implementing, verify dbt engine is Fusion (not Core), the warehouse is supported, and local execution is acceptable. Does not cover dbt Core.
---

# Cosmos + dbt Fusion: Implementation Checklist

Execute steps in order. This skill covers Fusion-specific constraints only.

> **Version note**: dbt Fusion support was introduced in Cosmos 1.11.0. Requires Cosmos ≥1.11.
>
> **Reference**: See **[reference/cosmos-config.md](reference/cosmos-config.md)** for ProfileConfig, operator_args, and Airflow 3 compatibility details.

> **Before starting**, confirm: (1) dbt engine = Fusion (not Core → use **cosmos-dbt-core**), (2) warehouse = Snowflake, Databricks, Bigquery and Redshift only.

### Fusion-Specific Constraints

| Constraint | Details |
|------------|---------|
| No async | `AIRFLOW_ASYNC` not supported |
| No virtualenv | Fusion is a binary, not a Python package |
| Warehouse support | Snowflake, Databricks, Bigquery and Redshift support [while in preview](https://github.com/dbt-labs/dbt-fusion) |

---

## 1. Confirm Cosmos Version

> **CRITICAL**: Cosmos 1.11.0 introduced dbt Fusion compatibility.

```bash
# Check installed version
pip show astronomer-cosmos

# Install/upgrade if needed
pip install "astronomer-cosmos>=1.11.0"
```

**Validate**: `pip show astronomer-cosmos` reports version ≥ 1.11.0

---

## 2. Install the dbt Fusion Binary (REQUIRED)

dbt Fusion is NOT bundled with Cosmos or dbt Core. Install it into the Airflow runtime/image.

Determine where to install the Fusion binary (Dockerfile / base image / runtime).

### Example Dockerfile Install

```dockerfile
USER root
RUN apt-get update && apt-get install -y curl
ENV SHELL=/bin/bash
RUN curl -fsSL https://public.cdn.getdbt.com/fs/install/install.sh | sh -s -- --update
USER astro
```

### Common Install Paths

| Environment | Typical path |
|-------------|--------------|
| Astro Runtime | `/home/astro/.local/bin/dbt` |
| System-wide | `/usr/local/bin/dbt` |

**Validate**: The `dbt` binary exists at the chosen path and `dbt --version` succeeds.

---

## 3. Choose Parsing Strategy (RenderConfig)

Parsing strategy is the same as dbt Core. Pick ONE:

| Load mode | When to use | Required inputs |
|-----------|-------------|-----------------|
| `dbt_manifest` | Large projects; fastest parsing | `ProjectConfig.manifest_path` |
| `dbt_ls` | Complex selectors; need dbt-native selection | Fusion binary accessible to scheduler |
| `automatic` | Simple setups; let Cosmos pick | (none) |

```python
from cosmos import RenderConfig, LoadMode

_render_config = RenderConfig(
    load_method=LoadMode.AUTOMATIC,  # or DBT_MANIFEST, DBT_LS
)
```

---

## 4. Configure Warehouse Connection (ProfileConfig)

> **Reference**: See **[reference/cosmos-config.md](reference/cosmos-config.md#profileconfig-warehouse-connection)** for full ProfileConfig options and examples.


```python
from cosmos import ProfileConfig
from cosmos.profiles import SnowflakeUserPasswordProfileMapping

_profile_config = ProfileConfig(
    profile_name="default",
    target_name="dev",
    profile_mapping=SnowflakeUserPasswordProfileMapping(
        conn_id="snowflake_default",
    ),
)
```

---

## 5. Configure ExecutionConfig (LOCAL Only)

> **CRITICAL**: dbt Fusion with Cosmos requires `ExecutionMode.LOCAL` with `dbt_executable_path` pointing to the Fusion binary.

```python
from cosmos import ExecutionConfig
from cosmos.constants import InvocationMode

_execution_config = ExecutionConfig(
    invocation_mode=InvocationMode.SUBPROCESS,
    dbt_executable_path="/home/astro/.local/bin/dbt",  # REQUIRED: path to Fusion binary
    # execution_mode is LOCAL by default - do not change
)
```

---

## 6. Configure Project (ProjectConfig)

```python
from cosmos import ProjectConfig

_project_config = ProjectConfig(
    dbt_project_path="/path/to/dbt/project",
    # manifest_path="/path/to/manifest.json",  # for dbt_manifest load mode
    # install_dbt_deps=False,  # if deps precomputed in CI
)
```

---

## 7. Assemble DAG / TaskGroup

### Option A: DbtDag (Standalone)

```python
from cosmos import DbtDag, ProjectConfig, ProfileConfig, ExecutionConfig, RenderConfig
from cosmos.profiles import SnowflakeUserPasswordProfileMapping
from pendulum import datetime

_project_config = ProjectConfig(
    dbt_project_path="/usr/local/airflow/dbt/my_project",
)

_profile_config = ProfileConfig(
    profile_name="default",
    target_name="dev",
    profile_mapping=SnowflakeUserPasswordProfileMapping(
        conn_id="snowflake_default",
    ),
)

_execution_config = ExecutionConfig(
    dbt_executable_path="/home/astro/.local/bin/dbt",  # Fusion binary
)

_render_config = RenderConfig()

my_fusion_dag = DbtDag(
    dag_id="my_fusion_cosmos_dag",
    project_config=_project_config,
    profile_config=_profile_config,
    execution_config=_execution_config,
    render_config=_render_config,
    start_date=datetime(2025, 1, 1),
    schedule="@daily",
)
```

### Option B: DbtTaskGroup (Inside Existing DAG)

```python
from airflow.sdk import dag, task  # Airflow 3.x
# from airflow.decorators import dag, task  # Airflow 2.x
from airflow.models.baseoperator import chain
from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig, ExecutionConfig
from pendulum import datetime

_project_config = ProjectConfig(dbt_project_path="/usr/local/airflow/dbt/my_project")
_profile_config = ProfileConfig(profile_name="default", target_name="dev")
_execution_config = ExecutionConfig(dbt_executable_path="/home/astro/.local/bin/dbt")

@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
def my_dag():
    @task
    def pre_dbt():
        return "some_value"

    dbt = DbtTaskGroup(
        group_id="dbt_fusion_project",
        project_config=_project_config,
        profile_config=_profile_config,
        execution_config=_execution_config,
    )

    @task
    def post_dbt():
        pass

    chain(pre_dbt(), dbt, post_dbt())

my_dag()
```

---

## 8. Final Validation

Before finalizing, verify:

- [ ] **Cosmos version**: ≥1.11.0
- [ ] **Fusion binary installed**: Path exists and is executable
- [ ] **Warehouse supported**: Snowflake, Databricks, Bigquery or Redshift only
- [ ] **Secrets handling**: Airflow connections or env vars, NOT plaintext

### Troubleshooting

If user reports dbt Core regressions after enabling Fusion:

```bash
AIRFLOW__COSMOS__PRE_DBT_FUSION=1
```

### User Must Test

- [ ] The DAG parses in the Airflow UI (no import/parse-time errors)
- [ ] A manual run succeeds against the target warehouse (at least one model)

---

## Reference

- Cosmos dbt Fusion docs: https://astronomer.github.io/astronomer-cosmos/configuration/dbt-fusion.html
- dbt Fusion install: https://docs.getdbt.com/docs/core/pip-install#dbt-fusion

---

## Related Skills

- **cosmos-dbt-core**: For dbt Core projects (not Fusion)
- **authoring-dags**: General DAG authoring patterns
- **testing-dags**: Testing DAGs after creation

Referenced files: 1

creating-openlineage-extractors12.9 KB

View saved version →

---
name: creating-openlineage-extractors
description: Create custom OpenLineage extractors for Airflow operators. Use when the user needs lineage from unsupported or third-party operators, wants column-level lineage, or needs complex extraction logic beyond what inlets/outlets provide.
---

# Creating OpenLineage Extractors

This skill guides you through creating custom OpenLineage extractors to capture lineage from Airflow operators that don't have built-in support.

> **Reference:** See the [OpenLineage provider developer guide](https://airflow.apache.org/docs/apache-airflow-providers-openlineage/stable/guides/developer.html) for the latest patterns and list of supported operators/hooks.

## When to Use Each Approach

| Scenario | Approach |
|----------|----------|
| Operator you own/maintain | **OpenLineage Methods** (recommended, simplest) |
| Third-party operator you can't modify | Custom Extractor |
| Need column-level lineage | OpenLineage Methods or Custom Extractor |
| Complex extraction logic | OpenLineage Methods or Custom Extractor |
| Simple table-level lineage | Inlets/Outlets (simplest, but lowest priority) |

> **Important:** Always prefer OpenLineage methods over custom extractors when possible. Extractors are harder to write, easier to diverge from operator behavior after changes, and harder to debug.

### On Astro

Astro includes built-in OpenLineage integration — no additional transport configuration is needed. Lineage events are automatically collected and displayed in the Astro UI's **Lineage tab**. Custom extractors deployed to an Astro project are automatically picked up, so you only need to register them in `airflow.cfg` or via environment variable and deploy.

---

## Two Approaches

### 1. OpenLineage Methods (Recommended)

Use when you can add methods directly to your custom operator. This is the **go-to solution** for operators you own.

### 2. Custom Extractors

Use when you need lineage from third-party or provider operators that you **cannot modify**.

---

## Approach 1: OpenLineage Methods (Recommended)

When you own the operator, add OpenLineage methods directly:

```python
from airflow.models import BaseOperator


class MyCustomOperator(BaseOperator):
    """Custom operator with built-in OpenLineage support."""

    def __init__(self, source_table: str, target_table: str, **kwargs):
        super().__init__(**kwargs)
        self.source_table = source_table
        self.target_table = target_table
        self._rows_processed = 0  # Set during execution

    def execute(self, context):
        # Do the actual work
        self._rows_processed = self._process_data()
        return self._rows_processed

    def get_openlineage_facets_on_start(self):
        """Called when task starts. Return known inputs/outputs."""
        # Import locally to avoid circular imports
        from openlineage.client.event_v2 import Dataset
        from airflow.providers.openlineage.extractors import OperatorLineage

        return OperatorLineage(
            inputs=[Dataset(namespace="postgres://db", name=self.source_table)],
            outputs=[Dataset(namespace="postgres://db", name=self.target_table)],
        )

    def get_openlineage_facets_on_complete(self, task_instance):
        """Called after success. Add runtime metadata."""
        from openlineage.client.event_v2 import Dataset
        from openlineage.client.facet_v2 import output_statistics_output_dataset
        from airflow.providers.openlineage.extractors import OperatorLineage

        return OperatorLineage(
            inputs=[Dataset(namespace="postgres://db", name=self.source_table)],
            outputs=[
                Dataset(
                    namespace="postgres://db",
                    name=self.target_table,
                    facets={
                        "outputStatistics": output_statistics_output_dataset.OutputStatisticsOutputDatasetFacet(
                            rowCount=self._rows_processed
                        )
                    },
                )
            ],
        )

    def get_openlineage_facets_on_failure(self, task_instance):
        """Called after failure. Optional - for partial lineage."""
        return None
```

### OpenLineage Methods Reference

| Method | When Called | Required |
|--------|-------------|----------|
| `get_openlineage_facets_on_start()` | Task enters RUNNING | No |
| `get_openlineage_facets_on_complete(ti)` | Task succeeds | No |
| `get_openlineage_facets_on_failure(ti)` | Task fails | No |

> Implement only the methods you need. Unimplemented methods fall through to Hook-Level Lineage or inlets/outlets.

---

## Approach 2: Custom Extractors

Use this approach only when you **cannot modify** the operator (e.g., third-party or provider operators).

### Basic Structure

```python
from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage
from openlineage.client.event_v2 import Dataset


class MyOperatorExtractor(BaseExtractor):
    """Extract lineage from MyCustomOperator."""

    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        """Return operator class names this extractor handles."""
        return ["MyCustomOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        """Called BEFORE operator executes. Use for known inputs/outputs."""
        # Access operator properties via self.operator
        source_table = self.operator.source_table
        target_table = self.operator.target_table

        return OperatorLineage(
            inputs=[
                Dataset(
                    namespace="postgres://mydb:5432",
                    name=f"public.{source_table}",
                )
            ],
            outputs=[
                Dataset(
                    namespace="postgres://mydb:5432",
                    name=f"public.{target_table}",
                )
            ],
        )

    def extract_on_complete(self, task_instance) -> OperatorLineage | None:
        """Called AFTER operator executes. Use for runtime-determined lineage."""
        # Access properties set during execution
        # Useful for operators that determine outputs at runtime
        return None
```

### OperatorLineage Structure

```python
from airflow.providers.openlineage.extractors.base import OperatorLineage
from openlineage.client.event_v2 import Dataset
from openlineage.client.facet_v2 import sql_job

lineage = OperatorLineage(
    inputs=[Dataset(namespace="...", name="...")],      # Input datasets
    outputs=[Dataset(namespace="...", name="...")],     # Output datasets
    run_facets={"sql": sql_job.SQLJobFacet(query="SELECT...")},  # Run metadata
    job_facets={},                                      # Job metadata
)
```

### Extraction Methods

| Method | When Called | Use For |
|--------|-------------|---------|
| `_execute_extraction()` | Before operator runs | Static/known lineage |
| `extract_on_complete(task_instance)` | After success | Runtime-determined lineage |
| `extract_on_failure(task_instance)` | After failure | Partial lineage on errors |

### Registering Extractors

**Option 1: Configuration file (`airflow.cfg`)**

```ini
[openlineage]
extractors = mypackage.extractors.MyOperatorExtractor;mypackage.extractors.AnotherExtractor
```

**Option 2: Environment variable**

```bash
AIRFLOW__OPENLINEAGE__EXTRACTORS='mypackage.extractors.MyOperatorExtractor;mypackage.extractors.AnotherExtractor'
```

> **Important:** The path must be importable from the Airflow worker. Place extractors in your DAGs folder or installed package.

---

## Common Patterns

### SQL Operator Extractor

```python
from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage
from openlineage.client.event_v2 import Dataset
from openlineage.client.facet_v2 import sql_job


class MySqlOperatorExtractor(BaseExtractor):
    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        return ["MySqlOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        sql = self.operator.sql
        conn_id = self.operator.conn_id

        # Parse SQL to find tables (simplified example)
        # In practice, use a SQL parser like sqlglot
        inputs, outputs = self._parse_sql(sql)

        namespace = f"postgres://{conn_id}"

        return OperatorLineage(
            inputs=[Dataset(namespace=namespace, name=t) for t in inputs],
            outputs=[Dataset(namespace=namespace, name=t) for t in outputs],
            job_facets={
                "sql": sql_job.SQLJobFacet(query=sql)
            },
        )

    def _parse_sql(self, sql: str) -> tuple[list[str], list[str]]:
        """Parse SQL to extract table names. Use sqlglot for real parsing."""
        # Simplified example - use proper SQL parser in production
        inputs = []
        outputs = []
        # ... parsing logic ...
        return inputs, outputs
```

### File Transfer Extractor

```python
from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage
from openlineage.client.event_v2 import Dataset


class S3ToSnowflakeExtractor(BaseExtractor):
    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        return ["S3ToSnowflakeOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        s3_bucket = self.operator.s3_bucket
        s3_key = self.operator.s3_key
        table = self.operator.table
        schema = self.operator.schema

        return OperatorLineage(
            inputs=[
                Dataset(
                    namespace=f"s3://{s3_bucket}",
                    name=s3_key,
                )
            ],
            outputs=[
                Dataset(
                    namespace="snowflake://myaccount.snowflakecomputing.com",
                    name=f"{schema}.{table}",
                )
            ],
        )
```

### Dynamic Lineage from Execution

```python
from openlineage.client.event_v2 import Dataset


class DynamicOutputExtractor(BaseExtractor):
    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        return ["DynamicOutputOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        # Only inputs known before execution
        return OperatorLineage(
            inputs=[Dataset(namespace="...", name=self.operator.source)],
        )

    def extract_on_complete(self, task_instance) -> OperatorLineage | None:
        # Outputs determined during execution
        # Access via operator properties set in execute()
        outputs = self.operator.created_tables  # Set during execute()

        return OperatorLineage(
            inputs=[Dataset(namespace="...", name=self.operator.source)],
            outputs=[Dataset(namespace="...", name=t) for t in outputs],
        )
```

---

## Common Pitfalls

### 1. Circular Imports

**Problem:** Importing Airflow modules at the top level causes circular imports.

```python
# ❌ BAD - can cause circular import issues
from airflow.models import TaskInstance
from openlineage.client.event_v2 import Dataset

class MyExtractor(BaseExtractor):
    ...
```

```python
# ✅ GOOD - import inside methods
class MyExtractor(BaseExtractor):
    def _execute_extraction(self):
        from openlineage.client.event_v2 import Dataset
        # ...
```

### 2. Wrong Import Path

**Problem:** Extractor path doesn't match actual module location.

```bash
# ❌ Wrong - path doesn't exist
AIRFLOW__OPENLINEAGE__EXTRACTORS='extractors.MyExtractor'

# ✅ Correct - full importable path
AIRFLOW__OPENLINEAGE__EXTRACTORS='dags.extractors.my_extractor.MyExtractor'
```

### 3. Not Handling None

**Problem:** Extraction fails when operator properties are None.

```python
# ✅ Handle optional properties
def _execute_extraction(self) -> OperatorLineage | None:
    if not self.operator.source_table:
        return None  # Skip extraction

    return OperatorLineage(...)
```

---

## Testing Extractors

### Unit Testing

```python
import pytest
from unittest.mock import MagicMock
from mypackage.extractors import MyOperatorExtractor


def test_extractor():
    # Mock the operator
    operator = MagicMock()
    operator.source_table = "input_table"
    operator.target_table = "output_table"

    # Create extractor
    extractor = MyOperatorExtractor(operator)

    # Test extraction
    lineage = extractor._execute_extraction()

    assert len(lineage.inputs) == 1
    assert lineage.inputs[0].name == "input_table"
    assert len(lineage.outputs) == 1
    assert lineage.outputs[0].name == "output_table"
```

---

## Precedence Rules

OpenLineage checks for lineage in this order:

1. **Custom Extractors** (highest priority)
2. **OpenLineage Methods** on operator
3. **Hook-Level Lineage** (from `HookLineageCollector`)
4. **Inlets/Outlets** (lowest priority)

If a custom extractor exists, it overrides built-in extraction and inlets/outlets.

---

## Related Skills

- **annotating-task-lineage**: For simple table-level lineage with inlets/outlets
- **tracing-upstream-lineage**: Investigate data origins
- **tracing-downstream-lineage**: Investigate data dependencies
dag-factory16.9 KB

View saved version →

---
name: dag-factory
description: Authors Apache Airflow DAGs declaratively from dag-factory YAML configs. Use when building DAGs declaratively from YAML via dag-factory; creating/editing dag-factory templates/YAML configs,reating/editing dag-factory YAML configs, defaults, dynamic tasks, datasets, or callbacks; or validating dag-factory configurations; upgrading or re-pinning dag-factory.
---
# DAG Factory

You are helping a user build Apache Airflow DAGs declaratively with **dag-factory**, a library that turns YAML configuration files into Airflow DAGs. Execute steps in order and prefer the simplest configuration that meets the user's needs.

> **Package**: `dag-factory` on PyPI
> **Repo**: https://github.com/astronomer/dag-factory
> **Docs**: https://astronomer.github.io/dag-factory/latest/
> **Targets**: dag-factory **v1.0+** only. For pre-1.0 projects, see [reference/migration.md](reference/migration.md) before applying any guidance from this skill.
> **Requires**: Python 3.10+, Airflow 2.4+ (Airflow 3 supported)

## Before Starting

Confirm with the user:
1. **Airflow version** ≥2.4
2. **Python version** ≥3.10
3. **dag-factory version**: this skill targets **v1.0+**. If the project is on <1.0, follow [reference/migration.md](reference/migration.md) to upgrade before continuing.
4. **Use case**: dag-factory is for declarative, low-code DAG authoring. If the user needs reusable, validated Pythonic templates with Pydantic, suggest **blueprint** instead. If they need full Python flexibility, suggest the **authoring-dags** skill.

---

## Determine What the User Needs

| User Request | Action |
|--------------|--------|
| "Create a YAML DAG" / "Convert this Python DAG to YAML" | Go to **Defining a DAG in YAML** |
| "Set up dag-factory in my project" | Go to **Project Setup** |
| "Share defaults across DAGs" / "Set start_date once" | Go to **Defaults** |
| "Use a custom operator" / "Use KPO / Slack / Snowflake" | Go to **Custom & Provider Operators** |
| "Dynamic / mapped tasks" / "expand / partial" | Go to **Dynamic Task Mapping** |
| "Schedule on dataset" / "Outlets and inlets" | Go to **Datasets** |
| "Add a callback" / "Slack on failure" | Go to **Callbacks** |
| "Use a timetable" / "datetime in YAML" / "timedelta in YAML" | Go to **Custom Python Objects (`__type__`)** |
| "Lint my YAML" / "Validate" | Go to **Validation Commands** |
| "Convert Airflow 2 YAML to Airflow 3" | Go to **Validation Commands** (`dagfactory convert`) |
| "Migrate from dag-factory <1.0" | See [reference/migration.md](reference/migration.md) |
| dag-factory errors / troubleshooting | Go to **Troubleshooting** |

---

## Project Setup

### 1. Install the Package

Add to `requirements.txt`:

```
dag-factory>=1.0.0
```

dag-factory **does not** install Airflow providers automatically. Install any provider packages your YAML references (e.g., `apache-airflow-providers-slack`, `apache-airflow-providers-cncf-kubernetes`).

### 2. Create the Loader

Create `dags/load_dags.py` so Airflow's DAG processor will pick it up:

```python
import os
from pathlib import Path

from dagfactory import load_yaml_dags

CONFIG_ROOT_DIR = Path(os.getenv("CONFIG_ROOT_DIR", "/usr/local/airflow/dags/"))

# Option A: load every *.yml / *.yaml under a folder
load_yaml_dags(globals_dict=globals(), dags_folder=str(CONFIG_ROOT_DIR))

# Option B: load a single file
# load_yaml_dags(globals_dict=globals(), config_filepath=str(CONFIG_ROOT_DIR / "my_dag.yml"))

# Option C: load from an in-Python dict
# load_yaml_dags(globals_dict=globals(), config_dict={...})
```

`globals_dict=globals()` is required so generated DAG objects are registered into the module namespace where Airflow can discover them.

### 3. Verify Installation

```bash
dagfactory --version
```

---

## Defining a DAG in YAML

Each top-level YAML key (other than `default`) defines a DAG. The key becomes the `dag_id`. **Use the list format for `tasks` and `task_groups`** — it is the recommended format since v1.0.0.

```yaml
# dags/example_dag_factory.yml
default:
  default_args:
    start_date: 2024-11-11

basic_example_dag:
  default_args:
    owner: "custom_owner"
  description: "this is an example dag"
  schedule: "0 3 * * *"
  catchup: false
  task_groups:
    - group_name: "example_task_group"
      tooltip: "this is an example task group"
      dependencies: [task_1]
  tasks:
    - task_id: "task_1"
      operator: airflow.operators.bash.BashOperator
      bash_command: "echo 1"
    - task_id: "task_2"
      operator: airflow.operators.bash.BashOperator
      bash_command: "echo 2"
      dependencies: [task_1]
    - task_id: "task_3"
      operator: airflow.operators.bash.BashOperator
      bash_command: "echo 3"
      dependencies: [task_1]
      task_group_name: "example_task_group"
```

### Key Fields

| Field | Where | Purpose |
|-------|-------|---------|
| `default` | top-level | Shared DAG-level args applied to every DAG in this file |
| `default_args` | DAG or `default` block | Standard Airflow `default_args` (owner, retries, start_date, ...) |
| `schedule` | DAG | Cron expression, preset (`@daily`), Dataset list, or `__type__` timetable |
| `catchup` / `description` / `tags` | DAG | Standard Airflow DAG kwargs |
| `tasks` | DAG | List of task dicts; each requires `task_id` and `operator` |
| `operator` | task | **Full import path** to operator class (e.g. `airflow.operators.bash.BashOperator`) |
| `dependencies` | task / task_group | List of upstream `task_id`s or `group_name`s |
| `task_groups` | DAG | List of group dicts; each requires `group_name` |
| `task_group_name` | task | Assigns a task to a task group |

Tasks do **not** need to be ordered by dependency in the YAML — dag-factory resolves the DAG topology.

### Dictionary Format (Legacy)

Pre-1.0 dictionary format (where `tasks` is a dict keyed by `task_id`) still works for backward compatibility, but prefer the list format for new code.

---

## Defaults

There are four ways to set defaults, in **precedence order** (highest first):

1. `default_args` / DAG-level keys inside an individual DAG
2. The top-level `default:` block in the same YAML file
3. `defaults_config_dict=` argument to `load_yaml_dags`
4. A `defaults.yml` (or `defaults.yaml`) file via `defaults_config_path=` (or auto-detected next to the DAG YAML)

> Note: loader argument names and several other field names changed in v1.0.0. See [reference/migration.md](reference/migration.md) if you're working on an older project.

### `default` Block in the Same File

Powerful for templating multiple DAGs from one file:

```yaml
default:
  default_args:
    owner: "data-team"
    start_date: 2025-01-01
    retries: 2
  catchup: false
  schedule: "@daily"

dag_one:
  description: "first DAG"
  tasks:
    - task_id: t1
      operator: airflow.operators.bash.BashOperator
      bash_command: "echo one"

dag_two:
  description: "second DAG"
  tasks:
    - task_id: t1
      operator: airflow.operators.bash.BashOperator
      bash_command: "echo two"
```

### `defaults.yml` File

Place a `defaults.yml` next to the DAG YAML, or point `defaults_config_path` at a parent directory. dag-factory **merges** all `defaults.yml` files walking up the directory tree, with the file closest to the DAG YAML winning. DAG-level args (e.g. `schedule`, `catchup`) go at the root of `defaults.yml`; per-task defaults go under `default_args`.

```yaml
# defaults.yml
schedule: 0 1 * * *
catchup: false
default_args:
  start_date: '2024-12-31'
  owner: data-team
```

---

## Custom & Provider Operators

Reference any operator by its **full Python import path**. dag-factory passes all other task keys as kwargs to that operator.

```yaml
tasks:
  - task_id: begin
    operator: airflow.providers.standard.operators.empty.EmptyOperator
  - task_id: make_bread
    operator: customized.operators.breakfast_operators.MakeBreadOperator
    bread_type: 'Sourdough'
```

The operator's package must be installed and importable. For Airflow 3, prefer `airflow.providers.standard.operators.*` over the legacy `airflow.operators.*` paths — the `dagfactory convert` CLI rewrites these automatically.

### KubernetesPodOperator

Specify the operator path and pass kwargs directly. As of v1.0, dag-factory no longer does legacy type casting — use `__type__` for nested k8s objects.

```yaml
tasks:
  - task_id: hello-world-pod
    operator: airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator
    image: "python:3.12-slim"
    cmds: ["python", "-c"]
    arguments: ["print('hi')"]
    name: example-pod
    namespace: default
    container_resources:
      __type__: kubernetes.client.models.V1ResourceRequirements
      limits: {cpu: "1", memory: "1024Mi"}
      requests: {cpu: "0.5", memory: "512Mi"}
```

---

## Dynamic Task Mapping

Use `expand` and `partial` keys on a task to map dynamically. dag-factory has two distinct ways to reference an upstream task's output:

- **`task_id.output`** — XCom-style reference, used inside `expand` `op_args` / `op_kwargs` (and the equivalent kwargs of other operators).
- **`+task_id`** — bare value reference, used when the value sits directly under `expand` (e.g. `expand: {number: +numbers_list}`) or as a TaskFlow decorator argument.

Don't mix them: `+request` won't resolve inside `op_args`, and `request.output` won't resolve as a bare `expand` value.

```yaml
dynamic_task_map:
  default_args:
    start_date: 2025-01-01
  schedule: "0 3 * * *"
  tasks:
    - task_id: request
      operator: airflow.providers.standard.operators.python.PythonOperator
      python_callable_name: make_list
      python_callable_file: $CONFIG_ROOT_DIR/expand_tasks.py
    - task_id: process
      operator: airflow.providers.standard.operators.python.PythonOperator
      python_callable_name: consume_value
      python_callable_file: $CONFIG_ROOT_DIR/expand_tasks.py
      partial:
        op_kwargs:
          fixed_param: "test"
      expand:
        op_args: request.output    # XCom-style — used inside op_args / op_kwargs
      dependencies: [request]
```

Bare-value form (TaskFlow `decorator` tasks, or any non-`op_args` mapping):

```yaml
tasks:
  - task_id: numbers_list
    decorator: airflow.sdk.definitions.decorators.task
    python_callable: sample.build_numbers_list
  - task_id: double_number
    decorator: airflow.sdk.definitions.decorators.task
    python_callable: sample.double
    expand:
      number: +numbers_list   # + resolves to upstream task `numbers_list`'s XComArg
```

For named map indices (Airflow 2.9+), set `map_index_template: "{{ task.custom_mapping_key }}"` and have the callable assign `context["custom_mapping_key"]`.

**Tested patterns**: simple mapping, task-generated mapping, repeated mapping, `partial`, multiple-parameter mapping, `map_index_template`.
**Unsupported / untested**: mapping over task groups, zipping, transforming expanding data.

---

## Datasets

Use `inlets` / `outlets` on tasks to declare dataset producers, and a list of dataset URIs as `schedule` to consume them.

```yaml
producer_dag:
  default_args:
    start_date: '2024-01-01'
  schedule: "0 5 * * *"
  catchup: false
  tasks:
    - task_id: task_1
      operator: airflow.operators.bash.BashOperator
      bash_command: "echo 1"
      outlets: ['s3://bucket_example/raw/dataset1.json']

consumer_dag:
  default_args:
    start_date: '2024-01-01'
  schedule: ['s3://bucket_example/raw/dataset1.json']
  catchup: false
  tasks:
    - task_id: task_1
      operator: airflow.operators.bash.BashOperator
      bash_command: "echo 'consumer'"
```

### Conditional Dataset Scheduling (Airflow 2.9+ / dag-factory 0.22+)

Nesting the logical operators `__and__` / `__or__` under `datasets` key.

```yaml
schedule:
  datasets:
    __or__:
      - __and__:
          - s3://bucket-cjmm/raw/dataset_custom_1
          - s3://bucket-cjmm/raw/dataset_custom_2
      - s3://bucket-cjmm/raw/dataset_custom_3
```

---

## Callbacks

Three styles, all valid at the DAG, TaskGroup, or Task level (or under `default_args`):

### 1. String pointing to a callable

```yaml
- task_id: task_1
  operator: airflow.operators.bash.BashOperator
  bash_command: "echo task_1"
  on_failure_callback: include.custom_callbacks.output_standard_message
```

With kwargs:

```yaml
- task_id: task_2
  operator: airflow.operators.bash.BashOperator
  bash_command: "echo task_2"
  on_success_callback:
    callback: include.custom_callbacks.output_custom_message
    param1: "Task status"
    param2: "Successful!"
```

### 2. File path + function name (no kwargs)

```yaml
- task_id: task_3
  operator: airflow.operators.bash.BashOperator
  bash_command: "echo task_3"
  on_retry_callback_name: output_standard_message
  on_retry_callback_file: /usr/local/airflow/include/custom_callbacks.py
```

### 3. Provider callbacks

```yaml
- task_id: task_4
  operator: airflow.operators.bash.BashOperator
  bash_command: "echo task_4"
  on_failure_callback:
    callback: airflow.providers.slack.notifications.slack.send_slack_notification
    slack_conn_id: slack_conn_id
    text: ":red_circle: Task Failed."
    channel: "#channel"
```

The provider package must be installed.

---

## Custom Python Objects (`__type__`)

For anything that isn't a simple scalar — `datetime`, `timedelta`, `Asset`, timetables, k8s objects — use the generalized object syntax:

```yaml
start_date:
  __type__: datetime.datetime
  year: 2025
  month: 1
  day: 1

execution_timeout:
  __type__: datetime.timedelta
  hours: 1

schedule:
  __type__: airflow.timetables.trigger.CronTriggerTimetable
  cron: "0 1 * * 3"
  timezone: UTC
```

- `__type__` is the **full import path** to the class
- `__args__` is a list of positional arguments
- Other keys become keyword arguments
- For lists of typed objects, use `__type__: builtins.list` with an `items:` key

### Reserved Keys

Don't use these YAML keys for your own data — dag-factory reserves them: `__type__`, `__args__`, `__join__`, `__and__`, `__or__`. The key `items` is also reserved when used inside a `__type__: builtins.list` block — don't add a custom field named `items` to a typed list construction.

---

## Validation Commands

After installing, the `dagfactory` CLI is on PATH:

| Command | When to Use |
|---------|-------------|
| `dagfactory --version` | Confirm install / version |
| `dagfactory lint <path>` | Validate YAML syntax for a file or directory |
| `dagfactory lint <path> --verbose` | Show a per-file table of results |
| `dagfactory convert <path>` | Show diffs to migrate Airflow 2 → 3 import paths |
| `dagfactory convert <path> --override` | Apply the conversions in place |

### Validation Workflow

```bash
# 1. Lint YAML
dagfactory lint dags/

# 2. Have Airflow parse to catch operator/import errors
#    (Astro CLI users)
astro dev parse
```

`dagfactory lint` only checks YAML syntax — operator import errors and missing kwargs surface at Airflow parse time.

---

## Troubleshooting

### "Operator not found" / `ModuleNotFoundError`

**Cause**: Provider package not installed, or wrong import path.

**Fix**: Install the provider (`pip install apache-airflow-providers-...`) and verify the path. For Airflow 3, run `dagfactory convert` to update legacy `airflow.operators.*` paths to `airflow.providers.standard.operators.*`.

### YAML parses but the DAG doesn't appear in Airflow

**Cause**: Loader file missing or `globals_dict=globals()` not passed.

**Fix**: Ensure a Python file in `dags/` calls `load_yaml_dags(globals_dict=globals(), ...)`. Check `astro dev parse` (or `airflow dags list-import-errors`) for parse errors.

### "Argument is not JSON-serializable" / wrong kwarg type

**Cause**: A scalar string is being passed where a Python object is expected (e.g. `start_date: "2025-01-01"` for a field that needs `datetime`).

**Fix**: Use `__type__: datetime.datetime` (or `datetime.timedelta` etc.) per **Custom Python Objects**.

### Conditional dataset schedule ignored

**Cause**: Airflow <2.9, dag-factory <0.22, or using legacy `!and`/`!or` keys.

**Fix**: Upgrade and rename to `__and__` / `__or__`.

### Multiple `defaults.yml` not merging as expected

**Cause**: `defaults_config_path` not pointing at a parent directory of the DAG YAML.

**Fix**: Set `defaults_config_path` to the highest ancestor folder you want included; dag-factory walks the tree from DAG file → ancestor and merges in that order, with files closer to the DAG winning.

---

## Verification Checklist

Before finishing, verify with the user:

- [ ] `dagfactory lint dags/` passes
- [ ] Loader file exists in `dags/` and calls `load_yaml_dags(globals_dict=globals(), ...)`
- [ ] Required Airflow providers are in `requirements.txt`
- [ ] DAG appears in Airflow UI without import errors

---

## Related Skills

- **authoring-dags** — Writing Airflow DAGs in pure Python with `af` CLI validation. Use when YAML can't express what you need.
- **testing-dags**: For testing DAGs, debugging failures, and the test -> fix -> retest loop
- **debugging-dags**: For troubleshooting failed DAGs

## Reference

- GitHub: https://github.com/astronomer/dag-factory
- Docs: https://astronomer.github.io/dag-factory/latest/
- PyPI: https://pypi.org/project/dag-factory/
- Migration Guide: https://astronomer.github.io/dag-factory/latest/migration_guide/

Referenced files: 1

debugging-dags7.32 KB

View saved version →

---
name: debugging-dags
description: Comprehensive DAG failure diagnosis and root-cause analysis with structured investigation and prevention recommendations. Use when deep failure investigation is needed, a DAG fails to import/parse or 'airflow dags list' errors on a file; a task or run is failing and must be diagnosed and fixed; requests like 'why did X fail', 'my dag keeps failing — find and fix it', or fixing a broken DAG so it loads cleanly. For simple 'why did it fail / show logs', the airflow skill handles it directly.
---

# DAG Diagnosis

You are a data engineer debugging a failed Airflow DAG. Follow this systematic approach to identify the root cause and provide actionable remediation.

## Running the CLI

These commands assume `af` is on PATH. Run via `astro otto` to get it automatically, or install standalone with `uv tool install astro-airflow-mcp`.

---

## Step 1: Identify the Failure

If a specific DAG was mentioned:
- Run `af runs diagnose <dag_id> <dag_run_id>` (if run_id is provided)
- If no run_id specified, run `af dags stats` to find recent failures

If no DAG was specified:
- Run `af health` to find recent failures across all DAGs
- Check for import errors with `af dags errors`
- Show DAGs with recent failures
- Ask which DAG to investigate further

## Step 2: Get the Error Details

Once you have identified a failed task:

1. **Get task logs** using `af tasks logs <dag_id> <dag_run_id> <task_id>`
2. **Look for the actual exception** - scroll past the Airflow boilerplate to find the real error
3. **Categorize the failure type**:
   - **Data issue**: Missing data, schema change, null values, constraint violation
   - **Code issue**: Bug, syntax error, import failure, type error
   - **Infrastructure issue**: Connection timeout, resource exhaustion, permission denied
   - **Dependency issue**: Upstream failure, external API down, rate limiting

## Step 3: Check Context

Gather additional context to understand WHY this happened:

1. **Recent changes**: Was there a code deploy? Check git history if available
2. **Package version changes**: Was a package upgraded — in the image, in a venv-style operator, or at the index? See [Package version changes](#package-version-changes) below.
3. **Data volume**: Did data volume spike? Run a quick count on source tables
4. **Upstream health**: Did upstream tasks succeed but produce unexpected data?
5. **Historical pattern**: Is this a recurring failure? Check if same task failed before
6. **Timing**: Did this fail at an unusual time? (resource contention, maintenance windows)

Use `af runs get <dag_id> <dag_run_id>` to compare the failed run against recent successful runs.

### Package version changes

A common cause of failures with no git activity is dependency drift — the user's code didn't change, but a package they depend on did. Check in this order:

1. **Worker image diff** (preferred when available). Every Astro deploy = new image tag, so the registry has a "before" and "after". Diff `pip freeze` between current and previous image — that's ground truth for what changed:
   ```
   docker run --rm <current_image> pip freeze > /tmp/now.txt
   docker run --rm <previous_image> pip freeze > /tmp/prev.txt
   diff /tmp/prev.txt /tmp/now.txt
   ```
   Also compare `docker run --rm <image> python --version` between the two — a Python minor-version bump (3.11 → 3.12, or even a patch) can break wheel compatibility even when `pip freeze` looks identical. `af config providers` lists currently installed provider versions, useful for cross-checking against modules named in the traceback.

2. **Venv-style operators bypass the worker image.** `@task.virtualenv`, `PythonVirtualenvOperator`, `ExternalPythonOperator`, and `KubernetesPodOperator` build their environment per task run, so an image diff won't catch failures inside them. If the failed task is one of these, read its `requirements` / `image` / `python_version` / `python` args directly:
   - Unbounded specifier (e.g. `pandas>=2.0.0` with no upper bound, or no specifier at all) → a new upstream release is the prime suspect.
   - `image="foo:latest"` or no tag → the image moved underneath you.
   - `python_version="3.11"` (on `@task.virtualenv` / `PythonVirtualenvOperator`) or a `python` path (on `ExternalPythonOperator`) resolving to a different interpreter than it used to — a Python minor-version change can break wheel compatibility for unchanged `requirements`. Same vector applies to the worker image itself if the base Python changed there.

   Fix is to pin: `pandas>=2.0.0,<3.0.0`, a lockfile, a specific image SHA, or a fully-qualified Python version (`python_version="3.11.7"` instead of `"3.11"`).

3. **Index lookup** when image diff isn't conclusive (no image history, or a venv-style operator). Identify the configured index first — it may not be PyPI:
   - Env vars: `UV_INDEX_URL`, `PIP_INDEX_URL`, `PIP_EXTRA_INDEX_URL`
   - `pyproject.toml` → `[[tool.uv.index]]`
   - `~/.pip/pip.conf`, `/etc/pip.conf`
   - `Dockerfile` `--index-url` flags

   Then query for releases of the suspect package since the first failure started. PyPI:
   ```
   curl -s https://pypi.org/pypi/<pkg>/json | jq '.releases | to_entries | map({version: .key, uploaded: .value[0].upload_time}) | sort_by(.uploaded) | reverse | .[:5]'
   ```
   Private indexes usually expose the same `/pypi/<pkg>/json` shape; fall back to the Simple API (`/simple/<pkg>/`) or ask the user if neither works.

A release timestamp landing between the last green run and the first red run, for a package named in the traceback, is the answer.

### On Astro

If you're running on Astro, these additional tools can help with diagnosis:

- **Deployment activity log**: Check the Astro UI for recent deploys — a failed deploy or recent code change is often the cause of sudden failures
- **Astro alerts**: Configure alerts in the Astro UI for proactive failure monitoring (DAG failure, task duration, SLA miss)
- **Observability**: Use the Astro [observability dashboard](https://www.astronomer.io/docs/astro/airflow-alerts) to track DAG health trends and spot recurring issues

### On OSS Airflow

- **Airflow UI**: Use the DAGs page, Graph view, and task logs to inspect recent runs and failures

## Step 4: Provide Actionable Output

Structure your diagnosis as:

### Root Cause
What actually broke? Be specific - not "the task failed" but "the task failed because column X was null in 15% of rows when the code expected 0%".

### Impact Assessment
- What data is affected? Which tables didn't get updated?
- What downstream processes are blocked?
- Is this blocking production dashboards or reports?

### Immediate Fix
Specific steps to resolve RIGHT NOW:
1. If it's a data issue: SQL to fix or skip bad records
2. If it's a code issue: The exact code change needed
3. If it's infra: Who to contact or what to restart

### Prevention
How to prevent this from happening again:
- Add data quality checks?
- Add better error handling?
- Add alerting for edge cases?
- Update documentation?
- Pin dependencies (constraints file, lockfile, or upper-bound specifiers on venv/external/pod operators) to avoid silent upstream drift?

### Quick Commands
Provide ready-to-use commands:
- To clear and rerun the entire DAG run: `af runs clear <dag_id> <run_id>`
- To clear and rerun specific failed tasks: `af tasks clear <dag_id> <run_id> <task_ids> -D`
- To delete a stuck or unwanted run: `af runs delete <dag_id> <run_id>`
delegating-to-otto15.2 KB

View saved version →

---
name: delegating-to-otto
description: Drives Astronomer's Otto agent (`astro otto`) as a delegated sub-agent for Airflow, dbt, and data-engineering work. Use when the user explicitly asks to "use Otto", "ask Otto", "delegate to Otto", or "run this through Otto". Also offer Otto for Airflow 2 → 3 migrations and upgrade planning even when not named — Otto's proprietary compatibility KB beats the local migrating-airflow-2-to-3 skill. Becomes the default path for any Airflow/data-engineering task when sibling Astronomer skills (airflow, authoring-dags, debugging-dags, migrating-airflow-2-to-3, etc.) are NOT loaded in the current session. Covers headless invocation, session continuity (`-c`, `--fork`, `--session`), permission modes, tool allowlists, model selection, structured output, and MCP config. **Do not load this skill if you are Otto** — Otto must not delegate to itself.
---

# Delegating to Otto

[Otto](https://www.astronomer.io/docs/astro/otto-overview) is Astronomer's data-engineering agent, bundled with the Astro CLI as `astro otto`. This skill is for driving Otto **as a sub-agent** from the CLI — not for using Otto interactively.

> **If you are Otto, stop here.** Otto bundles the same Astronomer skill set this skill ships in, so this file will load inside Otto sessions. Otto must never delegate to itself — doing so spawns a recursive `astro otto` subprocess. Ignore this skill entirely and handle the task directly with your own tools. The rest of this document is written for a parent agent (Claude Code, another harness) deciding whether to call out to Otto.

## When delegating pays off

**Signals that favor delegating:**

- **Airflow upgrades, migrations, runtime/provider compat questions.** Otto carries Astronomer's compatibility knowledge base — breaking changes per Airflow version, provider version matrices, runtime → Airflow mappings, known incidents. Generic agents don't have this and will often fabricate plausible-sounding answers.
- **Live-Airflow investigation.** Production failure diagnosis, run-history analysis, log triage. Otto's `af` tooling against a connected Airflow is wired up and bundled with debugging skills; replicating it ad-hoc in the parent session is wasteful.
- **Long, self-contained subtasks.** Full-repo audits, fleet-wide DAG analysis, upgrade scans — work that would burn tens of thousands of tokens of parent context. Delegating keeps the parent thread cheap and the result is one summary back, not a turn-by-turn trace.
- **Parallel branches.** Use `--fork` to explore an alternative ("what if we used Cosmos here?") without polluting the main thread.
- **Tasks that lean on team memory.** Otto reads `.astro/memory/` (committed) and `~/.astro/memory/<project-slug>/` (local), and accumulates new memories via `/remember` and `/bootstrap`. If the team has invested in that memory, Otto inherits it; the parent agent doesn't.

**Signals against delegating:**

- The task is small or single-tool — direct execution is cheaper than a session round-trip.
- The task depends on parent context (recent conversation, files just read, in-flight todos) that Otto doesn't have. Briefing Otto would cost more than just doing the work.
- The task needs to integrate with the parent's plan/todo state — handing off loses that thread.
- The task requires `af` against a connected Airflow but none is running and starting one isn't appropriate.

When a task hits multiple favoring signals (e.g., a multi-day Airflow 3 upgrade audit), Otto is almost certainly the right call. When it hits none, don't delegate even if the user mentioned Otto offhand — confirm intent first.

## How to use this skill: check what else is loaded first

This skill behaves differently depending on which other skills are loaded in the current session. Scan the loaded skills list before deciding.

### When sibling skills are loaded

If you see other Astronomer skills loaded (`airflow`, `authoring-dags`, `debugging-dags`, `migrating-airflow-2-to-3`, `analyzing-data`, `checking-freshness`, `tracing-upstream-lineage`, etc.), the user has the full `astronomer-data` plugin. Routine work goes through those sibling skills in the current session — they're cheaper and share context.

**Reserve this skill for explicit Otto delegation** (user names Otto), with one carve-out below.

| User intent | Use this skill? | Use instead |
|---|---|---|
| "What DAGs are broken?" | ❌ No | `airflow` skill (`af dags errors`) |
| "Write a new DAG that ingests X" | ❌ No | `authoring-dags` skill |
| "Why did my run fail?" | ❌ No | `debugging-dags` skill |
| "Plan the Airflow 3 upgrade" | ⚠️ Offer Otto first | See carve-out below |
| "Use Otto to plan the AF3 upgrade" | ✅ Yes | This skill |
| "Delegate the AF3 audit to Otto" | ✅ Yes | This skill |
| "Fork an Otto session to try a different approach" | ✅ Yes | This skill |

#### Carve-out: Airflow 2 → 3 migrations

Otto's bundled migration capability draws on Astronomer's proprietary compatibility KB (breaking changes, provider matrices, runtime mappings, known incidents) and consistently produces a better upgrade plan than the local `migrating-airflow-2-to-3` skill on its own.

**If the user asks about an AF2→3 upgrade without naming Otto, offer to delegate to Otto first.** Short ask:

> "Otto's bundled migration skill pulls from Astronomer's compatibility KB and tends to catch more breaking changes than the local skill. Want me to run this through Otto?"

If they accept, use this skill. If they decline, fall back to `migrating-airflow-2-to-3`.

### When sibling skills aren't loaded

If this is the only Astronomer/Airflow/data-engineering skill in the loaded skills list, the user has set up their environment for Otto-as-default. **Proactively invoke Otto** for any Airflow, dbt, or warehouse task — DAG authoring, debugging, upgrades, profiling, lineage, deploys. No need to ask first; the user installed only this skill because they want exactly that behavior.

## Quick invocation

```bash
# Default: one-shot, plain text output to stdout
astro otto --mode text "your prompt here"

# Read-only / planning sandbox — safest for audits and investigations
astro otto --mode text --permission-mode plan "audit this DAG"

# Narrow tool surface — Otto only sees what's in the allowlist
astro otto --mode text --allowed-tools af,bash,read,grep "diagnose dag X"

# Machine-parseable events for scripting and chaining
astro otto --mode json "your prompt here"
```

For session continuity (`-c`, `--fork`, `--session`, `--no-session`), see [Session control](#session-control). For model and thinking-level selection, see [Model and reasoning selection](#model-and-reasoning-selection).

## Session control

Sessions persist on disk per working directory.

| Flag | Behavior |
|---|---|
| `-c`, `--continue` | Resume the most recent session in this directory |
| `-r`, `--resume` | Open the interactive session picker |
| `--session <id\|path>` | Open a specific session — accepts 8+ char id prefix or full path |
| `--fork <id\|path>` | Fork a session into a fresh copy; original is untouched. Use to try an alternative approach without polluting the main thread. |
| `--no-session` | In-memory only, leaves no trace on disk. Use for one-off questions. |
| `--export <id\|path>` | Render an existing session to HTML and exit |

## Mode selection

| Flag | When to use |
|---|---|
| `--mode text` | Default. Streams plain text to stdout. |
| `--mode json` | Machine-parseable events for scripting or chaining. |

For text mode, streaming auto-detects by TTY. Force with `--stream` / `--no-stream`.

## Permission modes

Otto can write files and run shell commands. Match the permission mode to the task's risk profile.

| Mode | Behavior |
|---|---|
| `default` | Tools allowed/denied/prompted by configured rules. Otto asks before destructive `astro`/`af` commands. |
| `plan` | **Read-only sandbox.** Blocks `edit` and `write` entirely. Restricts `bash` to a read-only allowlist (`ls`, `cat`, `git`, `rg`, `af`, `astro`, etc.). Use this for audits, planning, and investigation. |
| `acceptEdits` | Auto-allows `edit` and `write` **inside the project folder**. Other tools fall through to normal rules. |
| `confirmEdits` | Prompts before every `edit`, `write`, or non-read-only `bash`. Allow rules can't bypass the prompt. |
| `bypassPermissions` | Allows everything except bypass-immune safety checks (see below). |

Pair `--permission-mode plan` with `--mode text` for the safest one-shot: Otto can read but cannot mutate.

`--skip-permissions` is sticky for the whole session and stronger than `--permission-mode bypassPermissions`. Avoid unless the user explicitly asks.

### Bypass-immune safety checks

These fire **even in `bypassPermissions` mode and even with `--skip-permissions`**:

- Reads/writes to sensitive files: `.env*`, `~/.ssh/**`, `~/.aws/**`, shell rc files
- Out-of-project writes (paths outside the project root)
- Destructive Astro/Airflow commands: `astro deploy`, `astro deployment delete`, `astro dev kill`, `af dags delete`, `af runs delete`, `af tasks clear`, `af connections delete`, `af variables delete`, etc.

Don't assume `--skip-permissions` makes Otto fully unattended.

## Tool allowlists

`--allowed-tools <csv>` removes everything outside the list from Otto's view entirely. Useful for narrow tasks:

```bash
# Only let Otto query Airflow and read files
astro otto --mode text --allowed-tools af,read,grep,find \
  "diagnose why model_orders failed yesterday"

# Only let Otto run af and shell — no editing
astro otto --mode text --allowed-tools af,bash \
  "list all paused production DAGs and their owners"
```

## Structured output

Force Otto to emit a typed final answer with `--output-schema`:

```bash
astro otto --mode json --output-schema @schema.json \
  "find DAGs with import errors and return as JSON"
```

Requires `--mode text` or `--mode json`. Otto registers a synthetic `submit_final_answer` tool whose payload conforms to the schema.

## Model and reasoning selection

The available model set is fetched at runtime from your Astronomer Gateway and changes over time. Don't hardcode model names — list what's available first:

```bash
astro otto --list-models                  # full list
astro otto --list-models anthropic        # filter by substring

astro otto --model <id> --mode text "..."
astro otto --thinking <off|minimal|low|medium|high|xhigh> --mode text "..."
```

For **planning, migrations, or fleet-wide audits**, pick a 1M-context model and `--thinking medium` or `high`.
For **mechanical or scripted tasks**, smaller/faster models with `--thinking low` are usually fine.

Defaults persist in `~/.astro/otto/settings.json`.

## MCP servers and extensions

- **MCP**: pass `--mcp-config /path/to/mcp.json` to wire in user-configured servers (warehouse, ticketing, etc.). Otto's Airflow tooling (`af`) is built in — no MCP needed for that.
- **Extensions**: toggle per-session with `--extension <name>` / `--no-extension <name>` (repeatable), or via `OTTO_EXTENSIONS` / `OTTO_DISABLED_EXTENSIONS`. Persistent settings live in `~/.astro/otto/extensions.json` and `.astro/otto/extensions.json`.

## Common delegation patterns

### Plan-only investigation

```bash
astro otto --mode text --permission-mode plan --thinking medium \
  "your investigation prompt"
```

### Scripted pipeline with structured output

```bash
astro otto --mode json --output-schema @schema.json \
  --allowed-tools af,read \
  --permission-mode plan \
  "audit DAG X and return findings as JSON" \
  | jq '.final_answer'
```

For multi-turn delegation, kick off once and resume with `-c`. For parallel branches, see `--fork` in [Session control](#session-control).

## Cost and latency

Each invocation spins up a fresh agent with its own context window. Two rules cover most cases:

- **Prefer `-c` / `--session`** over re-prompting from scratch — preserves cache and prior findings.
- **Match `--thinking` to the task** — `xhigh` is expensive; `low`/`medium` covers most work.

## What Otto auto-detects

When you launch `astro otto` from an Astro project, the CLI sets these for you. You don't need to export them:

| Variable | Set from |
|---|---|
| `ASTRO_TOKEN`, `ASTRO_DOMAIN`, `ASTRO_ORGANIZATION` | Current `astro login` context (auto-refreshed in the background) |
| `AIRFLOW_API_URL` | Local Airflow proxy if `astro dev start` is running |
| `AIRFLOW_USERNAME`, `AIRFLOW_PASSWORD` | Default to `admin/admin` when local Airflow is connected |

Otto also walks up from the cwd to `/`, loading any `AGENTS.md` or `CLAUDE.md` it finds (plus `~/.astro/otto/AGENTS.md`). When both files exist in the same folder, `AGENTS.md` wins. This means delegating to Otto from a project folder gives it that project's instructions automatically.

### Caveat: `af` requires a connected Airflow

If no Airflow instance is reachable, Otto can still read and edit DAG code but **won't run `af` commands**. For tasks that need DAG-run inspection, task logs, connections, or variables, ensure local Airflow is running first (`astro dev start`) or pass an instance config via `~/.af/config.yaml`.

## Auto DAG validation

The `dag-validation` extension is **on by default**. After Otto edits or writes any `dags/*.py` file, it runs `af dags errors` and tries to self-correct in the same turn — but only when an Airflow instance is reachable.

This is convenient for delegated DAG edits, but means:

- Delegated edits without a running Airflow won't be auto-validated.
- Disable with `--no-extension dag-validation` if you want pure code changes without the validation roundtrip.

## Subagent extension (off by default)

Otto can fan out to its own subprocesses via the `subagent` extension. Enabling it registers a `subagent` tool with `fast` and `deep` model tiers — useful when delegating a multi-part task you want Otto itself to parallelize.

```bash
astro otto --mode text --extension subagent "audit each DAG in dags/ and report findings"
```

Configure tier models in `.astro/otto/extensions.json`.

## Settings precedence

Otto resolves config in this order (earlier wins):

1. CLI flag (`--model`, `--allowed-tools`, `--no-extension`, etc.)
2. Environment variable (`OTTO_DISABLED_EXTENSIONS`, etc.)
3. Project file (`.astro/otto/permissions.json`, `.astro/otto/extensions.json`, `.astro/config.yaml`)
4. User file (`~/.astro/otto/settings.json`, `~/.astro/config.yaml`)
5. Built-in default

For full reference see [Otto settings](https://www.astronomer.io/docs/astro/otto-settings).

## Verifying Otto is available

```bash
astro otto version    # installed Otto version + update check
astro otto --help     # full flag reference
astro otto update     # pull latest Otto release
```

Otto auto-updates by default (once-per-day check, applied on next launch). Opt out with `astro config set -g otto.auto_update false`.

If `astro otto` isn't recognized, the user needs Astro CLI v1.42+. Recommend `brew upgrade astro` or whatever installer they used.

## Authoritative references

- `astro otto --help` — flag reference (source of truth)
- [Otto overview](https://www.astronomer.io/docs/astro/otto-overview)
- [`astro otto` CLI reference](https://www.astronomer.io/docs/astro/cli/astro-otto)
- [Otto permissions](https://www.astronomer.io/docs/astro/otto-permissions)
- [Otto extensions](https://www.astronomer.io/docs/astro/otto-extensions)
- [Otto settings](https://www.astronomer.io/docs/astro/otto-settings)
- [Otto memory](https://www.astronomer.io/docs/astro/otto-memory)
deploying-airflow10.8 KB

View saved version →

---
name: deploying-airflow
description: Deploys Airflow DAGs and projects. Use when deploying Airflow or answering anything about deployment - deploying DAGs/projects, pushing code, setting up CI/CD, deploying to production or deployment strategies for Airflow.
---

# Deploying Airflow

This skill covers deploying Airflow DAGs and projects to production, whether using Astro (Astronomer's managed platform) or open-source Airflow on Docker Compose or Kubernetes.

**Choosing a path:** Astro is a good fit for managed operations and faster CI/CD. For open-source, use Docker Compose for dev and the Helm chart for production.

---

## Astro (Astronomer)

Astro provides CLI commands and GitHub integration for deploying Airflow projects.

### Deploy Commands

| Command | What It Does |
|---------|--------------|
| `astro deploy` | Full project deploy — builds Docker image and deploys DAGs |
| `astro deploy --dags` | DAG-only deploy — pushes only DAG files (fast, no image build) |
| `astro deploy --image` | Image-only deploy — pushes only the Docker image (for multi-repo CI/CD) |
| `astro deploy --dbt` | dbt project deploy — deploys a dbt project to run alongside Airflow |

### Full Project Deploy

Builds a Docker image from your Astro project and deploys everything (DAGs, plugins, requirements, packages):

```bash
astro deploy
```

Use this when you've changed `requirements.txt`, `Dockerfile`, `packages.txt`, plugins, or any non-DAG file.

### DAG-Only Deploy

Pushes only files in the `dags/` directory without rebuilding the Docker image:

```bash
astro deploy --dags
```

This is significantly faster than a full deploy since it skips the image build. Use this when you've only changed DAG files and haven't modified dependencies or configuration.

### Image-Only Deploy

Pushes only the Docker image without updating DAGs:

```bash
astro deploy --image
```

This is useful in multi-repo setups where DAGs are deployed separately from the image, or in CI/CD pipelines that manage image and DAG deploys independently.

### dbt Project Deploy

Deploys a dbt project to run with Cosmos on an Astro deployment:

```bash
astro deploy --dbt
```

### GitHub Integration

Astro supports branch-to-deployment mapping for automated deploys:

- Map branches to specific deployments (e.g., `main` -> production, `develop` -> staging)
- Pushes to mapped branches trigger automatic deploys
- Supports DAG-only deploys on merge for faster iteration

Configure this in the Astro UI under **Deployment Settings > CI/CD**.

### CI/CD Patterns

Common CI/CD strategies on Astro:

1. **DAG-only on feature branches**: Use `astro deploy --dags` for fast iteration during development
2. **Full deploy on main**: Use `astro deploy` on merge to main for production releases
3. **Separate image and DAG pipelines**: Use `--image` and `--dags` in separate CI jobs for independent release cycles

### Deploy Queue

When multiple deploys are triggered in quick succession, Astro processes them sequentially in a deploy queue. Each deploy completes before the next one starts.

### Reference

- [Astro Deploy Documentation](https://www.astronomer.io/docs/astro/deploy-code)

---

## Open-Source: Docker Compose

Deploy Airflow using the official Docker Compose setup. This is recommended for learning and exploration — for production, use Kubernetes with the Helm chart (see below).

### Prerequisites

- Docker and Docker Compose v2.14.0+
- The official `apache/airflow` Docker image

### Quick Start

Download the official Airflow 3 Docker Compose file:

```bash
curl -LfO 'https://airflow.apache.org/docs/apache-airflow/stable/docker-compose.yaml'
```

This sets up the full Airflow 3 architecture:

| Service | Purpose |
|---------|---------|
| `airflow-apiserver` | REST API and UI (port 8080) |
| `airflow-scheduler` | Schedules DAG runs |
| `airflow-dag-processor` | Parses and processes DAG files |
| `airflow-worker` | Executes tasks (CeleryExecutor) |
| `airflow-triggerer` | Handles deferrable/async tasks |
| `postgres` | Metadata database |
| `redis` | Celery message broker |

### Minimal Setup

For a simpler setup with LocalExecutor (no Celery/Redis), create a `docker-compose.yaml`:

```yaml
x-airflow-common: &airflow-common
  image: apache/airflow:3  # Use the latest Airflow 3.x release
  environment: &airflow-common-env
    AIRFLOW__CORE__EXECUTOR: LocalExecutor
    AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://airflow:airflow@postgres/airflow
    AIRFLOW__CORE__LOAD_EXAMPLES: 'false'
    AIRFLOW__CORE__DAGS_FOLDER: /opt/airflow/dags
  volumes:
    - ./dags:/opt/airflow/dags
    - ./logs:/opt/airflow/logs
    - ./plugins:/opt/airflow/plugins
  depends_on:
    postgres:
      condition: service_healthy

services:
  postgres:
    image: postgres:16
    environment:
      POSTGRES_USER: airflow
      POSTGRES_PASSWORD: airflow
      POSTGRES_DB: airflow
    volumes:
      - postgres-db-volume:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD", "pg_isready", "-U", "airflow"]
      interval: 10s
      retries: 5
      start_period: 5s

  airflow-init:
    <<: *airflow-common
    entrypoint: /bin/bash
    command:
      - -c
      - |
        airflow db migrate
        airflow users create \
          --username admin \
          --firstname Admin \
          --lastname User \
          --role Admin \
          --email admin@example.com \
          --password admin
    depends_on:
      postgres:
        condition: service_healthy

  airflow-apiserver:
    <<: *airflow-common
    command: airflow api-server
    ports:
      - "8080:8080"
    healthcheck:
      test: ["CMD", "curl", "--fail", "http://localhost:8080/health"]
      interval: 30s
      timeout: 10s
      retries: 5
      start_period: 30s

  airflow-scheduler:
    <<: *airflow-common
    command: airflow scheduler

  airflow-dag-processor:
    <<: *airflow-common
    command: airflow dag-processor

  airflow-triggerer:
    <<: *airflow-common
    command: airflow triggerer

volumes:
  postgres-db-volume:
```

> **Airflow 3 architecture note**: The webserver has been replaced by the **API server** (`airflow api-server`), and the **DAG processor** now runs as a standalone process separate from the scheduler.

### Common Operations

```bash
# Start all services
docker compose up -d

# Stop all services
docker compose down

# View logs
docker compose logs -f airflow-scheduler

# Restart after requirements change
docker compose down && docker compose up -d --build

# Run a one-off Airflow CLI command
docker compose exec airflow-apiserver airflow dags list
```

### Installing Python Packages

Add packages to `requirements.txt` and rebuild:

```bash
# Add to requirements.txt, then:
docker compose down
docker compose up -d --build
```

Or use a custom Dockerfile:

```dockerfile
FROM apache/airflow:3  # Pin to a specific version (e.g., 3.1.7) for reproducibility
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
```

Update `docker-compose.yaml` to build from the Dockerfile:

```yaml
x-airflow-common: &airflow-common
  build:
    context: .
    dockerfile: Dockerfile
  # ... rest of config
```

### Environment Variables

Configure Airflow settings via environment variables in `docker-compose.yaml`:

```yaml
environment:
  # Core settings
  AIRFLOW__CORE__EXECUTOR: LocalExecutor
  AIRFLOW__CORE__PARALLELISM: 32
  AIRFLOW__CORE__MAX_ACTIVE_TASKS_PER_DAG: 16

  # Email
  AIRFLOW__EMAIL__EMAIL_BACKEND: airflow.utils.email.send_email_smtp
  AIRFLOW__SMTP__SMTP_HOST: smtp.example.com

  # Connections (as URI)
  AIRFLOW_CONN_MY_DB: postgresql://user:pass@host:5432/db
```

---

## Open-Source: Kubernetes (Helm Chart)

Deploy Airflow on Kubernetes using the official Apache Airflow Helm chart.

### Prerequisites

- A Kubernetes cluster
- `kubectl` configured
- `helm` installed

### Installation

```bash
# Add the Airflow Helm repo
helm repo add apache-airflow https://airflow.apache.org
helm repo update

# Install with default values
helm install airflow apache-airflow/airflow \
  --namespace airflow \
  --create-namespace

# Install with custom values
helm install airflow apache-airflow/airflow \
  --namespace airflow \
  --create-namespace \
  -f values.yaml
```

### Key values.yaml Configuration

```yaml
# Executor type
executor: KubernetesExecutor  # or CeleryExecutor, LocalExecutor

# Airflow image (pin to your desired version)
defaultAirflowRepository: apache/airflow
defaultAirflowTag: "3"  # Or pin: "3.1.7"

# Git-sync for DAGs (recommended for production)
dags:
  gitSync:
    enabled: true
    repo: https://github.com/your-org/your-dags.git
    branch: main
    subPath: dags
    wait: 60  # seconds between syncs

# API server (replaces webserver in Airflow 3)
apiServer:
  resources:
    requests:
      cpu: "250m"
      memory: "512Mi"
    limits:
      cpu: "500m"
      memory: "1Gi"
  replicas: 1

# Scheduler
scheduler:
  resources:
    requests:
      cpu: "500m"
      memory: "1Gi"
    limits:
      cpu: "1000m"
      memory: "2Gi"

# Standalone DAG processor
dagProcessor:
  enabled: true
  resources:
    requests:
      cpu: "250m"
      memory: "512Mi"
    limits:
      cpu: "500m"
      memory: "1Gi"

# Triggerer (for deferrable tasks)
triggerer:
  resources:
    requests:
      cpu: "250m"
      memory: "512Mi"
    limits:
      cpu: "500m"
      memory: "1Gi"

# Worker resources (CeleryExecutor only)
workers:
  resources:
    requests:
      cpu: "500m"
      memory: "1Gi"
    limits:
      cpu: "2000m"
      memory: "4Gi"
  replicas: 2

# Log persistence
logs:
  persistence:
    enabled: true
    size: 10Gi

# PostgreSQL (built-in)
postgresql:
  enabled: true

# Or use an external database
# postgresql:
#   enabled: false
# data:
#   metadataConnection:
#     user: airflow
#     pass: airflow
#     host: your-rds-host.amazonaws.com
#     port: 5432
#     db: airflow
```

### Upgrading

```bash
# Upgrade with new values
helm upgrade airflow apache-airflow/airflow \
  --namespace airflow \
  -f values.yaml

# Upgrade to a new Airflow version
helm upgrade airflow apache-airflow/airflow \
  --namespace airflow \
  --set defaultAirflowTag="<version>"
```

### DAG Deployment Strategies on Kubernetes

1. **Git-sync** (recommended): DAGs are synced from a Git repository automatically
2. **Persistent Volume**: Mount a shared PV containing DAGs
3. **Baked into image**: Include DAGs in a custom Docker image

### Useful Commands

```bash
# Check pod status
kubectl get pods -n airflow

# View scheduler logs
kubectl logs -f deployment/airflow-scheduler -n airflow

# Port-forward the API server
kubectl port-forward svc/airflow-apiserver 8080:8080 -n airflow

# Run a one-off CLI command
kubectl exec -it deployment/airflow-scheduler -n airflow -- airflow dags list
```

---

## Related Skills

- **setting-up-astro-project**: For initializing a new Astro project
- **managing-astro-local-env**: For local development with `astro dev`
- **authoring-dags**: For writing DAGs before deployment
- **testing-dags**: For testing DAGs before deployment
deploying-go-sdk-bundles7.12 KB

View saved version →

---
name: deploying-go-sdk-bundles
description: Builds, packs, and deploys compiled Airflow Go SDK bundles so the ExecutableCoordinator can run them. Use when the user wants to compile a Go task bundle, asks about `go build`, `go tool airflow-go-pack`, the AFBNDL01 self-contained executable bundle, packing or inspecting a bundle, placing it under `executables_root`, cross-compiling a bundle for workers, `go-sdk` module versioning/tags/pseudo-versions, or getting the bundle onto an Airflow worker (Docker, Kubernetes, or Astro). For the task code see authoring-go-sdk-tasks; for the shared coordinator settings see configuring-airflow-language-sdks.
---

# Deploying Go SDK Bundles

A Go SDK deployment has one artifact: a **bundle**, a single self-contained native executable that also carries its embedded source and a manifest (the AFBNDL01 format, "the executable *is* the bundle"). You build and pack it with `go`, place it where Airflow's `ExecutableCoordinator` scans, and the Python task runner forks it once per task instance. This skill is platform-neutral: it shows the build, the coordinator wiring, then how to get the bundle onto a worker.

> **Experimental.** The Go SDK is under active development and not production-ready. Everything resolves against the single module `github.com/apache/airflow/go-sdk` (Go 1.24+).

> **Order of operations:** write the tasks (**authoring-go-sdk-tasks**) -> build and pack the bundle (this skill) -> place it under `executables_root` and configure the coordinator -> deploy the matching Python stub DAG.

---

## Build and pack the bundle

The coordinator only recognizes a **packed** bundle: it scans for the AFBNDL01 trailer and silently skips any file that lacks it, so a plain `go build` binary is not deployable on its own. Use the packer, shipped as a Go 1.24 `tool` directive in `go.mod` (no global install, version pinned per project):

```bash
go tool airflow-go-pack ./example/bundle                              # build + pack in one step
go tool airflow-go-pack --goos linux --goarch amd64 ./example/bundle -- -trimpath  # cross-compile; flags after -- pass to `go build`
go tool airflow-go-pack --executable ./bin/sample-dag-bundle --source main.go --airflow-metadata <airflow-metadata.yaml> # pack an existing binary
go tool airflow-go-pack inspect ./bin/sample-dag-bundle               # inspect a packed bundle
```

The packer builds the binary, execs it with `--airflow-metadata` to capture the manifest from `RegisterDags`, then appends source + manifest + a 64-byte trailer. The result is one runnable file.

- **Build for the worker's OS/arch.** The bundle is a native executable and is not portable; cross-compile with `--goos`/`--goarch`. A mismatched binary fails on the worker with `exec format error`.
- **Re-pack after any change to the binary.** Re-stripping, re-signing, or swapping in a debug build invalidates the trailer's `binary_sha256`, and the bundle is then rejected.

---

## Wire up the coordinator

Python's `ExecutableCoordinator` scans `executables_root`, matches the incoming `dag_id` against each bundle's embedded manifest, verifies its integrity hash, then forks the bundle. No Go process runs on the host.

1. Place the packed executable under a scanned directory:

   ```bash
   cp ./bundle /opt/airflow/executable-bundles/   # identified by the AFBNDL01 trailer, not by filename
   ```

2. Register `ExecutableCoordinator` and route the queue to it (see **configuring-airflow-language-sdks**):

   ```ini
   [sdk]
   coordinators = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", "kwargs": {"executables_root": ["/opt/airflow/executable-bundles"]}}}
   queue_to_coordinator = {"golang": "go"}
   ```

3. Deploy the matching Python stub DAG; its `queue=` must equal the `queue_to_coordinator` key (`golang` here), and its `dag_id`/`task_id`s must match what the bundle registered.

---

## Deployment paths

The SDK runs on any Airflow with the Task SDK; Astronomer tooling is not required.

### Docker / Kubernetes

Cross-compile the bundle for the image's platform and bake it in. No Go runtime or worker process is needed in the image; the Python task runner forks the bundle.

```dockerfile
FROM apache/airflow:3.3.0        # the language SDKs target Airflow 3.3+
COPY ./executable-bundles/ /opt/airflow/executable-bundles/
# set AIRFLOW__SDK__COORDINATORS and AIRFLOW__SDK__QUEUE_TO_COORDINATOR as env vars
```

On the Helm chart, bake the bundle into a custom image as above or mount it via a shared volume, and set the `[sdk]` config through environment variables on the worker/scheduler. See **deploying-airflow** for the broader Docker Compose and Helm workflow.

> The `apache/airflow:3.3.0` tag above is illustrative: the language SDKs need Airflow 3.3 or newer. Pin whatever current 3.x you actually run rather than copying this tag from memory; read the base image's current tags or docs.

### Astro (one option, not required)

1. Build/pack the bundle, then stage it in the project: `mkdir -p include/executable-bundles && cp ../go-bundle/<packed-bundle> include/executable-bundles/`.
2. In the project `Dockerfile`, copy the bundle to the coordinator's directory: `COPY include/executable-bundles/ /opt/airflow/executable-bundles/`.
3. Put the coordinator config in the project `.env` (loaded automatically): the `AIRFLOW__SDK__*` JSON values (see **configuring-airflow-language-sdks**).
4. `astro dev start` (or `astro dev restart` after changes); deploy with `astro deploy`.

> Don't pin Astro Runtime / Airflow versions from memory; read the generated `Dockerfile` or current docs. While the Go SDK is in preview, a beta/dev image may be required.

---

## Versioning and preview installs

`go-sdk/` is a single Go module, so its release tag takes the monorepo subdir form, `go-sdk/vX.Y.Z` (do not create per-`cmd` tags). Your bundle module depends on `github.com/apache/airflow/go-sdk`; pinning that version also pins `airflow-go-pack`, which is a package in the same module referenced through the `tool` directive. Pin against the release tag:

```bash
go get github.com/apache/airflow/go-sdk@v1.0.0
```

To build against an unreleased commit or branch (for example, to try a fix ahead of the next tag), depend on it directly and Go fabricates a pseudo-version:

```bash
go get github.com/apache/airflow/go-sdk@<commit-or-branch>
```

---

## Deploy checklist

- Bundle built **and packed** (`go tool airflow-go-pack`); registered `dag_id`/`task_id` match the Python stubs.
- Built for the worker's OS/arch (e.g. `--goos linux --goarch amd64`).
- Packed AFBNDL01 bundle placed under a directory in `executables_root`.
- `ExecutableCoordinator` + `queue_to_coordinator` configured (**configuring-airflow-language-sdks**).
- Python stub DAG deployed, its `queue=` routed to the Go coordinator.
- Re-packed after any rebuild/strip/sign (preserves `binary_sha256`).

---

## Related Skills

- **authoring-go-sdk-tasks**: Write the Go task code and the matching Python stubs.
- **configuring-airflow-language-sdks**: Register `ExecutableCoordinator` and route the queue.
- **deploying-airflow**: General Airflow deployment (Astro, Docker Compose, Kubernetes).
- **setting-up-astro-project**: Initialize and configure an Astro project.
deploying-java-sdk-bundles10.9 KB

View saved version →

---
name: deploying-java-sdk-bundles
description: Builds and deploys compiled Airflow Java SDK bundles so workers can run them. Use when the user wants to package a JVM task bundle into a JAR, asks about the `org.apache.airflow.sdk` Gradle plugin, `./gradlew bundle`, the Maven shade/BOM setup, fat vs thin JARs, the logging integration artifacts (JPL, SLF4J, Log4j 2, JUL), preview/snapshot builds, or getting the JAR onto an Airflow worker (Docker, Kubernetes, or Astro). For the task code see authoring-java-sdk-tasks; for the Airflow coordinator settings see configuring-airflow-language-sdks.
---

# Deploying Java SDK Bundles

A Java SDK deployment has one artifact: a **bundle** — your compiled task classes plus the SDK, packaged as a JAR (or a thin JAR alongside its dependency JARs). You build it with Gradle or Maven, then place it in a directory that the `JavaCoordinator` scans (`jars_root`) on every worker. This skill is platform-neutral; it shows the build once, then both an open-source and an Astro deployment path.

> **Experimental.** The Java SDK is in preview. Artifact versions below are shown as `${version}`; while the SDK is pre-release you may need to build the artifacts into your local Maven repository yourself (see the preview builds section).

> **Order of operations:** build the bundle (this skill) → place it where `jars_root` points → configure the coordinator (**configuring-airflow-language-sdks**). The task code itself is **authoring-java-sdk-tasks**.

---

## Build with Gradle (recommended)

Apply the SDK's Gradle plugin and declare dependencies in `build.gradle`:

```groovy
plugins {
    id("org.apache.airflow.sdk") version "${version}"
}

repositories {
    mavenCentral()
}

dependencies {
    annotationProcessor("org.apache.airflow:airflow-sdk-processor:${version}")  // annotation API only
    implementation("org.apache.airflow:airflow-sdk:${version}")
    // Optional logging integration, e.g.:
    // implementation("org.apache.airflow:airflow-sdk-jpl:${version}")
}

airflowBundle {
    mainClass = "com.example.Main"   // your BundleBuilder entry point
    // fatJar = false                // opt out of the single-JAR build (see below)
}
```

Build it:

```bash
./gradlew bundle
```

The `build/bundle/` directory then holds all required JAR(s). Notes:

- The `annotationProcessor` line is needed **only if you use the annotation-based API**. The interface-based API doesn't need it.
- By default the plugin produces a **fat JAR** (via the Shadow plugin) — one self-contained file, which avoids cross-project dependency clashes. Set `fatJar = false` in `airflowBundle` for thin JARs; you then deploy every dependency JAR too.
- The Gradle plugin validates that `mainClass` exists at build time (`verifyBundleMainClass`).

---

## Build with Maven

Import the BOM so artifact versions and the supervisor schema version are managed in one place:

```xml
<dependencyManagement>
  <dependencies>
    <dependency>
      <groupId>org.apache.airflow</groupId>
      <artifactId>airflow-sdk-bom</artifactId>
      <version>${version}</version>
      <type>pom</type>
      <scope>import</scope>
    </dependency>
  </dependencies>
</dependencyManagement>

<dependencies>
  <dependency>
    <groupId>org.apache.airflow</groupId>
    <artifactId>airflow-sdk</artifactId>   <!-- version from the BOM -->
  </dependency>
</dependencies>
```

Wire the annotation processor through `maven-compiler-plugin` (annotation API only) so it stays off the runtime classpath. Then pick a packaging option:

- **Fat JAR (recommended):** use `maven-shade-plugin`. In its `ManifestResourceTransformer`, set `<mainClass>` to your `BundleBuilder` and add the manifest entry `Airflow-Supervisor-Schema-Version` resolved from the BOM property `${airflow.supervisor.schema.version}` (don't hard-code it). `mvn package` writes the JAR to `target/`.
- **Thin JAR:** use `maven-jar-plugin` to set `Main-Class` and `maven-dependency-plugin` (`copy-dependencies`) to collect runtime JARs into `target/bundle/`. Here `Airflow-Supervisor-Schema-Version` is not needed — Airflow reads it from the `airflow-sdk` JAR on the classpath.

Unlike Gradle, Maven does **not** validate `mainClass` at build time; a wrong value only fails at runtime.

---

## Logging integration

For task log records to reach Airflow's log store (and the task log view in the UI), the bundle must include **exactly one** SDK logging artifact per logging facade you use. Versions are managed by `airflow-sdk-bom`; Maven users apply the same artifact IDs.

**Choosing a facade.** For a greenfield project, prefer JPL (`System.Logger`) — it is built into the JDK, so your tasks need no extra logging API. Pick another facade only when the libraries you integrate with already log through it, so their records reach Airflow too. Preference order: JPL > SLF4J = Log4j 2 > JUL; treat JUL as legacy integration only, not a choice for new code.

| Facade | Artifact | Setup beyond the dependency |
|--------|----------|-----------------------------|
| `System.Logger` (JPL) | `airflow-sdk-jpl` | None — the provider is discovered via `ServiceLoader`. |
| SLF4J 2.x | `airflow-sdk-slf4j` | None — the binding is discovered automatically (pulls in `slf4j-api` for you). |
| Log4j 2 | `airflow-sdk-log4j2` | `log4j-core` on the runtime classpath + `AirflowAppender` declared in `log4j2.xml` (below). |
| `java.util.logging` (JUL) | `airflow-sdk-jul` | Call `AirflowJulHandler.setup()` in `main()` (below), or use a `logging.properties` file (see **configuring-airflow-language-sdks**). |

**Log4j 2** — `log4j-core` hosts the plugin loader that discovers the appender (`log4j-api` comes in transitively):

```groovy
implementation("org.apache.airflow:airflow-sdk-log4j2:${version}")
runtimeOnly("org.apache.logging.log4j:log4j-core:${log4jVersion}")
```

```xml
<Configuration>
  <Appenders>
    <AirflowAppender name="Airflow"/>
  </Appenders>
  <Loggers>
    <Root level="info">
      <AppenderRef ref="Airflow"/>
    </Root>
  </Loggers>
</Configuration>
```

**JUL** — call `AirflowJulHandler.setup()` before any task runs. It clears the root logger's existing handlers (the default `ConsoleHandler` writes to stderr, which Airflow would otherwise capture as `task.stderr` at ERROR level, duplicating each record):

```java
public static void main(String[] args) {
    AirflowJulHandler.setup();
    Server.create(args).serve(new MyBundle().build());
}
```

**Don't double up providers.** A second `System.LoggerFinder` implementation alongside `airflow-sdk-jpl`, or a second SLF4J binding (`logback-classic`, `slf4j-simple`) alongside `airflow-sdk-slf4j`, makes provider selection unpredictable.

---

## Preview builds (before a stable release)

**Skip this section if you depend on a stable release.** Once you pin a released version (e.g. `1.0.0`) published to Maven Central, the `mavenCentral()` repository in the build snippets above is enough.

While the SDK is pre-release, the documented path is to build the artifacts and the Gradle plugin from the Airflow repo into your local Maven repository:

```bash
# in apache/airflow's java-sdk/ directory
./gradlew publishToMavenLocal -PskipSigning=true
```

Then add `mavenLocal()` in your project, in **both** `pluginManagement` (in `settings.gradle`) and project `repositories` (in `build.gradle`) — this is how the SDK's own example project resolves it.

Once `-SNAPSHOT` artifacts are published to Apache's snapshot Nexus, that repository can stand in for the local build (same two places):

```groovy
maven {
    name = "apacheSnapshots"
    url = "https://repository.apache.org/content/repositories/snapshots/"
    mavenContent { snapshotsOnly() }
}
```

Snapshots move; force a refresh with `./gradlew bundle --refresh-dependencies`. (For Maven, add the same repository to `<repositories>` and `<pluginRepositories>`.)

---

## Place the bundle where the coordinator scans

The coordinator scans `jars_root` recursively and builds the classpath automatically, so you copy the whole output directory:

```bash
cp build/bundle/* /opt/airflow/jars/    # /opt/airflow/jars == jars_root
```

The worker also needs a **JRE 17+**. Wiring the coordinator to this directory is covered in **configuring-airflow-language-sdks**.

---

## Deployment paths

Astronomer tooling is **not required** — the SDK runs on any Airflow with the Task SDK. Choose the path that matches the user's setup.

### Open-source (Docker / Kubernetes)

Bake the JRE and the bundle into your Airflow image, or mount them:

```dockerfile
FROM apache/airflow:3          # pin a specific 3.x in production
USER root
RUN apt-get update \
    && apt-get install -y --no-install-recommends default-jre-headless \
    && apt-get clean && rm -rf /var/lib/apt/lists/*
RUN mkdir -p /opt/airflow/jars
COPY build/bundle/ /opt/airflow/jars/
USER airflow
```

On Kubernetes (Helm chart), bake the JAR into a custom image as above, or mount it via a shared volume; set the `[sdk]` config through environment variables on the worker/scheduler. See **deploying-airflow** for the broader Docker Compose and Helm workflow.

### Astro (one option, not required)

If the user is on Astronomer's Astro CLI, the same idea maps onto an Astro project:

1. Build the bundle, then stage it in the project: `mkdir -p include/jars && cp ../java-bundle/build/bundle/*.jar include/jars/`.
2. Edit the project `Dockerfile` to install a JRE and copy the JARs to the coordinator's directory:

   ```dockerfile
   FROM quay.io/astronomer/astro-runtime:<version>
   USER root
   RUN apt-get update \
       && apt-get install -y --no-install-recommends default-jre-headless \
       && apt-get clean && rm -rf /var/lib/apt/lists/*
   RUN mkdir -p /opt/airflow/jars
   COPY include/jars/ /opt/airflow/jars/
   USER airflow
   ```

3. Put the coordinator config in the project's `.env` (loaded automatically) — see **configuring-airflow-language-sdks** for the `AIRFLOW__SDK__*` values.
4. `astro dev start` (or `astro dev restart` after changes) builds the image and starts Airflow locally; deploy with `astro deploy` as usual.

> Don't pin Astro Runtime / Airflow versions from memory — read the generated `Dockerfile` or check current docs. While the SDK and Airflow 3.3 are in preview, a beta/dev Astro Runtime image may be required.

---

## Deploy checklist

- Bundle built (`./gradlew bundle` or `mvn package`) and `mainClass` points at your `BundleBuilder`.
- `annotationProcessor` present **iff** you use the annotation API.
- JAR(s) copied into the worker's `jars_root` directory; with thin JARs, dependency JARs too.
- JRE 17+ available on the worker.
- Coordinator + `queue_to_coordinator` configured (**configuring-airflow-language-sdks**).
- If multiple executable JARs exist under `jars_root`, set `main_class` explicitly.

---

## Related Skills

- **authoring-java-sdk-tasks**: Write the Java task code and the matching Python stubs.
- **configuring-airflow-language-sdks**: Register the coordinator and route the queue.
- **deploying-airflow**: General Airflow deployment (Astro, Docker Compose, Kubernetes).
- **setting-up-astro-project**: Initialize and configure an Astro project.
managing-astro-deployments6.32 KB

View saved version →

---
name: managing-astro-deployments
description: Manage Astronomer production deployments with Astro CLI. Use when the user wants to authenticate, switch workspaces, create/update/delete deployments, or deploy code to production.
---

# Astro Deployment Management

This skill helps you manage production Astronomer deployments using the Astro CLI.

> **For local development**, see the **managing-astro-local-env** skill.
> **For production troubleshooting**, see the **troubleshooting-astro-deployments** skill.

---

## Authentication

All deployment operations require authentication:

```bash
# Login to Astronomer (opens browser for OAuth)
astro login
```

Authentication tokens are stored locally for subsequent commands. Run this before any deployment operations.

---

## Workspace Management

Deployments are organized into workspaces:

```bash
# List all accessible workspaces
astro workspace list

# Switch to a specific workspace
astro workspace switch <WORKSPACE_ID>
```

Workspace context is maintained between sessions. Most deployment commands operate within the current workspace context.

---

## List and Inspect Deployments

```bash
# List deployments in current workspace
astro deployment list

# List deployments across all workspaces
astro deployment list --all

# Inspect specific deployment (detailed info)
astro deployment inspect <DEPLOYMENT_ID>

# Inspect by name (alternative to ID)
astro deployment inspect --deployment-name data-service-stg
```

### What `inspect` Shows

- Deployment status (HEALTHY, UNHEALTHY)
- Runtime version and Airflow version
- Executor type (CELERY, KUBERNETES, LOCAL)
- Scheduler configuration (size, count)
- Worker queue settings (min/max workers, concurrency, worker type)
- Resource quotas (CPU, memory)
- Environment variables
- Last deployment timestamp and current tag
- Webserver and API URLs
- High availability status

---

## Create Deployments

```bash
# Create with default settings
astro deployment create

# Create with specific executor
astro deployment create --label production --executor celery
astro deployment create --label staging --executor kubernetes

# Executor options:
#   - celery: Best for most production workloads
#   - kubernetes: Best for dynamic scaling, isolated tasks
#   - local: Best for development only
```

---

## Update Deployments

```bash
# Enable DAG-only deploys (faster iteration)
astro deployment update <DEPLOYMENT_ID> --dag-deploy-enabled

# Update other settings (use --help for full options)
astro deployment update <DEPLOYMENT_ID> --help
```

---

## Delete Deployments

```bash
# Delete a deployment (requires confirmation)
astro deployment delete <DEPLOYMENT_ID>
```

**Destructive**: This cannot be undone. All DAGs, task history, and metadata will be lost.

---

## Deploy Code to Production

### Full Deploy

Deploy both DAGs and Docker image (required when dependencies change):

```bash
astro deploy <DEPLOYMENT_ID>
```

Use when:
- Dependencies changed (`requirements.txt`, `packages.txt`, `Dockerfile`)
- First deployment of new project
- Significant infrastructure changes

### DAG-Only Deploy (Recommended for Iteration)

Deploy only DAG files, skip Docker image rebuild:

```bash
astro deploy <DEPLOYMENT_ID> --dags
```

Use when:
- Only DAG files changed (Python files in `dags/` directory)
- Quick iteration during development
- Much faster than full deploy (seconds vs minutes)

**Requires**: `--dag-deploy-enabled` flag set on deployment (see Update Deployments)

### Image-Only Deploy

Deploy only Docker image, skip DAG sync:

```bash
astro deploy <DEPLOYMENT_ID> --image-only
```

Use when:
- Only dependencies changed
- Dockerfile or requirements updated
- No DAG changes

### Force Deploy

Bypass safety checks and deploy:

```bash
astro deploy <DEPLOYMENT_ID> --force
```

**Caution**: Skips validation that could prevent broken deployments.

---

## Deployment API Tokens

Manage API tokens for programmatic access to deployments:

```bash
# List tokens for a deployment
astro deployment token list --deployment-id <DEPLOYMENT_ID>

# Create a new token
astro deployment token create \
  --deployment-id <DEPLOYMENT_ID> \
  --name "CI/CD Pipeline" \
  --role DEPLOYMENT_ADMIN

# Create token with expiration
astro deployment token create \
  --deployment-id <DEPLOYMENT_ID> \
  --name "Temporary Access" \
  --role DEPLOYMENT_ADMIN \
  --expiry 30  # Days until expiration (0 = never expires)
```

**Roles**:
- `DEPLOYMENT_ADMIN`: Full access to deployment

**Note**: Token value is only shown at creation time. Store it securely.

---

## Common Workflows

### First-Time Production Deployment

```bash
# 1. Login
astro login

# 2. Switch to production workspace
astro workspace list
astro workspace switch <PROD_WORKSPACE_ID>

# 3. Create deployment
astro deployment create --label production --executor celery

# 4. Note the deployment ID, then deploy
astro deploy <DEPLOYMENT_ID>
```

### Iterative DAG Development

```bash
# 1. Enable fast deploys (one-time setup)
astro deployment update <DEPLOYMENT_ID> --dag-deploy-enabled

# 2. Make DAG changes locally

# 3. Deploy quickly
astro deploy <DEPLOYMENT_ID> --dags
```

### Promoting Code from Staging to Production

```bash
# 1. Deploy to staging first
astro workspace switch <STAGING_WORKSPACE_ID>
astro deploy <STAGING_DEPLOYMENT_ID>

# 2. Test in staging

# 3. Deploy same code to production
astro workspace switch <PROD_WORKSPACE_ID>
astro deploy <PROD_DEPLOYMENT_ID>
```

---

## Configuration Management

```bash
# View CLI configuration
astro config get

# Set configuration value
astro config set <KEY> <VALUE>

# Check CLI version
astro version

# Upgrade CLI to latest version
astro upgrade
```

---

## Tips

- Use `--dags` flag for fast iteration (seconds vs minutes)
- Always test in staging workspace before production
- Use `deployment inspect` to verify deployment health before deploying
- Deployment IDs are permanent, names can change
- Most commands work with deployment ID; `inspect` also accepts `--deployment-name`
- Set `--dag-deploy-enabled` once per deployment for fast deploys
- Keep workspace context visible with `astro workspace list` (shows asterisk for current)

---

## Related Skills

- **troubleshooting-astro-deployments**: Investigate deployment issues, view logs, manage environment variables
- **managing-astro-local-env**: Manage local Airflow development environment
- **setting-up-astro-project**: Initialize and configure Astro projects
managing-astro-local-env7.87 KB

View saved version →

---
name: managing-astro-local-env
description: Manage local Airflow environment with Astro CLI (Docker and standalone modes). Use when the user wants to start, stop, or restart Airflow, view logs, query the Airflow API, troubleshoot, or fix environment issues. For project setup, see setting-up-astro-project.
---

# Astro Local Environment

This skill helps you manage your local Airflow environment using the Astro CLI.

Two modes: **Docker** (default, uses containers) and **Standalone** (Docker-free, uses a local venv — requires Airflow 3 + `uv`).

> **To set up a new project**, see the **setting-up-astro-project** skill.
> **When Airflow is running**, use MCP tools from **authoring-dags** and **testing-dags** skills.

---

## Start / Stop / Restart (Docker)

```bash
# Start local Airflow (webserver at http://localhost:8080)
astro dev start

# Stop containers (preserves data)
astro dev stop

# Kill and remove volumes (clean slate)
astro dev kill

# Restart all containers
astro dev restart

# Restart specific component
astro dev restart --scheduler
astro dev restart --webserver
```

**Default credentials:** admin / admin

**Restart after modifying:** `requirements.txt`, `packages.txt`, `Dockerfile`

> **Standalone mode?** See the next section.

---

## Standalone Mode

Docker-free local development. Runs Airflow directly on your machine in a `.venv/` managed by `uv`.

**Requirements:** Airflow 3 (runtime 3.x), `uv` on PATH. Not supported on Windows.

> Plain `astro dev init` already pins a runtime 3.x image, so no version flag is needed. See **setting-up-astro-project** for project initialization.

### Start

```bash
# One-time: set standalone as default mode
astro config set dev.mode standalone

# Or use the flag per invocation
astro dev start --standalone
```

| Flag | Description |
|------|-------------|
| `--foreground` / `-f` | Stream output in foreground |
| `--port` / `-p` | Override webserver port (default: 8080) |
| `--no-proxy` | Disable reverse proxy |

### Stop / Kill / Restart

```bash
# Stop (preserves .venv)
astro dev stop

# Kill (removes .venv and .astro/standalone/ — clean slate)
astro dev kill

# Restart (preserves .venv for fast restart, use -k to kill first)
astro dev restart
```

> If you used `--standalone` on start instead of setting the config, pass `--standalone` on every subsequent command too (stop, kill, restart, bash, run, logs, etc.).

**State locations:** venv in `.venv/`, database and logs in `.astro/standalone/`, DAGs from `dags/`.

---

## Reverse Proxy

Run multiple Airflow projects locally without port conflicts. Works in both Docker and standalone modes.

Each project gets a hostname like `<project-name>.localhost:6563`. Visit `http://localhost:6563` to see all active projects.

```bash
# Check proxy status and active routes
astro dev proxy status

# Force-stop proxy (auto-restarts on next astro dev start)
astro dev proxy stop
```

| Config | Command |
|--------|---------|
| Change proxy port | `astro config set proxy.port <port>` |
| Disable per-start | `astro dev start --no-proxy` |

Default proxy port: **6563**

---

## Check Status

```bash
astro dev ps
```

---

## View Logs

```bash
# All logs
astro dev logs

# Specific component
astro dev logs --scheduler
astro dev logs --webserver

# Follow in real-time
astro dev logs -f
```

**Standalone:** `astro dev logs` works the same but shows a unified log (no per-component filtering).

---

## Run Airflow CLI Commands

```bash
# Open a shell with Airflow environment
astro dev bash

# Run Airflow CLI commands
astro dev run airflow info
astro dev run airflow dags list
```

**Standalone:** Same commands work — `bash` opens a venv-activated shell, `run` executes in the venv.

---

## Querying the Airflow API

Use `astro api airflow` to query a running local Airflow instance. Prefer operation IDs over URL paths.

**Defaults:** localhost:8080, admin/admin (auto-detected). Override with `--api-url`, `--username`, `--password`.

### Discovery

```bash
# List all endpoints
astro api airflow ls

# Filter by keyword
astro api airflow ls dags
astro api airflow ls task

# Show params and schema for an operation
astro api airflow describe get_dag
```

### Key Flags

| Flag | Purpose |
|------|---------|
| `-p key=value` | Path parameters |
| `-F key=value` | Body/query fields (auto-converts booleans/numbers) |
| `-q` / `--jq` | jq filter on response |
| `--paginate` | Fetch all pages |
| `-X` / `--method` | Override HTTP method |
| `--generate` | Output curl command instead of executing |

### DAGs

```bash
# List all DAGs
astro api airflow get_dags

# Filter by pattern (SQL LIKE — use % wildcards)
astro api airflow get_dags -F dag_id_pattern=%etl%

# Get a specific DAG
astro api airflow get_dag -p dag_id=my_dag

# Get full details (schedule, params, etc.)
astro api airflow get_dag_details -p dag_id=my_dag

# Pause / unpause
astro api airflow patch_dag -p dag_id=my_dag -F is_paused=true
astro api airflow patch_dag -p dag_id=my_dag -F is_paused=false

# View DAG source code
astro api airflow get_dag_source -p dag_id=my_dag

# Check import errors
astro api airflow get_import_errors
```

### DAG Runs

```bash
# List runs for a DAG
astro api airflow get_dag_runs -p dag_id=my_dag

# Trigger a run
astro api airflow trigger_dag_run -p dag_id=my_dag

# Trigger with config
astro api airflow trigger_dag_run -p dag_id=my_dag -F conf[key]=value

# Get a specific run
astro api airflow get_dag_run -p dag_id=my_dag -p dag_run_id=manual__2026-04-07

# Clear (re-run) a DAG run
astro api airflow clear_dag_run -p dag_id=my_dag -p dag_run_id=manual__2026-04-07 -F dry_run=false
```

### Task Instances

```bash
# List task instances for a run
astro api airflow get_task_instances -p dag_id=my_dag -p dag_run_id=manual__2026-04-07

# Use ~ as wildcard (all DAGs or all runs)
astro api airflow get_task_instances -p dag_id=my_dag -p dag_run_id=~

# Get a specific task instance
astro api airflow get_task_instance -p dag_id=my_dag -p dag_run_id=manual__2026-04-07 -p task_id=extract

# Clear/retry failed tasks
astro api airflow post_clear_task_instances -p dag_id=my_dag \
  -F dag_run_id=manual__2026-04-07 -F only_failed=true -F dry_run=false

# Get task logs
astro api airflow get_log -p dag_id=my_dag -p dag_run_id=manual__2026-04-07 \
  -p task_id=extract -p try_number=1
```

### Config & Connections

```bash
astro api airflow get_connections
astro api airflow get_variables
astro api airflow get_config
```

### Filtering with jq

```bash
# List only DAG IDs
astro api airflow get_dags -q '.dags[].dag_id'

# Get failed task IDs from a run
astro api airflow get_task_instances -p dag_id=my_dag -p dag_run_id=~ \
  -q '[.task_instances[] | select(.state=="failed") | .task_id]'
```

---

## Troubleshooting

| Issue | Solution |
|-------|----------|
| Port 8080 in use | Stop other containers or edit `.astro/config.yaml` |
| Container won't start | `astro dev kill` then `astro dev start` |
| Package install failed | Check `requirements.txt` syntax |
| DAG not appearing | Run `astro dev parse` to check for import errors |
| Out of disk space | `docker system prune` |
| Standalone won't start | Ensure `uv` is on PATH and runtime is 3.x |
| Proxy port conflict | `astro config set proxy.port <port>` |
| `.venv` corrupted | `astro dev kill` then `astro dev start --standalone` |

### Reset Environment

When things are broken:

```bash
astro dev kill
astro dev start
```

---

## Upgrade Airflow

### Test compatibility first

```bash
astro dev upgrade-test
```

### Change version

1. Edit `Dockerfile`:
   ```dockerfile
   FROM quay.io/astronomer/astro-runtime:13.0.0
   ```

2. Restart:
   ```bash
   astro dev kill && astro dev start
   ```

---

## Related Skills

- **setting-up-astro-project**: Initialize projects and configure dependencies
- **authoring-dags**: Write DAGs (uses MCP tools, requires running Airflow)
- **testing-dags**: Test DAGs (uses MCP tools, requires running Airflow)
- **deploying-airflow**: Deploy DAGs to production (Astro, Docker Compose, Kubernetes)
migrating-airflow-2-to-38.91 KB

View saved version →

---
name: migrating-airflow-2-to-3
description: Guide for migrating Apache Airflow 2.x projects to Airflow 3.x. Use when the user mentions Airflow 3 migration, upgrade, compatibility issues, breaking changes, or wants to modernize their Airflow codebase. If you detect Airflow 2.x code that needs migration, prompt the user and ask if they want you to help upgrade. Always load this skill as the first step for any migration-related request.
hooks:
  PostToolUse:
    - matcher: "Edit"
      hooks:
        - type: command
          command: "echo 'Consider running: ruff check --preview --select AIR .'"
---

# Airflow 2 to 3 Migration

This skill helps migrate **Airflow 2.x DAG code** to **Airflow 3.x**, focusing on code changes (imports, operators, hooks, context, API usage).

**Important**: Before migrating to Airflow 3, strongly recommend upgrading to Airflow 2.11 first, then to at least Airflow 3.0.11 (ideally directly to 3.1). Other upgrade paths would make rollbacks impossible. See: https://www.astronomer.io/docs/astro/airflow3/upgrade-af3#upgrade-your-airflow-2-deployment-to-airflow-3. Additionally, early 3.0 versions have many bugs - 3.1 provides a much better experience.

## Migration at a Glance

1. Run Ruff's Airflow migration rules to auto-fix detectable issues (AIR30/AIR301/AIR302/AIR31/AIR311/AIR312).
   - `ruff check --preview --select AIR --fix --unsafe-fixes .`
2. Scan for remaining issues using the manual search checklist in [reference/migration-checklist.md](reference/migration-checklist.md).
   - Focus on: direct metadata DB access, legacy imports, scheduling/context keys, XCom pickling, datasets-to-assets, REST API/auth, plugins, and file paths.
   - Hard behavior/config gotchas to explicitly review:
     - Cron scheduling semantics: consider `AIRFLOW__SCHEDULER__CREATE_CRON_DATA_INTERVAL=True` if you need Airflow 2-style cron data intervals.
     - `.airflowignore` syntax changed from regexp to glob; set `AIRFLOW__CORE__DAG_IGNORE_FILE_SYNTAX=regexp` if you must keep regexp behavior.
     - OAuth callback URLs add an `/auth/` prefix (e.g. `/auth/oauth-authorized/google`).
     - **Shared utility imports**: Bare imports like `import common` from `dags/common/` no longer work on Astro. Use fully qualified imports: `import dags.common`.
3. Plan changes per file and issue type:
   - Fix imports - update operators/hooks/providers - refactor metadata access to using the Airflow client instead of direct access - fix use of outdated context variables - fix scheduling logic.
4. Implement changes incrementally, re-running Ruff and code searches after each major change.
5. Explain changes to the user and caution them to test any updated logic such as refactored metadata, scheduling logic and use of the Airflow context.

---

## Architecture & Metadata DB Access

Airflow 3 changes how components talk to the metadata database:

- Workers no longer connect directly to the metadata DB.
- Task code runs via the **Task Execution API** exposed by the **API server**.
- The **DAG processor** runs as an independent process **separate from the scheduler**.
- The **Triggerer** uses the task execution mechanism via an **in-process API server**.

**Trigger implementation gotcha**: If a trigger calls hooks synchronously inside the asyncio event loop, it may fail or block. Prefer calling hooks via `sync_to_async(...)` (or otherwise ensure hook calls are async-safe).

**Key code impact**: Task code can still import ORM sessions/models, but **any attempt to use them to talk to the metadata DB will fail** with:

```text
RuntimeError: Direct database access via the ORM is not allowed in Airflow 3.x
```

### Patterns to search for

When scanning DAGs, custom operators, and `@task` functions, look for:

- Session helpers: `provide_session`, `create_session`, `@provide_session`
- Sessions from settings: `from airflow.settings import Session`
- Engine access: `from airflow.settings import engine`
- ORM usage with models: `session.query(DagModel)...`, `session.query(DagRun)...`

### Replacement: Airflow Python client

Preferred for rich metadata access patterns. Add to `requirements.txt`:

```text
apache-airflow-client==<your-airflow-runtime-version>
```

Example usage:

```python
import os
from airflow.sdk import BaseOperator
import airflow_client.client
from airflow_client.client.api.dag_api import DAGApi

_HOST = os.getenv("AIRFLOW__API__BASE_URL", "https://<your-org>.astronomer.run/<deployment>/")
_TOKEN = os.getenv("DEPLOYMENT_API_TOKEN")

class ListDagsOperator(BaseOperator):
    def execute(self, context):
        config = airflow_client.client.Configuration(host=_HOST, access_token=_TOKEN)
        with airflow_client.client.ApiClient(config) as api_client:
            dag_api = DAGApi(api_client)
            dags = dag_api.get_dags(limit=10)
            self.log.info("Found %d DAGs", len(dags.dags))
```

### Replacement: Direct REST API calls

For simple cases, call the REST API directly using `requests`:

```python
from airflow.sdk import task
import os
import requests

_HOST = os.getenv("AIRFLOW__API__BASE_URL", "https://<your-org>.astronomer.run/<deployment>/")
_TOKEN = os.getenv("DEPLOYMENT_API_TOKEN")

@task
def list_dags_via_api() -> None:
    response = requests.get(
        f"{_HOST}/api/v2/dags",
        headers={"Accept": "application/json", "Authorization": f"Bearer {_TOKEN}"},
        params={"limit": 10}
    )
    response.raise_for_status()
    print(response.json())
```

---

## Ruff Airflow Migration Rules

Use Ruff's Airflow rules to detect and fix many breaking changes automatically.

- **AIR30 / AIR301 / AIR302**: Removed code and imports in Airflow 3 - **must be fixed**.
- **AIR31 / AIR311 / AIR312**: Deprecated code and imports - still work but will be removed in future versions; **should be fixed**.

Commands to run (via `uv`) against the project root:

```bash
# Auto-fix all detectable Airflow issues (safe + unsafe)
ruff check --preview --select AIR --fix --unsafe-fixes .

# Check remaining Airflow issues without fixing
ruff check --preview --select AIR .
```

---

## Reference Files

For detailed code examples and migration patterns, see:

- **[reference/config-changes.md](reference/config-changes.md)** - `airflow.cfg` section moves, renames, and removals
- **[reference/migration-patterns.md](reference/migration-patterns.md)** - Code examples for imports, scheduling, XCom, Assets, DAG bundles, runtime behavior changes
- **[reference/removed-methods.md](reference/removed-methods.md)** - Removed model methods with SDK/API migration paths
- **[reference/migration-checklist.md](reference/migration-checklist.md)** - Search patterns and fixes for issues Ruff doesn't catch

---

## Quick Reference Tables

### Key Import Changes

| Airflow 2.x | Airflow 3 |
|-------------|-----------|
| `airflow.operators.dummy_operator.DummyOperator` | `airflow.providers.standard.operators.empty.EmptyOperator` |
| `airflow.operators.bash.BashOperator` | `airflow.providers.standard.operators.bash.BashOperator` |
| `airflow.operators.python.PythonOperator` | `airflow.providers.standard.operators.python.PythonOperator` |
| `airflow.decorators.dag` | `airflow.sdk.dag` |
| `airflow.decorators.task` | `airflow.sdk.task` |
| `airflow.datasets.Dataset` | `airflow.sdk.Asset` |

### Context Key Changes

| Removed Key | Replacement |
|-------------|-------------|
| `execution_date` | `context["dag_run"].logical_date` |
| `tomorrow_ds` / `yesterday_ds` | Use `ds` with date math: `macros.ds_add(ds, 1)` / `macros.ds_add(ds, -1)` |
| `prev_ds` / `next_ds` | `prev_start_date_success` or timetable API |
| `triggering_dataset_events` | `triggering_asset_events` |
| `templates_dict` | `context["params"]` |

**Asset-triggered runs**: `logical_date` may be `None`; use `context["dag_run"].logical_date` defensively.

**Cannot trigger with future `logical_date`**: Use `logical_date=None` and rely on `run_id` instead.

Cron note: for scheduled runs using cron, `logical_date` semantics differ under `CronTriggerTimetable` (aligning `logical_date` with `run_after`). If you need Airflow 2-style cron data intervals, consider `AIRFLOW__SCHEDULER__CREATE_CRON_DATA_INTERVAL=True`.

### Default Behavior Changes

| Setting | Airflow 2 Default | Airflow 3 Default |
|---------|-------------------|-------------------|
| `schedule` | `timedelta(days=1)` | `None` |
| `catchup` | `True` | `False` |

### Callback Behavior Changes

- `on_success_callback` no longer runs on skip; use `on_skipped_callback` if needed.
- `@teardown` with `TriggerRule.ALWAYS` not allowed; teardowns now execute even if DAG run terminated early.

---

## Resources

- [Astronomer Airflow 3 Upgrade Guide](https://www.astronomer.io/docs/astro/airflow3/upgrade-af3)
- [Airflow 3 Release Notes](https://airflow.apache.org/docs/apache-airflow/stable/release_notes.html)
- [Ruff Airflow Rules](https://docs.astral.sh/ruff/rules/#airflow-air)

---

## Related Skills

- **testing-dags**: For testing DAGs after migration
- **debugging-dags**: For troubleshooting migration issues
- **deploying-airflow**: For deploying migrated DAGs to production

Referenced files: 4

migrating-ai-sdk-to-common-ai18.5 KB

View saved version →

---
name: migrating-ai-sdk-to-common-ai
description: Migrates Airflow projects from airflow-ai-sdk to apache-airflow-providers-common-ai 0.4.0+. Use when replacing airflow-ai-sdk with the official Airflow AI provider - migrating LLM decorators (@task.llm, @task.agent, @task.llm_branch, @task.embed), switching from model strings/objects to connection-based LLM configuration, updating imports from airflow_ai_sdk to the new provider, or upgrading an existing common-ai 0.1.x setup to 0.4.x (multimodal prompts, toolsets, embedding operators); also when common-ai provider, AIP-99, a pydanticai connection or migrating away from airflow-ai-sdk come up.
---

# Migrate airflow-ai-sdk to apache-airflow-providers-common-ai

This skill migrates Airflow projects from `airflow-ai-sdk` to `apache-airflow-providers-common-ai` (target **0.4.0+**), the official Airflow AI provider built on PydanticAI. It also covers upgrading projects already on common-ai 0.1.x, since several capabilities (multimodal prompts, `toolsets`, embedding operators, structured-output XCom behavior) changed between 0.1.0 and 0.4.0.

> **CRITICAL**: The new provider requires **Airflow 3.0+** and (for 0.4.0) **pydantic-ai-slim >= 1.71.0**. The API surface has changed: LLM configuration moves from code (model strings/objects) to Airflow connections (`pydanticai` type). There is no `@task.embed` in the new provider; embeddings move to the LlamaIndex integration or a plain `@task` (see Step 3).

## Before starting

Use the Grep tool with the pattern below to inventory everything that needs to migrate:

```
airflow_ai_sdk|airflow-ai-sdk|ai_sdk|@task\.llm|@task\.agent|@task\.llm_branch|@task\.embed
```

From the results, capture:

1. All files importing `airflow-ai-sdk` / `airflow_ai_sdk`
2. Which decorators are in use: `@task.llm`, `@task.agent`, `@task.llm_branch`, `@task.embed`
3. The model configuration pattern (string names like `"gpt-5"`, or `OpenAIModel(...)` objects)
4. Any `airflow_ai_sdk.BaseModel` subclasses used as `output_type`

Use this inventory to drive the steps below.

---

## Step 1: Update requirements.txt

**Remove:**
```
airflow-ai-sdk[openai]
# or any variant: airflow-ai-sdk[openai]==0.1.7, airflow-ai-sdk[anthropic], etc.
```

**Add:**
```
apache-airflow-providers-common-ai[openai]>=0.4.0
```

Use the latest available 0.x version unless the user has pinned a specific one. Available extras (0.4.0): `[openai]`, `[anthropic]`, `[google]`, `[bedrock]`, `[llamaindex]`, `[langchain]`, `[mcp]`, plus file-format extras (`[pdf]`, `[docx]`, `[parquet]`, `[avro]`) for `DocumentLoaderOperator` and `[sql]`/`[common-sql]` for the SQL operators. There are no `[groq]`/`[mistral]` extras; for those providers install the matching `pydantic-ai-slim` extra yourself.

Add `[llamaindex]` if the project migrates `@task.embed` to the `LlamaIndexEmbeddingOperator` (recommended, see Step 3). In that case `sentence-transformers` and `torch` can usually be **removed**, which shrinks the image considerably. Keep them only if the project stays on local sentence-transformers embeddings via plain `@task`.

---

## Step 2: Create PydanticAI connection

The new provider uses an Airflow connection instead of model strings or objects in code.

**Connection type:** `pydanticai`
**Default connection ID:** `pydanticai_default`

### Via environment variable (.env)

```bash
AIRFLOW_CONN_PYDANTICAI_DEFAULT='{
    "conn_type": "pydanticai",
    "password": "<api-key>",
    "extra": {
        "model": "<provider>:<model-name>"
    }
}'
```

### Model format

The model field uses `provider:model` format:

| Provider | Example model value |
|----------|-------------------|
| OpenAI | `openai:gpt-5` |
| Anthropic | `anthropic:claude-sonnet-4-20250514` |
| Google | `google:gemini-2.5-pro` |
| Groq | `groq:llama-3.3-70b-versatile` |
| Mistral | `mistral:mistral-large-latest` |
| Bedrock | `bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0` |

### Custom endpoints (Ollama, vLLM, Snowflake Cortex, etc.)

Set `host` to the base URL:
```bash
AIRFLOW_CONN_PYDANTICAI_CORTEX='{
    "conn_type": "pydanticai",
    "password": "<api-key>",
    "host": "https://my-endpoint.com/v1",
    "extra": {
        "model": "openai:<model-name>"
    }
}'
```

Use the `openai:` prefix for any OpenAI-compatible API, regardless of the actual provider.

### Connection ID convention

The env var name determines the connection ID:
- `AIRFLOW_CONN_PYDANTICAI_DEFAULT` creates `pydanticai_default`
- `AIRFLOW_CONN_PYDANTICAI_CORTEX` creates `pydanticai_cortex`

### Model resolution priority

1. `model_id` parameter on the decorator/operator (highest)
2. `model` in connection's extra JSON (fallback)

### Other connection types (0.4.0)

Besides `pydanticai`, the provider registers vendor-specific connection types: `pydanticai-azure` (Azure OpenAI: host = endpoint, extra `api_version`), `pydanticai-bedrock` (AWS credentials/region in extra), and `pydanticai-vertex` (GCP project/location in extra). The LlamaIndex and LangChain hooks read API key/host/extra from whatever connection ID they are given, so a single `pydanticai_default` connection can serve LLM calls **and** embeddings: one API key entry for the whole project.

---

## Step 3: Migrate decorators

### @task.llm

```python
# BEFORE (airflow-ai-sdk)
import airflow_ai_sdk as ai_sdk

class MyOutput(ai_sdk.BaseModel):
    field: str

@task.llm(
    model="gpt-5",                    # or model=OpenAIModel(...)
    system_prompt="You are helpful.",
    output_type=MyOutput,
)
def my_task(text: str) -> str:
    return text

# AFTER (apache-airflow-providers-common-ai)
from pydantic import BaseModel

class MyOutput(BaseModel):
    field: str

@task.llm(
    llm_conn_id="pydanticai_default",  # Airflow connection ID
    system_prompt="You are helpful.",
    output_type=MyOutput,
)
def my_task(text: str) -> str:
    return text
```

**Parameter mapping:**

| airflow-ai-sdk | common-ai provider | Notes |
|----------------|-------------------|-------|
| `model="gpt-5"` | `llm_conn_id="pydanticai_default"` | Model specified in connection |
| `model=OpenAIModel(...)` | `llm_conn_id="pydanticai_default"` | Model + endpoint in connection |
| `system_prompt="..."` | `system_prompt="..."` | Unchanged |
| `output_type=MyModel` | `output_type=MyModel` | Unchanged |
| `result_type=MyModel` | `output_type=MyModel` | `result_type` was already deprecated |
| (not available) | `model_id="openai:gpt-5"` | Override connection's model |
| (not available) | `require_approval=True` | Built-in HITL review |
| (not available) | `agent_params={...}` | Extra kwargs for pydantic-ai Agent |
| (not available) | `serialize_output=True` | Force dict shape for BaseModel output |

**Multimodal prompts (0.4.0+):** the translation function may return a `Sequence[UserContent]` instead of a string, e.g. for vision:

```python
@task.llm(llm_conn_id="pydanticai_default", system_prompt="...", output_type=ReviewAnalysis)
def analyze(text: str, image_path: str | None = None):
    if image_path:
        with open(image_path, "rb") as f:
            return [text, BinaryContent(data=f.read(), media_type="image/jpeg")]
    return text
```

This matches the old airflow-ai-sdk vision pattern, so vision code migrates unchanged. Note: common-ai **0.1.x only accepted strings** — if a project disabled vision to migrate to 0.1.0, re-enable it when bumping to 0.4.0. Non-string prompts are incompatible with `require_approval=True` / `enable_hitl_review=True` (both render the prompt as text).

**Structured output via XCom (0.4.0 behavior change):** with `output_type=<BaseModel subclass>`, the model **instance** flows through XCom on Airflow cores whose task SDK has `SUPPORTS_OPERATOR_DESERIALIZATION_WALKER` (attribute access downstream); on older cores (including Astro Runtime 3.2 task SDK 1.2.x) the provider automatically dumps to a **dict** (subscript access). Check which shape arrives at runtime before choosing attribute vs dict access downstream, or set `serialize_output=True` to force the dict shape everywhere. The `output_type` class must be defined at **module scope** (nested classes cannot be deserialized from XCom).

### @task.llm_branch

```python
# BEFORE
@task.llm_branch(
    model="gpt-5",
    system_prompt="Choose a team...",
    allow_multiple_branches=False,
)
def route(text: str) -> str:
    return text

# AFTER
@task.llm_branch(
    llm_conn_id="pydanticai_default",
    system_prompt="Choose a team...",
    allow_multiple_branches=False,    # same parameter, unchanged
)
def route(text: str) -> str:
    return text
```

Only change: `model=` becomes `llm_conn_id=`.

### @task.agent

This has the biggest API change. The Agent is no longer pre-built in user code.

```python
# BEFORE (airflow-ai-sdk) - Agent built at module level
from pydantic_ai import Agent

my_agent = Agent(
    "gpt-5",
    system_prompt="You are a research assistant.",
    tools=[search_tool, lookup_tool],
)

@task.agent(agent=my_agent)
def research(question: str) -> str:
    return question

# AFTER (common-ai provider) - No Agent object, config via parameters
from pydantic_ai.toolsets import FunctionToolset

@task.agent(
    llm_conn_id="pydanticai_default",
    system_prompt="You are a research assistant.",
    toolsets=[FunctionToolset(tools=[search_tool, lookup_tool])],
)
def research(question: str) -> str:
    return question
```

**Parameter mapping:**

| airflow-ai-sdk | common-ai provider | Notes |
|----------------|-------------------|-------|
| `agent=Agent(model, ...)` | `llm_conn_id="..."` | Model from connection |
| Agent's `system_prompt` | `system_prompt="..."` | Now a decorator param |
| Agent's `tools=[...]` | `toolsets=[FunctionToolset(tools=[...])]` | Preferred: gets automatic tool-call logging |
| Agent's `tools=[...]` | `agent_params={"tools": [...]}` | Also works, but no tool-call logging |
| Agent's `output_type` | `output_type=MyModel` | Now a decorator param |
| (not available) | `durable=True` | Step-level caching (needs `[common.ai] durable_cache_path`) |
| (not available) | `enable_hitl_review=True` | Iterative human review loop (see below) |

**Key insight:** Everything that was configured on the `Agent()` constructor now goes into either a top-level decorator parameter or `agent_params`. The `agent_params` dict is passed directly to pydantic-ai's `Agent` constructor. Prefer `toolsets` over `agent_params["tools"]`: the operator wraps each toolset in a `LoggingToolset`, so every tool call appears in the task log with timing.

**enable_hitl_review behavior:** the task generates a first draft, then **blocks** until a human acts. The reviewer uses the **HITL Review** tab/extra link on the task instance (chat UI from the provider's auto-registered `hitl_review` plugin) to request changes (agent regenerates with the feedback in its message history) or approve. Constraints: requires a string prompt, incompatible with `durable=True`, and the final (possibly regenerated) output is what flows to XCom. Warn users that the Dag run waits indefinitely at this task unless `hitl_timeout` is set. For headless testing, the plugin exposes REST endpoints under `/hitl-review`: `GET /sessions/find`, `POST /sessions/feedback`, `POST /sessions/approve`, `POST /sessions/reject` (query params `dag_id`, `task_id`, `run_id`, `map_index`).

### @task.embed (NO EQUIVALENT — three replacement options)

The new provider does NOT include an embed decorator. Pick the replacement based on what the project needs:

**Option A (recommended): `LlamaIndexEmbeddingOperator`** (0.4.0, `[llamaindex]` extra). Connection-based, one task embeds the whole document list, and with `persist_dir` the resulting vector index is persisted for retrieval (pairs with `LlamaIndexRetrievalOperator`):

```python
from airflow.providers.common.ai.operators.llamaindex_embedding import LlamaIndexEmbeddingOperator

_embeddings = LlamaIndexEmbeddingOperator(
    task_id="create_embeddings",
    documents=[{"text": "...", "metadata": {"id": 1}}, ...],  # templated, accepts XComArg
    llm_conn_id="pydanticai_default",   # reuses the same connection (API key only)
    embed_model="text-embedding-3-small",
    persist_dir=f"{AIRFLOW_HOME}/include/my_index",  # optional; local path or s3://, gs://, ...
)
```

The operator returns `{"chunks": [{"text", "metadata", "vector"}], ...}`. Put a stable key into each document's `metadata` — it round-trips through chunking, so vectors can be mapped back to source records.

**Option B: `LlamaIndexHook` for raw vectors** (no operator, no persisted index). Shortest path when vectors go straight to a database:

```python
@task
def create_embeddings(rows):
    from airflow.providers.common.ai.hooks.llamaindex import LlamaIndexHook
    embed_model = LlamaIndexHook(
        llm_conn_id="pydanticai_default",
        embed_model="text-embedding-3-small",
    ).get_embedding_model()
    vectors = embed_model.get_text_embedding_batch([r["text"] for r in rows])
    return list(zip([r["id"] for r in rows], vectors))
```

**Option C: plain `@task` with sentence-transformers** (keeps the old local/offline behavior, no API cost; requires keeping `sentence-transformers` + `torch` in requirements):

```python
@task
def embed_texts(texts: list[str]) -> list[list[float]]:
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer("all-MiniLM-L6-v2")
    return model.encode(texts, normalize_embeddings=True).tolist()
```

Note on dimensions: switching from `all-MiniLM-L6-v2` (384) to `text-embedding-3-small` (1536) changes vector size — existing stored embeddings must be regenerated, and fixed-size vector columns (e.g. pgvector `vector(384)`) need a schema change. Embed all texts in one task/batch call rather than `.expand()` per text: batching is one API round-trip and avoids per-task model loading.

---

## Step 4: Update imports

| Old import | New import |
|-----------|-----------|
| `import airflow_ai_sdk as ai_sdk` | Remove entirely |
| `from airflow_ai_sdk import BaseModel` | `from pydantic import BaseModel` |
| `from airflow_ai_sdk.models.base import BaseModel` | `from pydantic import BaseModel` |
| `class Foo(ai_sdk.BaseModel):` | `class Foo(BaseModel):` |
| `from pydantic_ai import Agent` | Remove if Agent was only used for `@task.agent` |
| `from pydantic_ai.models.openai import OpenAIModel` | Remove (model config in connection now) |
| (new) | `from pydantic_ai.toolsets import FunctionToolset` for `@task.agent` toolsets |

The `@task.llm`, `@task.agent`, `@task.llm_branch` decorators are auto-registered by the provider. No explicit import needed beyond `from airflow.sdk import task`.

`pydantic_ai` imports for non-decorator usage (e.g., `BinaryContent` for multimodal) are still valid since the new provider depends on `pydantic-ai-slim` (>= 1.71.0 for provider 0.4.0).

---

## Step 5: Update connections.yaml (if used for local testing)

```yaml
pydanticai_default:
  conn_type: pydanticai
  password: <api-key>
  extra:
    model: "openai:gpt-5"
```

For custom endpoints:
```yaml
pydanticai_cortex:
  conn_type: pydanticai
  password: <api-key>
  host: https://my-endpoint.com/v1
  extra:
    model: "openai:llama3.1-8b"
```

---

## Step 6: Clean up env vars

The new provider reads model config from the `pydanticai` connection, so env vars that previously fed the model in code are usually redundant. Before removing any of them, grep the project (and any sibling scripts/services) to confirm nothing else still references them:

```
OPENAI_API_KEY|OPENAI_BASE_URL|ANTHROPIC_API_KEY|GOOGLE_API_KEY
```

Candidates for removal **only if no other code references them**:
- `OPENAI_API_KEY` (now in the pydanticai connection's password field)
- `OPENAI_BASE_URL` (now in the connection's host field)
- Custom model name vars (now in the connection's extra.model)

If anything outside the migrated DAGs still uses them (other DAGs not yet migrated, helper scripts, non-Airflow services sharing the `.env`), leave them in place.

**Keep** `AIRFLOW_CONN_*` env vars for all connections.

---

## Step 7: Verify

After migration, grep the codebase to confirm no stale references remain:

```
airflow_ai_sdk|airflow-ai-sdk|ai_sdk\.BaseModel|from pydantic_ai import Agent|from pydantic_ai.models
```

Verify:
- [ ] No imports from `airflow_ai_sdk`
- [ ] No `Agent()` objects created for `@task.agent` (unless used outside decorators)
- [ ] No `model=` parameter on LLM decorators (should be `llm_conn_id=`)
- [ ] All `@task.embed` replaced (LlamaIndex operator/hook or plain `@task`); stored embeddings regenerated if the model/dimensions changed
- [ ] Vision translation functions return `[text, BinaryContent(...)]` again if they were string-only-restricted under common-ai 0.1.x
- [ ] Downstream consumers of `output_type=BaseModel` results use the XCom shape that actually arrives (dict on older cores, instance on newer; `serialize_output=True` pins it)
- [ ] `pydanticai` connection configured in `.env` or connections.yaml
- [ ] `requirements.txt` has `apache-airflow-providers-common-ai[...]` instead of `airflow-ai-sdk[...]`; `torch`/`sentence-transformers` removed if no longer used
- [ ] Run the Dags end-to-end: tasks with `enable_hitl_review=True` or `require_approval=True` wait for human input, so the test plan must include acting on them (UI tab or `/hitl-review` REST)

---

## Quick reference: New features in common-ai provider

These features are available after migration but have no airflow-ai-sdk equivalent:

| Feature | Parameter / API | Since | Description |
|---------|-----------------|-------|-------------|
| HITL approval | `require_approval=True` on `@task.llm` | 0.1.0 | Pause for human review before returning |
| HITL review loop | `enable_hitl_review=True` on `@task.agent` | 0.1.0 | Iterative review with regeneration (chat UI via `hitl_review` plugin) |
| Durable execution | `durable=True` on `@task.agent` | 0.1.0 | Step-level caching for resilience |
| Tool logging | `enable_tool_logging=True` on `@task.agent` | 0.1.0 | INFO-level tool call logs (default: on; requires `toolsets`) |
| Model override | `model_id="openai:gpt-5"` | 0.1.0 | Override connection's model per-task |
| File analysis | `@task.llm_file_analysis` | 0.1.0 | Analyze files/images via ObjectStoragePath |
| NL-to-SQL | `@task.llm_sql` | 0.1.0 | Generate SQL from natural language |
| Multimodal prompts | Translation function returns `Sequence[UserContent]` | 0.4.0 | Vision and other binary content in `@task.llm` / `@task.agent` / `@task.llm_branch` |
| Pydantic instance via XCom | `output_type=BaseModel` (with `serialize_output` opt-out) | 0.4.0 | Instance flows through XCom on capable cores; dict fallback otherwise |
| Embeddings | `LlamaIndexEmbeddingOperator` (+ `persist_dir`) | 0.4.0 | Connection-based embeddings + persisted vector index |
| Retrieval | `LlamaIndexRetrievalOperator` | 0.4.0 | Top-k similarity search over a persisted index |
migrating-dagster-to-airflow14.8 KB

View saved version →

---
name: migrating-dagster-to-airflow
description: Guide for migrating Dagster projects to Apache Airflow 3 on Astro. Use when the user mentions migrating, converting, or porting Dagster (or Dagster+) code to Airflow or Astro, wants to plan or assess such a migration, or asks what a Dagster construct maps to in Airflow. Covers assets, partitions, schedules, sensors, declarative automation, resources, IO managers, ops/jobs, dbt, Pipes, Components, and Dagster+ platform config. Always load this skill as the first step for any Dagster-to-Airflow request.
hooks:
  PostToolUse:
    - matcher: "Edit"
      hooks:
        - type: command
          command: "echo 'Migrated DAG edited: consider re-running gates 1-3: python3 scripts/validate_dag.py <astro_project> --manifest manifest.json'"
---

# Dagster → Airflow 3 (Astro) migration

Migrate a Dagster project to Airflow 3 on Astro Runtime, honestly. The migration is asset-first (Dagster asset graphs translate to Airflow assets and asset-aware schedules, not flattened DAGs), incremental (domain by domain, Dagster stays authoritative until parity), and honest (every definition gets an explicit disposition; semantic deltas are documented, never papered over).

First time driving this? Read `reference/quickstart.md` first: hour-one commands, the glossary, and what can and cannot break.

## Migration at a glance

1. Baseline the source project's tests, then inventory it read-only (`scripts/inventory.py` → manifest).
2. Review classifications (MECH/JUDG/REDESIGN/NONE per `reference/mapping.md`); make the go/no-go call (three outcomes; migrate-with-conditions is the common case, stay is the narrow one); plan DAG boundaries, per-edge IO decisions, and Gate 3 expectations into the manifest.
3. Trial-migrate 2-3 representative units end-to-end through every validation gate.
4. Migrate domain by domain through the six-gate ladder (`reference/validation.md`), tracking per-unit state (`scripts/status.py`); fix failure classes via `reference/troubleshooting.md`, never stub.
5. Map the platform layer (secrets, alerts, CI/CD, Deployments) per `reference/astro-deployment.md`.
6. Run side by side, then cut over per domain (consumers unpause first; see the checklist), keeping rollback one step away.
7. Deliver the migration report: every definition dispositioned, an equivalence row per trigger, losses stated plainly.

## Version drift

Verified against Airflow 3.3.0 / Astro Runtime 3.3-2 / astronomer-cosmos 1.15 / Dagster 1.13 (2026-07). Version-sensitive rows in the references carry their floor (notably the 3.2-vs-3.3 partition surface). Before relying on a version-gated claim: check the target (`airflow version`, `astro deployment inspect`), probe imports for sdk surface (`python3 -c "from airflow.sdk import X"`), and prefer `--help` / API spec discovery over assuming verbatim CLI/REST contracts on newer versions. Playbook entries are version-scoped per entry.

## Requirements

- Target **Astro Runtime 3.3+** (Airflow 3.3+); the native asset-partition surface requires it. Below 3.2 the mapping degrades badly; say so and recommend upgrading before migrating.
- The Dagster repo, and ideally a running Dagster instance (its materialization metadata provides parity-test fixtures).
- `astro` CLI for the target project.

## Hard rules

1. **Never stub.** A translated unit either works through its validation gate or is deferred with a written reason. Fake-success bodies and workaround code with long justifying comments are failures.
2. **No silent omissions.** Every record in the inventory manifest ends `complete` or `deferred (reason)`. `scripts/status.py summary` exits nonzero otherwise; run it before claiming done.
3. **Equivalence rows for every trigger.** Each schedule/sensor/automation condition gets a report row: source spelling, target spelling, delta in one sentence. Semantic deltas exist (catchup, on_cron inversion, eager guarantees); the sin is not the delta, it is the undocumented delta.
4. **Fix classes, not instances.** When a translation pattern fails validation, fix the pattern (and record it in `reference/troubleshooting.md`), then re-apply; do not hand-patch one unit.
5. **Do not invent APIs.** The references contain verified names only. Anything not covered there gets verified against official docs before use.

## Workflow

### Phase 0: Preflight

Confirm target Runtime version, `astro` CLI presence, and repo access. Detect the project layout: classic (`@repository`/`workspace.yaml`), modern (`Definitions`), or Components (`pyproject.toml [tool.dg]`, `defs.yaml` files); all three occur, sometimes together. Baseline the source project's test suite now: pre-existing failures are recorded and excluded from migration blame.

### Phase 1: Inventory (read-only)

```
python3 scripts/inventory.py <dagster_repo> --out manifest.json          # static scan
python3 scripts/inventory.py <dagster_repo> --runtime --out manifest.json # + runtime introspection when the project imports
```

The manifest lists every definition with file:line, captured params, current-vs-deprecated spelling, and dependency edges with their IO manager; every record starts `classification: "pending"`. Classifying is YOUR first judgment task: assign each record MECH / JUDG / REDESIGN / NONE from its row in `reference/mapping.md` and write it into the manifest. The scanner enumerates (deterministic completeness); the agent classifies (judgment). A record you cannot map to a mapping.md row is itself a finding: record it, do not guess. Also grep for `DAGSTER_CLOUD_` and `EnvVar(` (platform layer, Phase 5).

Manifest conventions: the canonical manifest lives in the migration run directory. Once the Astro project exists (Phase 2 scaffold), copy the manifest to its `include/inventory/manifest.json` so the Gate 3 pytest and `status.py` defaults find it; until then it just stays in the run dir (keep the two in sync afterward, the run-dir copy wins). Static records are the canonical migration units; runtime-mode enrichment merges into them, and only genuinely runtime-only definitions (factory-generated) become new units.

Emit the migration report skeleton now: one section per manifest record, plus the secrets/env naming map from `reference/astro-deployment.md`. Scale the skeleton to the project: a secretless local project gets a one-line "no secrets/platform layer" note, not empty boilerplate sections.

### Phase 1.5: Go/no-go (the honest gate)

Before translating anything, answer the project-level question the inventory makes answerable: **what does this team give up by migrating, and does each loss have an acceptable answer?** Assess the NONE and REDESIGN rows against what is load-bearing for THIS team, evaluating the mitigation, not just the loss:

| If load-bearing | The Airflow-world answer | Stay-signal only if |
|---|---|---|
| dbt rebuild-on-code-change (`code_version_changed()`) | State-aware dbt builds on a cron (Fusion / dbt State skip unchanged models per run, so the post-deploy tick rebuilds exactly what changed), and/or CI-triggered `dbt build` on merge (PR-gated, often an upgrade) | The team can neither run a state-aware dbt stack nor dbt from CI |
| Freshness driving materialization | Astro Observe freshness SLAs / Timeliness alerts + scheduled runs sized to the SLA | Freshness-triggered compute is genuinely irreplaceable by schedule+alerting |
| Per-asset cost accounting (Insights) | Astro Observe pipeline-level warehouse cost management; per-asset granularity is lost | Per-ASSET chargeback is a contractual/organizational requirement |
| Asset catalog / column-level lineage as daily tools | Airflow 3 asset views + OpenLineage/Astro lineage (asset-level) | Column-level lineage is embedded in daily workflows with no external catalog |
| Deep AutomationCondition compositions, `can_subset`, selective per-partition materialization | Most decompose to cron/asset schedules (see `reference/automation.md`); the residue is redesigned per domain | Multiple domains depend on compositions that decompose to nothing |
| Sensor cursor transactionality, run-scoped teardown | Idempotent consumers + `max_active_runs`; context managers in task bodies | Exactly-once event coalescing is a correctness requirement that idempotency cannot absorb |

One rule the table implies, stated plainly: **no dbt-only condition reaches "stay."** Between Cosmos, state-aware dbt builds, and CI-triggered builds, every dbt-workflow loss has an accepted-practice mitigation (execution-proven in this skill's eval program, including on a real warehouse); dbt items are conditions to record, never blockers. The observability rows (per-asset cost, column-level lineage) are separate conditions and are evaluated on their own, even for dbt-heavy teams.

The gate's outcome is three-valued, and the middle one is the common case:

- **Migrate**: no stay-signals; proceed to Phase 2.
- **Migrate with conditions** (most real projects): losses exist, mitigations are named and accepted in writing in the report's first section, specific domains may carry REDESIGN work; proceed to Phase 2 with those conditions recorded.
- **Stay on Dagster, today**: reserved for the case where MULTIPLE stay-signal conditions in the right column genuinely hold at once and the migration is not externally mandated. Then the honest deliverable is that recommendation, in writing, with the specific unmitigated losses named, and the run stops there. A migration guide that cannot say "don't" cannot be trusted when it says "do", but "don't" is earned by unmitigatable losses, not by the mere existence of deltas.

### Phase 2: Plan

- **DAG boundaries**: decide which asset-dependency edges become asset-aware schedules (cross-DAG) vs task ordering (intra-DAG). Group by domain/schedule cadence/team ownership; `define_asset_job` selections usually name the natural domains.
- **Per-edge IO decisions** via the tree in `reference/io-and-data-passing.md` (fuse / explicit storage / XCom).
- **Order**: leaf domains first, dependency order after; the platform layer last.
- **Fill each planned unit's target expectations into the manifest**: `dag_id`, `task_count`, `edges`, `schedule`, `asset_outlets` per unit. Gate 3 asserts against exactly these fields; a unit without them is skipped by validation, so an unenriched manifest means Gate 3 checks nothing (validate_dag reports skipped counts loudly, do not ignore them).
- Scaffold the target: `astro dev init`, shared helpers under `include/`. House conventions the scaffold imposes (e.g. a test demanding `retries >= 2`) do NOT override source fidelity: source behavior wins; convention adoption is a post-cutover improvement listed in the report, and the scaffold test gets skipped with an explicit reason.

### Phase 3: Trial

Migrate 2-3 representative units end-to-end through every gate before fanning out. Pick one MECH asset, one partitioned asset, one JUDG case, or the nearest available mix (small projects may have no partitioned or no MECH assets; pick one full path through a real DAG instead). What the trial teaches goes into `reference/troubleshooting.md` before scaling; if the trial fails structurally, stop and rework the plan, not the units.

### Phase 4: Migrate, domain by domain

Per unit, the state machine (tracked in the manifest):

```
pending → translate → fix-import → fix-lint → fix-tests → verify-parity → complete
                                    ↘ deferred (reason required)
```

- Translate using the reference file for the construct (routing table below). Rich context beats cleverness: read the source unit, its mapping rows, and a nearby already-migrated example.
- Validate through the gates: `python3 scripts/validate_dag.py <astro_project> --manifest manifest.json` (gates 1-3), then execution and parity per `reference/validation.md`.
- On gate failure, retry with the latest validator output in context (cap ~10 attempts, then defer with the failure class).
- Advance state only on gate pass: `python3 scripts/status.py advance <unit-id> ...`. A wrong disposition is corrected with `status.py reopen <unit-id> --reason ...`. The no-hand-editing rule applies to the STATE field (`status`) only; the PLAN fields (`dag_id`, `task_count`, `edges`, `schedule`, `asset_outlets`, `target`) are the planner's to write in Phase 2.
- Units that deliberately translate to NO DAG of their own (helpers absorbed into tasks, policies that became alerts, resources that became connections) are dispositioned `complete` with `target: "none"` and evidence naming where they went; Gate 3 skips them by design.
- Commit per unit, atomically.

### Phase 5: Platform layer

`reference/astro-deployment.md`: Deployments topology, secrets/connection naming map, CI/CD and preview Deployments, alert-policy mapping, Observe/lineage expectations, the `DAGSTER_CLOUD_*` in-code rewrite checklist.

### Phase 6: Side-by-side and cutover

Dagster remains authoritative. Run migrated DAGs shadowed/paused; compare outputs over the same logical window (row counts + checksums; recompute expected values from the Dagster-produced output itself, using recorded materialization metadata only opportunistically, per `reference/validation.md` Gate 5). Flip schedules one domain per change window: pause the Dagster schedule, unpause the Airflow DAG; rollback is the reverse. Keep Dagster readable after cutover (run history does not migrate).

### Phase 7: Final report

`scripts/status.py summary` must pass. The report contains: the go/no-go assessment (Phase 1.5) and its rationale, disposition table for every definition, all equivalence rows, the NONE/REDESIGN losses stated plainly (lineage depth, code_version triggers, Insights cost accounting, sensor cursor transactionality), the secrets map, and the deferred list with reasons. Spot-check ten `complete` claims before delivering it.

## Reference routing

| Construct encountered | Read |
|---|---|
| First hour, glossary, what can break | `reference/quickstart.md` |
| Anything (first stop: one row per construct) | `reference/mapping.md` |
| Asset-key → URI convention, translation granularity, external/observable assets | `reference/assets.md` |
| Asset deps, IO managers, XCom, storage decisions | `reference/io-and-data-passing.md` |
| Any `partitions_def`, partition mappings, backfills | `reference/partitions.md` |
| Schedules, sensors, AutomationCondition, freshness | `reference/automation.md` |
| `@dbt_assets`, translators, dbt Cloud | `reference/dbt.md` |
| Components (defs.yaml), custom Component subclasses, dynamic generation | `reference/components.md` |
| dagster_cloud.yaml, secrets, alerts, CI/CD, cutover | `reference/astro-deployment.md` |
| Gates, parity testing, state machine | `reference/validation.md` |
| Failure classes seen before | `reference/troubleshooting.md` |

## Scripts

| Script | Purpose |
|---|---|
| `scripts/inventory.py` | Scan the Dagster repo → JSON manifest (static + optional runtime mode) |
| `scripts/validate_dag.py` | Gates 1-3 against the generated Astro project |
| `scripts/status.py` | Per-unit state machine + completeness gate |

Referenced files: 22

profiling-tables3.77 KB

View saved version →

---
name: profiling-tables
description: Deep-dive data profiling for a specific table. Use when the user asks to profile a table, wants statistics about a dataset, asks about data quality, or needs to understand a table's structure and content. Requires a table name.
---

# Data Profile

Generate a comprehensive profile of a table that a new team member could use to understand the data.

## Step 1: Basic Metadata

Query column metadata:

```sql
SELECT COLUMN_NAME, DATA_TYPE, COMMENT
FROM <database>.INFORMATION_SCHEMA.COLUMNS
WHERE TABLE_SCHEMA = '<schema>' AND TABLE_NAME = '<table>'
ORDER BY ORDINAL_POSITION
```

If the table name isn't fully qualified, search INFORMATION_SCHEMA.TABLES to locate it first.

## Step 2: Size and Shape

Run via `run_sql`:

```sql
SELECT
    COUNT(*) as total_rows,
    COUNT(*) / 1000000.0 as millions_of_rows
FROM <table>
```

## Step 3: Column-Level Statistics

For each column, gather appropriate statistics based on data type:

### Numeric Columns
```sql
SELECT
    MIN(column_name) as min_val,
    MAX(column_name) as max_val,
    AVG(column_name) as avg_val,
    STDDEV(column_name) as std_dev,
    PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY column_name) as median,
    SUM(CASE WHEN column_name IS NULL THEN 1 ELSE 0 END) as null_count,
    COUNT(DISTINCT column_name) as distinct_count
FROM <table>
```

### String Columns
```sql
SELECT
    MIN(LEN(column_name)) as min_length,
    MAX(LEN(column_name)) as max_length,
    AVG(LEN(column_name)) as avg_length,
    SUM(CASE WHEN column_name IS NULL OR column_name = '' THEN 1 ELSE 0 END) as empty_count,
    COUNT(DISTINCT column_name) as distinct_count
FROM <table>
```

### Date/Timestamp Columns
```sql
SELECT
    MIN(column_name) as earliest,
    MAX(column_name) as latest,
    DATEDIFF('day', MIN(column_name), MAX(column_name)) as date_range_days,
    SUM(CASE WHEN column_name IS NULL THEN 1 ELSE 0 END) as null_count
FROM <table>
```

## Step 4: Cardinality Analysis

For columns that look like categorical/dimension keys:

```sql
SELECT
    column_name,
    COUNT(*) as frequency,
    ROUND(COUNT(*) * 100.0 / SUM(COUNT(*)) OVER(), 2) as percentage
FROM <table>
GROUP BY column_name
ORDER BY frequency DESC
LIMIT 20
```

This reveals:
- High-cardinality columns (likely IDs or unique values)
- Low-cardinality columns (likely categories or status fields)
- Skewed distributions (one value dominates)

## Step 5: Sample Data

Get representative rows:

```sql
SELECT *
FROM <table>
LIMIT 10
```

If the table is large and you want variety, sample from different time periods or categories.

## Step 6: Data Quality Assessment

Summarize quality across dimensions:

### Completeness
- Which columns have NULLs? What percentage?
- Are NULLs expected or problematic?

### Uniqueness
- Does the apparent primary key have duplicates?
- Are there unexpected duplicate rows?

### Freshness
- When was data last updated? (MAX of timestamp columns)
- Is the update frequency as expected?

### Validity
- Are there values outside expected ranges?
- Are there invalid formats (dates, emails, etc.)?
- Are there orphaned foreign keys?

### Consistency
- Do related columns make sense together?
- Are there logical contradictions?

## Step 7: Output Summary

Provide a structured profile:

### Overview
2-3 sentences describing what this table contains, who uses it, and how fresh it is.

### Schema
| Column | Type | Nulls% | Distinct | Description |
|--------|------|--------|----------|-------------|
| ... | ... | ... | ... | ... |

### Key Statistics
- Row count: X
- Date range: Y to Z
- Last updated: timestamp

### Data Quality Score
- Completeness: X/10
- Uniqueness: X/10
- Freshness: X/10
- Overall: X/10

### Potential Issues
List any data quality concerns discovered.

### Recommended Queries
3-5 useful queries for common questions about this data.
setting-up-astro-project3.18 KB

View saved version →

---
name: setting-up-astro-project
description: Initialize and configure Astro/Airflow projects. Use when the user wants to create a new project, set up dependencies, configure connections/variables, or understand project structure. For running the local environment, see managing-astro-local-env.
---

# Astro Project Setup

This skill helps you initialize and configure Airflow projects using the Astro CLI.

> **To run the local environment**, see the **managing-astro-local-env** skill.
> **To write DAGs**, see the **authoring-dags** skill.
> **Open-source alternative:** If the user isn't on Astro, guide them to Apache Airflow's Docker Compose quickstart for local dev and the Helm chart for production. For deployment strategies, use the `deploying-airflow` skill.

---

## Initialize a New Project

```bash
astro dev init
```

> **Don't pass `--airflow-version` or `--runtime-version` unless the user explicitly asks for a specific pin.** Plain `astro dev init` resolves to the latest Astro Runtime — that's the right default. Specifying a version risks pinning to a stale value from training data. If the user wants to know what was installed, read the generated `Dockerfile` afterward instead of guessing.

Creates this structure:
```
project/
├── dags/                # DAG files
├── include/             # SQL, configs, supporting files
├── plugins/             # Custom Airflow plugins
├── tests/               # Unit tests
├── Dockerfile           # Image customization
├── packages.txt         # OS-level packages
├── requirements.txt     # Python packages
└── airflow_settings.yaml # Connections, variables, pools
```

---

## Adding Dependencies

### Python Packages (requirements.txt)

```
apache-airflow-providers-snowflake==5.3.0
pandas==2.1.0
requests>=2.28.0
```

### OS Packages (packages.txt)

```
gcc
libpq-dev
```

### Custom Dockerfile

For complex setups (private PyPI, custom scripts):

```dockerfile
FROM quay.io/astronomer/astro-runtime:12.4.0

RUN pip install --extra-index-url https://pypi.example.com/simple my-package
```

**After modifying dependencies:** Run `astro dev restart`

---

## Configuring Connections & Variables

### airflow_settings.yaml

Loaded automatically on environment start:

```yaml
airflow:
  connections:
    - conn_id: my_postgres
      conn_type: postgres
      host: host.docker.internal
      port: 5432
      login: user
      password: pass
      schema: mydb

  variables:
    - variable_name: env
      variable_value: dev

  pools:
    - pool_name: limited_pool
      pool_slot: 5
```

### Export/Import

```bash
# Export from running environment
astro dev object export --connections --file connections.yaml

# Import to environment
astro dev object import --connections --file connections.yaml
```

---

## Validate Before Running

Parse DAGs to catch errors without starting the full environment:

```bash
astro dev parse
```

---

## Related Skills

- **managing-astro-local-env**: Start, stop, and troubleshoot the local environment
- **authoring-dags**: Write and validate DAGs (uses MCP tools)
- **testing-dags**: Test DAGs (uses MCP tools)
- **deploying-airflow**: Deploy DAGs to production (Astro, Docker Compose, Kubernetes)
testing-dags10.3 KB

View saved version →

---
name: testing-dags
description: Complex DAG testing workflows with debugging and fixing cycles. Use for multi-step testing requests like "test this dag and fix it if it fails", "test and debug", "run the pipeline and troubleshoot issues". For simple test requests ("test dag", "run dag"), the airflow entrypoint skill handles it directly. This skill is for iterative test-debug-fix cycles.
---

# DAG Testing Skill

Use `af` commands to test, debug, and fix DAGs in iterative cycles.

## Running the CLI

These commands assume `af` is on PATH. Run via `astro otto` to get it automatically, or install standalone with `uv tool install astro-airflow-mcp`.

---

## Quick Validation with Astro CLI

If the user has the Astro CLI available, these commands provide fast feedback without needing a running Airflow instance:

```bash
# Parse DAGs to catch import errors, syntax issues, and DAG-level problems
astro dev parse

# Run pytest against DAGs (runs tests in tests/ directory)
astro dev pytest
```

Use these for quick validation during development. For full end-to-end testing against a live Airflow instance, continue to the trigger-and-wait workflow below.

---

## FIRST ACTION: Just Trigger the DAG

When the user asks to test a DAG, your **FIRST AND ONLY action** should be:

```bash
af runs trigger-wait <dag_id>
```

**DO NOT:**
- Call `af dags list` first
- Call `af dags get` first
- Call `af dags errors` first
- Use `grep` or `ls` or any other bash command
- Do any "pre-flight checks"

**Just trigger the DAG.** If it fails, THEN debug.

---

## Testing Workflow Overview

```
┌─────────────────────────────────────┐
│ 1. TRIGGER AND WAIT                 │
│    Run DAG, wait for completion     │
└─────────────────────────────────────┘
                 ↓
        ┌───────┴───────┐
        ↓               ↓
   ┌─────────┐    ┌──────────┐
   │ SUCCESS │    │ FAILED   │
   │ Done!   │    │ Debug... │
   └─────────┘    └──────────┘
                       ↓
        ┌─────────────────────────────────────┐
        │ 2. DEBUG (only if failed)           │
        │    Get logs, identify root cause    │
        └─────────────────────────────────────┘
                       ↓
        ┌─────────────────────────────────────┐
        │ 3. FIX AND RETEST                   │
        │    Apply fix, restart from step 1   │
        └─────────────────────────────────────┘
```

**Philosophy: Try first, debug on failure.** Don't waste time on pre-flight checks — just run the DAG and diagnose if something goes wrong.

---

## Phase 1: Trigger and Wait

Use `af runs trigger-wait` to test the DAG:

### Primary Method: Trigger and Wait

```bash
af runs trigger-wait <dag_id> --timeout 300
```

**Example:**

```bash
af runs trigger-wait my_dag --timeout 300
```

**Why this is the preferred method:**
- Single command handles trigger + monitoring
- Returns immediately when DAG completes (success or failure)
- Includes failed task details if run fails
- No manual polling required

### Response Interpretation

**Success:**
```json
{
  "dag_run": {
    "dag_id": "my_dag",
    "dag_run_id": "manual__2025-01-14T...",
    "state": "success",
    "start_date": "...",
    "end_date": "..."
  },
  "timed_out": false,
  "elapsed_seconds": 45.2
}
```

**Failure:**
```json
{
  "dag_run": {
    "state": "failed"
  },
  "timed_out": false,
  "elapsed_seconds": 30.1,
  "failed_tasks": [
    {
      "task_id": "extract_data",
      "state": "failed",
      "try_number": 2
    }
  ]
}
```

**Timeout:**
```json
{
  "dag_id": "my_dag",
  "dag_run_id": "manual__...",
  "state": "running",
  "timed_out": true,
  "elapsed_seconds": 300.0,
  "message": "Timed out after 300 seconds. DAG run is still running."
}
```

### Alternative: Trigger and Monitor Separately

Use this only when you need more control:

```bash
# Step 1: Trigger
af runs trigger my_dag
# Returns: {"dag_run_id": "manual__...", "state": "queued"}

# Step 2: Check status
af runs get my_dag manual__2025-01-14T...
# Returns current state
```

---

## Handling Results

### If Success

The DAG ran successfully. Summarize for the user:
- Total elapsed time
- Number of tasks completed
- Any notable outputs (if visible in logs)

**You're done!**

### If Timed Out

The DAG is still running. Options:
1. Check current status: `af runs get <dag_id> <dag_run_id>`
2. Ask user if they want to continue waiting
3. Increase timeout and try again

### If Failed

Move to Phase 2 (Debug) to identify the root cause.

---

## Phase 2: Debug Failures (Only If Needed)

When a DAG run fails, use these commands to diagnose:

### Get Comprehensive Diagnosis

```bash
af runs diagnose <dag_id> <dag_run_id>
```

Returns in one call:
- Run metadata (state, timing)
- All task instances with states
- Summary of failed tasks
- State counts (success, failed, skipped, etc.)

### Get Task Logs

```bash
af tasks logs <dag_id> <dag_run_id> <task_id>
```

**Example:**

```bash
af tasks logs my_dag manual__2025-01-14T... extract_data
```

**For specific retry attempt:**

```bash
af tasks logs my_dag manual__2025-01-14T... extract_data --try 2
```

**Look for:**
- Exception messages and stack traces
- Connection errors (database, API, S3)
- Permission errors
- Timeout errors
- Missing dependencies

### Check Upstream Tasks

If a task shows `upstream_failed`, the root cause is in an upstream task. Use `af runs diagnose` to find which task actually failed.

### Check Import Errors (If DAG Didn't Run)

If the trigger failed because the DAG doesn't exist:

```bash
af dags errors
```

This reveals syntax errors or missing dependencies that prevented the DAG from loading.

---

## Phase 3: Fix and Retest

Once you identify the issue:

### Common Fixes

| Issue | Fix |
|-------|-----|
| Missing import | Add to DAG file |
| Missing package | Add to `requirements.txt` |
| Connection error | Check `af config connections`, verify credentials |
| Variable missing | Check `af config variables`, create if needed |
| Timeout | Increase task timeout or optimize query |
| Permission error | Check credentials in connection |

### After Fixing

1. Save the file
2. **Retest:** `af runs trigger-wait <dag_id>`

**Repeat the test → debug → fix loop until the DAG succeeds.**

---

## CLI Quick Reference

| Phase | Command | Purpose |
|-------|---------|---------|
| Test | `af runs trigger-wait <dag_id>` | **Primary test method — start here** |
| Test | `af runs trigger <dag_id>` | Start run (alternative) |
| Test | `af runs get <dag_id> <run_id>` | Check run status |
| Debug | `af runs diagnose <dag_id> <run_id>` | Comprehensive failure diagnosis |
| Debug | `af tasks logs <dag_id> <run_id> <task_id>` | Get task output/errors |
| Debug | `af dags errors` | Check for parse errors (if DAG won't load) |
| Debug | `af dags get <dag_id>` | Verify DAG config |
| Debug | `af dags explore <dag_id>` | Full DAG inspection |
| Config | `af config connections` | List connections |
| Config | `af config variables` | List variables |

---

## Testing Scenarios

### Scenario 1: Test a DAG (Happy Path)

```bash
af runs trigger-wait my_dag
# Success! Done.
```

### Scenario 2: Test a DAG (With Failure)

```bash
# 1. Run and wait
af runs trigger-wait my_dag
# Failed...

# 2. Find failed tasks
af runs diagnose my_dag manual__2025-01-14T...

# 3. Get error details
af tasks logs my_dag manual__2025-01-14T... extract_data

# 4. [Fix the issue in DAG code]

# 5. Retest
af runs trigger-wait my_dag
```

### Scenario 3: DAG Doesn't Exist / Won't Load

```bash
# 1. Trigger fails - DAG not found
af runs trigger-wait my_dag
# Error: DAG not found

# 2. Find parse error
af dags errors

# 3. [Fix the issue in DAG code]

# 4. Retest
af runs trigger-wait my_dag
```

### Scenario 4: Debug a Failed Scheduled Run

```bash
# 1. Get failure summary
af runs diagnose my_dag scheduled__2025-01-14T...

# 2. Get error from failed task
af tasks logs my_dag scheduled__2025-01-14T... failed_task_id

# 3. [Fix the issue]

# 4. Retest
af runs trigger-wait my_dag
```

### Scenario 5: Test with Custom Configuration

```bash
af runs trigger-wait my_dag --conf '{"env": "staging", "batch_size": 100}' --timeout 600
```

### Scenario 6: Long-Running DAG

```bash
# Wait up to 1 hour
af runs trigger-wait my_dag --timeout 3600

# If timed out, check current state
af runs get my_dag manual__2025-01-14T...
```

---

## Debugging Tips

### Common Error Patterns

**Connection Refused / Timeout:**
- Check `af config connections` for correct host/port
- Verify network connectivity to external system
- Check if connection credentials are correct

**ModuleNotFoundError:**
- Package missing from `requirements.txt`
- After adding, may need environment restart

**PermissionError:**
- Check IAM roles, database grants, API keys
- Verify connection has correct credentials

**Task Timeout:**
- Query or operation taking too long
- Consider adding timeout parameter to task
- Optimize underlying query/operation

### Reading Task Logs

Task logs typically show:
1. Task start timestamp
2. Any print/log statements from task code
3. Return value (for @task decorated functions)
4. Exception + full stack trace (if failed)
5. Task end timestamp and duration

**Focus on the exception at the bottom of failed task logs.**

### On Astro

Astro deployments support environment promotion, which helps structure your testing workflow:

- **Dev deployment**: Test DAGs freely with `astro deploy --dags` for fast iteration
- **Staging deployment**: Run integration tests against production-like data
- **Production deployment**: Deploy only after validation in lower environments
- Use separate Astro deployments for each environment and promote code through them

---

## Related Skills

- **authoring-dags**: For creating new DAGs (includes validation before testing)
- **debugging-dags**: For general Airflow troubleshooting
- **deploying-airflow**: For deploying DAGs to production after testing
tracing-downstream-lineage4.86 KB

View saved version →

---
name: tracing-downstream-lineage
description: Trace downstream data lineage and impact analysis. Use when the user asks what depends on this data, what breaks if something changes, downstream dependencies, or needs to assess change risk before modifying a table or DAG.
---

# Downstream Lineage: Impacts

Answer the critical question: "What breaks if I change this?"

Use this BEFORE making changes to understand the blast radius.

## Impact Analysis

### Step 1: Identify Direct Consumers

Find everything that reads from this target:

**For Tables:**

1. **Search DAG source code**: Look for DAGs that SELECT from this table
   - Use `af dags list` to get all DAGs
   - Use `af dags source <dag_id>` to search for table references
   - Look for: `FROM target_table`, `JOIN target_table`

2. **Check for dependent views**:
   ```sql
   -- Snowflake
   SELECT * FROM information_schema.view_table_usage
   WHERE table_name = '<target_table>'

   -- Or check SHOW VIEWS and search definitions
   ```

3. **Look for BI tool connections**:
   - Dashboards often query tables directly
   - Check for common BI patterns in table naming (rpt_, dashboard_)

### On Astro

If you're running on Astro, the **Lineage tab** in the Astro UI provides visual dependency graphs across DAGs and datasets, making downstream impact analysis faster. It shows which DAGs consume a given dataset and their current status, reducing the need for manual source code searches.

**For DAGs:**

1. **Check what the DAG produces**: Use `af dags source <dag_id>` to find output tables
2. **Then trace those tables' consumers** (recursive)

### Step 2: Build Dependency Tree

Map the full downstream impact:

```
SOURCE: fct.orders
    |
    +-- TABLE: agg.daily_sales --> Dashboard: Executive KPIs
    |       |
    |       +-- TABLE: rpt.monthly_summary --> Email: Monthly Report
    |
    +-- TABLE: ml.order_features --> Model: Demand Forecasting
    |
    +-- DIRECT: Looker Dashboard "Sales Overview"
```

### Step 3: Categorize by Criticality

**Critical** (breaks production):
- Production dashboards
- Customer-facing applications
- Automated reports to executives
- ML models in production
- Regulatory/compliance reports

**High** (causes significant issues):
- Internal operational dashboards
- Analyst workflows
- Data science experiments
- Downstream ETL jobs

**Medium** (inconvenient):
- Ad-hoc analysis tables
- Development/staging copies
- Historical archives

**Low** (minimal impact):
- Deprecated tables
- Unused datasets
- Test data

### Step 4: Assess Change Risk

For the proposed change, evaluate:

**Schema Changes** (adding/removing/renaming columns):
- Which downstream queries will break?
- Are there SELECT * patterns that will pick up new columns?
- Which transformations reference the changing columns?

**Data Changes** (values, volumes, timing):
- Will downstream aggregations still be valid?
- Are there NULL handling assumptions that will break?
- Will timing changes affect SLAs?

**Deletion/Deprecation**:
- Full dependency tree must be migrated first
- Communication needed for all stakeholders

### Step 5: Find Stakeholders

Identify who owns downstream assets:

1. **DAG owners**: Check `owners` field in DAG definitions
2. **Dashboard owners**: Usually in BI tool metadata
3. **Team ownership**: Look for team naming patterns or documentation

## Output: Impact Report

### Summary
"Changing `fct.orders` will impact X tables, Y DAGs, and Z dashboards"

### Impact Diagram
```
                    +--> [agg.daily_sales] --> [Executive Dashboard]
                    |
[fct.orders] -------+--> [rpt.order_details] --> [Ops Team Email]
                    |
                    +--> [ml.features] --> [Demand Model]
```

### Detailed Impacts

| Downstream | Type | Criticality | Owner | Notes |
|------------|------|-------------|-------|-------|
| agg.daily_sales | Table | Critical | data-eng | Updated hourly |
| Executive Dashboard | Dashboard | Critical | analytics | CEO views daily |
| ml.order_features | Table | High | ml-team | Retraining weekly |

### Risk Assessment

| Change Type | Risk Level | Mitigation |
|-------------|------------|------------|
| Add column | Low | No action needed |
| Rename column | High | Update 3 DAGs, 2 dashboards |
| Delete column | Critical | Full migration plan required |
| Change data type | Medium | Test downstream aggregations |

### Recommended Actions

Before making changes:
1. [ ] Notify owners: @data-eng, @analytics, @ml-team
2. [ ] Update downstream DAG: `transform_daily_sales`
3. [ ] Test dashboard: Executive KPIs
4. [ ] Schedule change during low-impact window

### Related Skills
- Trace where data comes from: **tracing-upstream-lineage** skill
- Check downstream freshness: **checking-freshness** skill
- Debug any broken DAGs: **debugging-dags** skill
- Add manual lineage annotations: **annotating-task-lineage** skill
- Build custom lineage extractors: **creating-openlineage-extractors** skill
tracing-upstream-lineage4.39 KB

View saved version →

---
name: tracing-upstream-lineage
description: Trace upstream data lineage. Use when the user asks where data comes from, what feeds a table, upstream dependencies, data sources, or needs to understand data origins.
---

# Upstream Lineage: Sources

Trace the origins of data - answer "Where does this data come from?"

## Lineage Investigation

### Step 1: Identify the Target Type

Determine what we're tracing:
- **Table**: Trace what populates this table
- **Column**: Trace where this specific column comes from
- **DAG**: Trace what data sources this DAG reads from

### Step 2: Find the Producing DAG

Tables are typically populated by Airflow DAGs. Find the connection:

1. **Search DAGs by name**: Use `af dags list` and look for DAG names matching the table name
   - `load_customers` -> `customers` table
   - `etl_daily_orders` -> `orders` table

2. **Explore DAG source code**: Use `af dags source <dag_id>` to read the DAG definition
   - Look for INSERT, MERGE, CREATE TABLE statements
   - Find the target table in the code

3. **Check DAG tasks**: Use `af tasks list <dag_id>` to see what operations the DAG performs

### On Astro

If you're running on Astro, the **Lineage tab** in the Astro UI provides visual lineage exploration across DAGs and datasets. Use it to quickly trace upstream dependencies without manually searching DAG source code.

### On OSS Airflow

Use DAG source code and task logs to trace lineage (no built-in cross-DAG UI).

### Step 3: Trace Data Sources

From the DAG code, identify source tables and systems:

**SQL Sources** (look for FROM clauses):
```python
# In DAG code:
SELECT * FROM source_schema.source_table  # <- This is an upstream source
```

**External Sources** (look for connection references):
- `S3Operator` -> S3 bucket source
- `PostgresOperator` -> Postgres database source
- `SalesforceOperator` -> Salesforce API source
- `HttpOperator` -> REST API source

**File Sources**:
- CSV/Parquet files in object storage
- SFTP drops
- Local file paths

### Step 4: Build the Lineage Chain

Recursively trace each source:

```
TARGET: analytics.orders_daily
    ^
    +-- DAG: etl_daily_orders
            ^
            +-- SOURCE: raw.orders (table)
            |       ^
            |       +-- DAG: ingest_orders
            |               ^
            |               +-- SOURCE: Salesforce API (external)
            |
            +-- SOURCE: dim.customers (table)
                    ^
                    +-- DAG: load_customers
                            ^
                            +-- SOURCE: PostgreSQL (external DB)
```

### Step 5: Check Source Health

For each upstream source:
- **Tables**: Check freshness with the **checking-freshness** skill
- **DAGs**: Check recent run status with `af dags stats`
- **External systems**: Note connection info from DAG code

## Lineage for Columns

When tracing a specific column:

1. Find the column in the target table schema
2. Search DAG source code for references to that column name
3. Trace through transformations:
   - Direct mappings: `source.col AS target_col`
   - Transformations: `COALESCE(a.col, b.col) AS target_col`
   - Aggregations: `SUM(detail.amount) AS total_amount`

## Output: Lineage Report

### Summary
One-line answer: "This table is populated by DAG X from sources Y and Z"

### Lineage Diagram
```
[Salesforce] --> [raw.opportunities] --> [stg.opportunities] --> [fct.sales]
                        |                        |
                   DAG: ingest_sfdc         DAG: transform_sales
```

### Source Details

| Source | Type | Connection | Freshness | Owner |
|--------|------|------------|-----------|-------|
| raw.orders | Table | Internal | 2h ago | data-team |
| Salesforce | API | salesforce_conn | Real-time | sales-ops |

### Transformation Chain
Describe how data flows and transforms:
1. Raw data lands in `raw.orders` via Salesforce API sync
2. DAG `transform_orders` cleans and dedupes into `stg.orders`
3. DAG `build_order_facts` joins with dimensions into `fct.orders`

### Data Quality Implications
- Single points of failure?
- Stale upstream sources?
- Complex transformation chains that could break?

### Related Skills
- Check source freshness: **checking-freshness** skill
- Debug source DAG: **debugging-dags** skill
- Trace downstream impacts: **tracing-downstream-lineage** skill
- Add manual lineage annotations: **annotating-task-lineage** skill
- Build custom lineage extractors: **creating-openlineage-extractors** skill
troubleshooting-astro-deployments8.35 KB

View saved version →

---
name: troubleshooting-astro-deployments
description: Troubleshoot Astronomer production deployments with Astro CLI. Use when investigating deployment issues, viewing production logs, analyzing failures, or managing deployment environment variables.
---

# Astro Deployment Troubleshooting

This skill helps you diagnose and troubleshoot production Astronomer deployments using the Astro CLI.

> **For deployment management**, see the **managing-astro-deployments** skill.
> **For local development**, see the **managing-astro-local-env** skill.

---

## Quick Health Check

Start with these commands to get an overview:

```bash
# 1. List deployments to find target
astro deployment list

# 2. Get deployment overview
astro deployment inspect <DEPLOYMENT_ID>

# 3. Check for errors
astro deployment logs <DEPLOYMENT_ID> --error -c 50
```

---

## Viewing Deployment Logs

Use `-c` to control log count (default: 500). Log flags cannot be combined — use one component or level flag per command.

### Component-Specific Logs

View logs from specific Airflow components:

```bash
# Scheduler logs (DAG processing, task scheduling)
astro deployment logs <DEPLOYMENT_ID> --scheduler -c 50

# Worker logs (task execution)
astro deployment logs <DEPLOYMENT_ID> --workers -c 30

# Webserver logs (UI access, health checks)
astro deployment logs <DEPLOYMENT_ID> --webserver -c 30

# Triggerer logs (deferrable operators)
astro deployment logs <DEPLOYMENT_ID> --triggerer -c 30
```

### Log Level Filtering

Filter by severity:

```bash
# Error logs only (most useful for troubleshooting)
astro deployment logs <DEPLOYMENT_ID> --error -c 30

# Warning logs
astro deployment logs <DEPLOYMENT_ID> --warn -c 50

# Info-level logs
astro deployment logs <DEPLOYMENT_ID> --info -c 50
```

### Search Logs

Search for specific keywords:

```bash
# Search for specific error
astro deployment logs <DEPLOYMENT_ID> --keyword "ConnectionError"

# Search for specific DAG
astro deployment logs <DEPLOYMENT_ID> --keyword "my_dag_name" -c 100

# Find import errors
astro deployment logs <DEPLOYMENT_ID> --error --keyword "ImportError"

# Find task failures
astro deployment logs <DEPLOYMENT_ID> --error --keyword "Task failed"
```

---

## Complete Investigation Workflow

### Step 1: Identify the Problem

```bash
# List deployments with status
astro deployment list

# Get deployment details
astro deployment inspect <DEPLOYMENT_ID>
```

Look for:
- Status: HEALTHY vs UNHEALTHY
- Runtime version compatibility
- Resource limits (CPU, memory)
- Recent deployment timestamp

### Step 2: Check Error Logs

```bash
# Start with errors
astro deployment logs <DEPLOYMENT_ID> --error -c 50
```

Look for:
- Recurring error patterns
- Specific DAGs failing repeatedly
- Import errors or syntax errors
- Connection or credential errors

### Step 3: Review Scheduler Logs

```bash
# Check DAG processing
astro deployment logs <DEPLOYMENT_ID> --scheduler -c 30
```

Look for:
- DAG parse errors
- Scheduling delays
- Task queueing issues

### Step 4: Check Worker Logs

```bash
# Check task execution
astro deployment logs <DEPLOYMENT_ID> --workers -c 30
```

Look for:
- Task execution failures
- Resource exhaustion
- Timeout errors

### Step 5: Verify Configuration

```bash
# Check environment variables
astro deployment variable list --deployment-id <DEPLOYMENT_ID>

# Verify deployment settings
astro deployment inspect <DEPLOYMENT_ID>
```

Look for:
- Missing or incorrect environment variables
- Secrets configuration (AIRFLOW__SECRETS__BACKEND)
- Connection configuration

---

## Common Investigation Patterns

### Recurring DAG Failures

Follow the complete investigation workflow above, then narrow to the specific DAG:

```bash
astro deployment logs <DEPLOYMENT_ID> --keyword "my_dag_name" -c 100
```

### Resource Issues

```bash
# 1. Check deployment resource allocation
astro deployment inspect <DEPLOYMENT_ID>
# Look for: resource_quota_cpu, resource_quota_memory
# Worker queue: max_worker_count, worker_type

# 2. Check for worker scaling issues
astro deployment logs <DEPLOYMENT_ID> --workers -c 50

# 3. Look for out-of-memory errors
astro deployment logs <DEPLOYMENT_ID> --error --keyword "memory"
```

### Configuration Problems

```bash
# 1. Review environment variables
astro deployment variable list --deployment-id <DEPLOYMENT_ID>

# 2. Check for secrets backend configuration
# Look for: AIRFLOW__SECRETS__BACKEND, AIRFLOW__SECRETS__BACKEND_KWARGS

# 3. Verify deployment settings
astro deployment inspect <DEPLOYMENT_ID>

# 4. Check webserver logs for auth issues
astro deployment logs <DEPLOYMENT_ID> --webserver -c 30
```

### Import Errors

```bash
# 1. Find import errors
astro deployment logs <DEPLOYMENT_ID> --error --keyword "ImportError"

# 2. Check scheduler for parse failures
astro deployment logs <DEPLOYMENT_ID> --scheduler --keyword "Failed to import" -c 50

# 3. Verify dependencies were deployed
astro deployment inspect <DEPLOYMENT_ID>
# Check: current_tag, last deployment timestamp
```

---

## Environment Variables Management

### List Variables

```bash
# List all variables for deployment
astro deployment variable list --deployment-id <DEPLOYMENT_ID>

# Find specific variable
astro deployment variable list --deployment-id <DEPLOYMENT_ID> --key AWS_REGION

# Export variables to file
astro deployment variable list --deployment-id <DEPLOYMENT_ID> --save --env .env.backup
```

### Create Variables

```bash
# Create regular variable
astro deployment variable create --deployment-id <DEPLOYMENT_ID> \
  --key API_ENDPOINT \
  --value https://api.example.com

# Create secret (masked in UI and logs)
astro deployment variable create --deployment-id <DEPLOYMENT_ID> \
  --key API_KEY \
  --value secret123 \
  --secret
```

### Update Variables

```bash
# Update existing variable
astro deployment variable update --deployment-id <DEPLOYMENT_ID> \
  --key API_KEY \
  --value newsecret
```

### Delete Variables

```bash
# Delete variable
astro deployment variable delete --deployment-id <DEPLOYMENT_ID> --key OLD_KEY
```

**Note**: Variables are available to DAGs as environment variables. Changes require no redeployment.

---

## Key Metrics from `deployment inspect`

Focus on these fields when troubleshooting:

- **status**: HEALTHY vs UNHEALTHY
- **runtime_version**: Airflow version compatibility
- **scheduler_size/scheduler_count**: Scheduler capacity
- **executor**: CELERY, KUBERNETES, or LOCAL
- **worker_queues**: Worker scaling limits and types
  - `min_worker_count`, `max_worker_count`
  - `worker_concurrency`
  - `worker_type` (resource class)
- **resource_quota_cpu/memory**: Overall resource limits
- **dag_deploy_enabled**: Whether DAG-only deploys work
- **current_tag**: Last deployment version
- **is_high_availability**: Redundancy enabled

---

## Investigation Best Practices

1. **Always start with error logs** - Most obvious failures appear here
2. **Check error logs for patterns** - Same DAG failing repeatedly? Timing patterns?
3. **Component-specific troubleshooting**:
   - Worker logs → task execution details
   - Scheduler logs → DAG processing and scheduling
   - Webserver logs → UI issues and health checks
   - Triggerer logs → deferrable operator issues
4. **Use `--keyword` for targeted searches** - More efficient than reading all logs
5. **The `inspect` command is your health dashboard** - Check it first
6. **Environment variables in `inspect` output** - May reveal configuration issues
7. **Log count default is 500** - Adjust with `-c` based on needs
8. **Don't forget to check deployment time** - Recent deploy might have introduced issue

---

## Troubleshooting Quick Reference

| Symptom | Command |
|---------|---------|
| Deployment shows UNHEALTHY | `astro deployment inspect <ID>` + `--error` logs |
| DAG not appearing | `--error` logs for import errors, check `--scheduler` logs |
| Tasks failing | `--workers` logs + search for DAG with `--keyword` |
| Slow scheduling | `--scheduler` logs + check `inspect` for scheduler resources |
| UI not responding | `--webserver` logs |
| Connection issues | Check variables, search logs for connection name |
| Import errors | `--error --keyword "ImportError"` + `--scheduler` logs |
| Out of memory | `inspect` for resources + `--workers --keyword "memory"` |

---

## Related Skills

- **managing-astro-deployments**: Create, update, delete deployments, deploy code
- **managing-astro-local-env**: Manage local Airflow development environment
- **setting-up-astro-project**: Initialize and configure Astro projects
warehouse-init10.3 KB

View saved version →

---
name: warehouse-init
description: Initialize warehouse schema discovery. Generates .astro/warehouse.md with all table metadata for instant lookups. Run once per project, refresh when schema changes. Use when user says "/astronomer-data:warehouse-init" or asks to set up data discovery.
---

# Initialize Warehouse Schema

Generate a comprehensive, user-editable schema reference file for the data warehouse.

**All CLI commands below are relative to this skill's directory.** Before running any `scripts/cli.py` command, `cd` to the directory containing this file.

## What This Does

1. Discovers all databases, schemas, tables, and columns from the warehouse
2. **Enriches with codebase context** (dbt models, gusty SQL, schema docs)
3. Records row counts and identifies large tables
4. Generates `.astro/warehouse.md` - a version-controllable, team-shareable reference
5. Enables instant concept→table lookups without warehouse queries

## Process

### Step 1: Read Warehouse Configuration

```bash
cat ~/.astro/agents/warehouse.yml
```

Get the list of databases to discover (e.g., `databases: [HQ, ANALYTICS, RAW]`).

### Step 2: Search Codebase for Context (Parallel)

**Launch a subagent to find business context in code:**

```
Task(
    subagent_type="Explore",
    prompt="""
    Search for data model documentation in the codebase:

    1. dbt models: **/models/**/*.yml, **/schema.yml
       - Extract table descriptions, column descriptions
       - Note primary keys and tests

    2. Gusty/declarative SQL: **/dags/**/*.sql with YAML frontmatter
       - Parse frontmatter for: description, primary_key, tests
       - Note schema mappings

    3. AGENTS.md or CLAUDE.md files with data layer documentation

    Return a mapping of:
      table_name -> {description, primary_key, important_columns, layer}
    """
)
```

### Step 3: Parallel Warehouse Discovery

**Launch one subagent per database** using the Task tool:

```
For each database in configured_databases:
    Task(
        subagent_type="general-purpose",
        prompt="""
        Discover all metadata for database {DATABASE}.

        Use the CLI to run SQL queries:
        uv run scripts/cli.py exec "df = run_sql('...')"
        uv run scripts/cli.py exec "print(df)"

        1. Query schemas:
           SELECT SCHEMA_NAME FROM {DATABASE}.INFORMATION_SCHEMA.SCHEMATA

        2. Query tables with row counts:
           SELECT TABLE_SCHEMA, TABLE_NAME, ROW_COUNT, COMMENT
           FROM {DATABASE}.INFORMATION_SCHEMA.TABLES
           ORDER BY TABLE_SCHEMA, TABLE_NAME

        3. For important schemas (MODEL_*, METRICS_*, MART_*), query columns:
           SELECT TABLE_NAME, COLUMN_NAME, DATA_TYPE, COMMENT
           FROM {DATABASE}.INFORMATION_SCHEMA.COLUMNS
           WHERE TABLE_SCHEMA = 'X'

        Return a structured summary:
        - Database name
        - List of schemas with table counts
        - For each table: name, row_count, key columns
        - Flag any tables with >100M rows as "large"
        """
    )
```

**Run all subagents in parallel** (single message with multiple Task calls).

### Step 4: Discover Categorical Value Families

For key categorical columns (like OPERATOR, STATUS, TYPE, FEATURE), discover value families:

```bash
uv run scripts/cli.py exec "df = run_sql('''
SELECT DISTINCT column_name, COUNT(*) as occurrences
FROM table
WHERE column_name IS NOT NULL
GROUP BY column_name
ORDER BY occurrences DESC
LIMIT 50
''')"
uv run scripts/cli.py exec "print(df)"
```

Group related values into families by common prefix/suffix (e.g., `Export*` for ExportCSV, ExportJSON, ExportParquet).

### Step 5: Merge Results

Combine warehouse metadata + codebase context:

1. **Quick Reference table** - concept → table mappings (pre-populated from code if found)
2. **Categorical Columns** - value families for key filter columns
3. **Database sections** - one per database
4. **Schema subsections** - tables grouped by schema
5. **Table details** - columns, row counts, **descriptions from code**, warnings

### Step 6: Generate warehouse.md

Write the file to:
- `.astro/warehouse.md` (default - project-specific, version-controllable)
- `~/.astro/agents/warehouse.md` (if `--global` flag)

## Output Format

```markdown
# Warehouse Schema

> Generated by `/astronomer-data:warehouse-init` on {DATE}. Edit freely to add business context.

## Quick Reference

| Concept | Table | Key Column | Date Column |
|---------|-------|------------|-------------|
| customers | HQ.MODEL_ASTRO.ORGANIZATIONS | ORG_ID | CREATED_AT |
<!-- Add your concept mappings here -->

## Categorical Columns

When filtering on these columns, explore value families first (values often have variants):

| Table | Column | Value Families |
|-------|--------|----------------|
| {TABLE} | {COLUMN} | `{PREFIX}*` ({VALUE1}, {VALUE2}, ...) |
<!-- Populated by /astronomer-data:warehouse-init from actual warehouse data -->

## Data Layer Hierarchy

Query downstream first: `reporting` > `mart_*` > `metric_*` > `model_*` > `IN_*`

| Layer | Prefix | Purpose |
|-------|--------|---------|
| Reporting | `reporting.*` | Dashboard-optimized |
| Mart | `mart_*` | Combined analytics |
| Metric | `metric_*` | KPIs at various grains |
| Model | `model_*` | Cleansed sources of truth |
| Raw | `IN_*` | Source data - avoid |

## {DATABASE} Database

### {SCHEMA} Schema

#### {TABLE_NAME}
{DESCRIPTION from code if found}

| Column | Type | Description |
|--------|------|-------------|
| COL1 | VARCHAR | {from code or inferred} |

- **Rows:** {ROW_COUNT}
- **Key column:** {PRIMARY_KEY from code or inferred}
{IF ROW_COUNT > 100M: - **⚠️ WARNING:** Large table - always add date filters}

## Relationships

```
{Inferred relationships based on column names like *_ID}
```
```

## Command Options

| Option | Effect |
|--------|--------|
| `/astronomer-data:warehouse-init` | Generate .astro/warehouse.md |
| `/astronomer-data:warehouse-init --refresh` | Regenerate, preserving user edits |
| `/astronomer-data:warehouse-init --database HQ` | Only discover specific database |
| `/astronomer-data:warehouse-init --global` | Write to ~/.astro/agents/ instead |

### Step 7: Pre-populate Cache

After generating warehouse.md, populate the concept cache:

```bash
uv run scripts/cli.py concept import -p .astro/warehouse.md
uv run scripts/cli.py concept learn customers HQ.MART_CUST.CURRENT_ASTRO_CUSTS -k ACCT_ID
```

### Step 8: Offer CLAUDE.md Integration (Ask User)

**Ask the user:**

> Would you like to add the Quick Reference table to your CLAUDE.md file?
>
> This ensures the schema mappings are always in context for data queries, improving accuracy from ~25% to ~100% for complex queries.
>
> Options:
> 1. **Yes, add to CLAUDE.md** (Recommended) - Append Quick Reference section
> 2. **No, skip** - Use warehouse.md and cache only

**If user chooses Yes:**

1. Check if `.claude/CLAUDE.md` or `CLAUDE.md` exists
2. If exists, append the Quick Reference section (avoid duplicates)
3. If not exists, create `.claude/CLAUDE.md` with just the Quick Reference

**Quick Reference section to add:**

```markdown
## Data Warehouse Quick Reference

When querying the warehouse, use these table mappings:

| Concept | Table | Key Column | Date Column |
|---------|-------|------------|-------------|
{rows from warehouse.md Quick Reference}

**Large tables (always filter by date):** {list tables with >100M rows}

> Auto-generated by `/astronomer-data:warehouse-init`. Run `/astronomer-data:warehouse-init --refresh` to update.
```
**If yes:** Append the Quick Reference section to `.claude/CLAUDE.md` or `CLAUDE.md`.

## After Generation

Tell the user:

```
Generated .astro/warehouse.md

Summary:
  - {N} databases, {N} schemas, {N} tables
  - {N} tables enriched with code descriptions
  - {N} concepts cached for instant lookup

Next steps:
  1. Edit .astro/warehouse.md to add business context
  2. Commit to version control
  3. Run /astronomer-data:warehouse-init --refresh when schema changes
```

## Refresh Behavior

When `--refresh` is specified:

1. Read existing warehouse.md
2. Preserve all HTML comments (`<!-- ... -->`)
3. Preserve Quick Reference table entries (user-added)
4. Preserve user-added descriptions
5. Update row counts and add new tables
6. Mark removed tables with `<!-- REMOVED -->` comment

## Cache Staleness & Schema Drift

The runtime cache has a **7-day TTL** by default. After 7 days, cached entries expire and will be re-discovered on next use.

### When to Refresh

Run `/astronomer-data:warehouse-init --refresh` when:
- **Schema changes**: Tables added, renamed, or removed
- **Column changes**: New columns added or types changed
- **After deployments**: If your data pipeline deploys schema migrations
- **Weekly**: As a good practice, even if no known changes

### Signs of Stale Cache

Watch for these indicators:
- Queries fail with "table not found" errors
- Results seem wrong or outdated
- New tables aren't being discovered

### Manual Cache Reset

If you suspect cache issues:

```bash
uv run scripts/cli.py cache status
uv run scripts/cli.py cache clear --stale-only
uv run scripts/cli.py cache clear
```

## Codebase Patterns Recognized

| Pattern | Source | What We Extract |
|---------|--------|-----------------|
| `**/models/**/*.yml` | dbt | table/column descriptions, tests |
| `**/dags/**/*.sql` | gusty | YAML frontmatter (description, primary_key) |
| `AGENTS.md`, `CLAUDE.md` | docs | data layer hierarchy, conventions |
| `**/docs/**/*.md` | docs | business context |

## Example Session

```
User: /astronomer-data:warehouse-init

Agent:
→ Reading warehouse configuration...
→ Found 1 warehouse with databases: HQ, PRODUCT

→ Searching codebase for data documentation...
  Found: AGENTS.md with data layer hierarchy
  Found: 45 SQL files with YAML frontmatter in dags/declarative/

→ Launching parallel warehouse discovery...
  [Database: HQ] Discovering schemas...
  [Database: PRODUCT] Discovering schemas...

→ HQ: Found 29 schemas, 401 tables
→ PRODUCT: Found 1 schema, 0 tables

→ Merging warehouse metadata with code context...
  Enriched 45 tables with descriptions from code

→ Generated .astro/warehouse.md

Summary:
  - 2 databases
  - 30 schemas
  - 401 tables
  - 45 tables enriched with code descriptions
  - 8 large tables flagged (>100M rows)

Next steps:
  1. Review .astro/warehouse.md
  2. Add concept mappings to Quick Reference
  3. Commit to version control
  4. Run /astronomer-data:warehouse-init --refresh when schema changes
```

Referenced files: 7

Technical details
First seen
Sep 30, 2026 · 22:02 UTC
Last seen
Oct 1, 2026 · 12:00 UTC
Collection status
Collected

plugins_6ab2f25e4928819184294ebadcbe38ab

Download listing JSON