DAG Authoring Approaches

from __future__ import annotations

import json
from datetime import datetime, timedelta

from airflow.decorators import dag, task
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.http.hooks.http import HttpHook


@dag(
    dag_id="taskflow_etl",
    start_date=datetime(2024, 1, 1),
    schedule="@daily",
    catchup=False,
    default_args={
        "retries": 3,
        "retry_delay": timedelta(minutes=2),
        "execution_timeout": timedelta(minutes=30),
    },
    tags=["etl"],
    doc_md="Daily ETL pipeline using TaskFlow API.",
)
def taskflow_etl():

    @task()
    def extract(**context) -> list[dict]:
        """Fetch data using Airflow HTTP connection."""
        hook = HttpHook(http_conn_id="api_default", method="GET")
        ds = context["ds"]
        response = hook.run(endpoint=f"/data?date={ds}")
        return response.json()

    @task()
    def transform(records: list[dict]) -> list[dict]:
        """Filter active records — pure function, no I/O."""
        return [r for r in records if r.get("status") == "active"]

    @task()
    def load(records: list[dict], **context) -> None:
        """Write to S3 with partitioned key for idempotency."""
        ds = context["ds"]
        hook = S3Hook(aws_conn_id="aws_default")
        hook.load_string(
            string_data=json.dumps(records),
            key=f"daily/{ds}/data.json",
            bucket_name="my-bucket",
            replace=True,
        )

    # Dependencies inferred from function arguments
    data = extract()
    filtered = transform(data)
    load(filtered)


taskflow_etl()

if __name__ == "__main__":
    taskflow_etl().test()

from airflow.sdk import Asset, asset


@asset(
    uri="s3://my-bucket/clean/customers.parquet",
    schedule="@daily",
)
def clean_customers(self, context):
    """Produce a clean customers dataset."""
    # self.uri contains the asset URI
    # Task logic here
    pass

from airflow.sdk import Asset, asset
from airflow.decorators import dag, task

# Producer: creates asset event on success
customers_asset = Asset("s3://my-bucket/clean/customers.parquet")


@asset(uri="s3://my-bucket/clean/customers.parquet", schedule="@daily")
def produce_customers(self, context):
    pass


# Consumer: triggered when asset is updated
@dag(
    dag_id="consume_customers",
    schedule=[customers_asset],  # Triggered by asset update
    catchup=False,
)
def consume_customers():
    @task()
    def process():
        pass

    process()


consume_customers()

from airflow.sdk import Asset
from airflow.decorators import dag, task

my_asset = Asset("s3://bucket/data")


@dag(dag_id="producer", schedule="@daily")
def producer():
    @task(outlets=[my_asset])  # Marks this task as producing the asset
    def produce(**context):
        pass

    produce()


producer()

a = Asset("sales")
b = Asset("sales")
# a and b are different Python objects but THE SAME asset in Airflow

@asset(schedule="@daily")
def daily_sales():
    pass

# These two are identical — same asset:
@dag(schedule=daily_sales)          # implicit: function reference
@dag(schedule=Asset("daily_sales")) # explicit: Asset object with same name

from datetime import datetime, timedelta

from airflow.sdk import DAG
from airflow.providers.standard.operators.python import PythonOperator
from airflow.providers.standard.operators.bash import BashOperator


def _extract(**context):
    pass


def _transform(**context):
    ti = context["ti"]
    data = ti.xcom_pull(task_ids="extract")
    # Manual XCom pull required
    return [r for r in data if r["status"] == "active"]


with DAG(
    dag_id="traditional_etl",
    start_date=datetime(2024, 1, 1),
    schedule="@daily",
    catchup=False,
    default_args={"retries": 3, "retry_delay": timedelta(minutes=2)},
) as dag:

    extract = PythonOperator(
        task_id="extract",
        python_callable=_extract,
    )

    transform = PythonOperator(
        task_id="transform",
        python_callable=_transform,
    )

    load = BashOperator(
        task_id="load",
        bash_command="echo 'loading data'",
    )

    extract >> transform >> load

from __future__ import annotations

from datetime import timedelta

from airflow.decorators import dag, task
from airflow.providers.openai.hooks.openai import OpenAIHook


@dag(
    dag_id="genai_inference",
    schedule=None,  # API-triggered only
    catchup=False,
    max_consecutive_failed_dag_runs=10,
    default_args={
        "retries": 2,
        "retry_delay": timedelta(minutes=1),
        "execution_timeout": timedelta(minutes=5),
    },
    tags=["inference", "genai"],
)
def genai_inference():

    @task()
    def process_prompt(**context) -> dict:
        """Read user prompt from dag_run.conf and call LLM."""
        conf = context["dag_run"].conf or {}
        user_prompt = conf.get("prompt", "")
        hook = OpenAIHook(conn_id="openai_default")
        client = hook.get_conn()
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": user_prompt},
            ],
        )
        return {"result": response.choices[0].message.content}

    process_prompt()


genai_inference()

curl -X POST "https://airflow.example.com/api/v2/dags/genai_inference/dagRuns" \
  -H "Content-Type: application/json" \
  -d '{"conf": {"prompt": "Summarize this quarter sales data"}}'

from airflow.sdk import Asset, AssetWatcher
from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger

sqs_trigger = SqsSensorTrigger(
    sqs_queue="inference-requests",
    aws_conn_id="aws_default",
)

request_event = Asset(
    "sqs://inference-requests",
    watchers=[AssetWatcher(trigger=sqs_trigger)],
)


@dag(schedule=[request_event], catchup=False)
def process_inference_request():
    ...