Advanced Airflow: XComs, Dynamic DAGs, and the TaskFlow API

from datetime import datetime, timedelta
from typing import List, Dict
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.empty import EmptyOperator
from airflow.utils.task_group import TaskGroup
import yaml
import logging

logger = logging.getLogger(__name__)


def create_etl_dag(
    dag_id: str,
    source_config: Dict,
    target_config: Dict,
    schedule: str = "@daily",
) -> DAG:
    """Factory function to create an ETL DAG from configuration."""

    default_args = {
        "owner": "data-engineering",
        "retries": 3,
        "retry_delay": timedelta(minutes=5),
        "execution_timeout": timedelta(hours=1),
    }

    def extract(**context):
        import pandas as pd
        import sqlalchemy

        engine = sqlalchemy.create_engine(source_config["connection_string"])
        query = source_config["query"].replace(
            "{{ds}}", context["ds"]
        )
        df = pd.read_sql(query, engine)
        output_path = f"s3://data-lake/raw/{dag_id}/{context['ds_nodash']}/data.parquet"
        df.to_parquet(output_path, index=False)
        context["ti"].xcom_push(key="input_path", value=output_path)
        context["ti"].xcom_push(key="row_count", value=len(df))
        logger.info(f"Extracted {len(df)} rows to {output_path}")

    def transform(**context):
        import pandas as pd

        input_path = context["ti"].xcom_pull(
            task_ids=f"{dag_id}.etl.extract", key="input_path"
        )
        df = pd.read_parquet(input_path)

        for step in source_config.get("transform_steps", []):
            if step["type"] == "filter":
                df = df.query(step["condition"])
            elif step["type"] == "rename":
                df = df.rename(columns=step["mapping"])
            elif step["type"] == "drop_nulls":
                df = df.dropna(subset=step["columns"])

        output_path = f"s3://data-lake/staging/{dag_id}/{context['ds_nodash']}/data.parquet"
        df.to_parquet(output_path, index=False)
        context["ti"].xcom_push(key="output_path", value=output_path)
        context["ti"].xcom_push(key="transformed_count", value=len(df))

    def load(**context):
        import pandas as pd
        import sqlalchemy

        output_path = context["ti"].xcom_pull(
            task_ids=f"{dag_id}.etl.transform", key="output_path"
        )
        df = pd.read_parquet(output_path)
        engine = sqlalchemy.create_engine(target_config["connection_string"])
        df.to_sql(
            target_config["table"],
            engine,
            if_exists="append",
            index=False,
            chunksize=1000,
        )
        logger.info(f"Loaded {len(df)} rows to {target_config['table']}")

    with DAG(
        dag_id=dag_id,
        default_args=default_args,
        schedule_interval=schedule,
        start_date=datetime(2024, 1, 1),
        catchup=False,
        max_active_runs=1,
        tags=["dynamic", "etl"],
    ) as dag:

        start = EmptyOperator(task_id="start")

        with TaskGroup("etl") as etl_group:
            extract_task = PythonOperator(
                task_id="extract", python_callable=extract
            )
            transform_task = PythonOperator(
                task_id="transform", python_callable=transform
            )
            load_task = PythonOperator(
                task_id="load", python_callable=load
            )
            extract_task >> transform_task >> load_task

        end = EmptyOperator(task_id="end", trigger_rule="none_failed")

        start >> etl_group >> end

    return dag


# Dynamic DAG generation from YAML config
with open("/opt/airflow/config/pipeline_configs.yaml") as f:
    configs = yaml.safe_load(f)

for pipeline in configs["pipelines"]:
    dag = create_etl_dag(
        dag_id=f"dynamic_{pipeline['name']}",
        source_config=pipeline["source"],
        target_config=pipeline["target"],
        schedule=pipeline.get("schedule", "@daily"),
    )
    globals()[dag.dag_id] = dag

from airflow import DAG
from airflow.decorators import task, dag
from datetime import datetime
from typing import List

default_args = {
    "owner": "data-engineering",
    "retries": 2,
    "retry_delay": 300,
}


@dag(
    schedule="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    default_args=default_args,
    tags=["taskflow", "dynamic-mapping"],
)
def dynamic_mapping_etl():

    @task
    def get_table_list() -> List[str]:
        """Dynamically discover tables to process."""
        import sqlalchemy

        engine = sqlalchemy.create_engine("postgresql://...")
        with engine.connect() as conn:
            result = conn.execute(
                "SELECT table_name FROM information_schema.tables "
                "WHERE table_schema = 'public' AND table_name LIKE 'events_%'"
            )
            return [row[0] for row in result]

    @task
    def extract_table(table_name: str) -> str:
        """Extract a single table to a parquet file."""
        import pandas as pd
        import sqlalchemy

        engine = sqlalchemy.create_engine("postgresql://...")
        df = pd.read_sql(f"SELECT * FROM {table_name}", engine)
        output_path = f"s3://raw/{table_name}/latest.parquet"
        df.to_parquet(output_path, index=False)
        return output_path

    @task
    def validate_table(table_name: str, file_path: str) -> dict:
        """Validate extracted data quality."""
        import pandas as pd

        df = pd.read_parquet(file_path)
        null_rate = df.isnull().mean().mean()
        duplicate_rate = 1 - (len(df.drop_duplicates()) / len(df))

        return {
            "table": table_name,
            "row_count": len(df),
            "null_rate": float(null_rate),
            "duplicate_rate": float(duplicate_rate),
            "is_valid": null_rate < 0.1 and duplicate_rate < 0.05,
        }

    @task(trigger_rule="none_failed")
    def aggregate_results(validations: List[dict]) -> dict:
        """Aggregate validation results across all tables."""
        total_tables = len(validations)
        valid_tables = sum(1 for v in validations if v["is_valid"])
        return {
            "total_tables": total_tables,
            "valid_tables": valid_tables,
            "validity_rate": valid_tables / total_tables if total_tables > 0 else 0,
        }

    # Dynamic task mapping
    tables = get_table_list()
    raw_files = extract_table.expand(table_name=tables)
    validations = validate_table.expand(
        table_name=tables, file_path=raw_files
    )
    aggregate_results(validations)


dynamic_mapping_etl()