Apache Airflow in Production

from airflow.sensors.base import BaseSensorOperator
from airflow.utils.decorators import apply_defaults

class SmartS3Sensor(BaseSensorOperator):
    """A batched S3 sensor that satisfies up to N sensor
    instances with a single S3 API call per polling cycle."""

    @apply_defaults
    def __init__(self, bucket, key_prefix, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.bucket = bucket
        self.key_prefix = key_prefix

    def poke(self, context):
        # Check partition completion via Hive metastore
        partition = context["execution_date"].strftime("%Y-%m-%d")
        key = f"{self.key_prefix}/dt={partition}/_SUCCESS"

        from airflow.hooks.S3_hook import S3Hook
        hook = S3Hook(aws_conn_id="aws_default")
        return hook.check_for_key(key, self.bucket)

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
import yaml

def generate_dag_from_config(config):
    """Generates a DAG dynamically from a config dict."""
    dag = DAG(
        dag_id=config["dag_id"],
        schedule_interval=config["schedule"],
        default_args=config["default_args"],
        catchup=False,
    )

    with dag:
        extract = PythonOperator(
            task_id="extract",
            python_callable=extract_data,
            op_kwargs={"source": config["source"]},
        )
        transform = PythonOperator(
            task_id="transform",
            python_callable=transform_data,
            op_kwargs={"transformations": config["transformations"]},
        )
        load = PythonOperator(
            task_id="load",
            python_callable=load_data,
            op_kwargs={"destination": config["destination"]},
        )

        extract >> transform >> load
    return dag

# Generate DAGs from a YAML of pipeline configs. Airbnb has 800+ DAGs
# auto-built from a single YAML, eliminating boilerplate and enforcing
# standards.
for config in load_pipeline_configs():
    globals()[config["dag_id"]] = generate_dag_from_config(config)

from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
from kubernetes.client import models as k8s

# Run a Spark job in an isolated K8s pod
spark_task = KubernetesPodOperator(
    task_id="spark_aggregation",
    name="spark-aggregation",
    namespace="dataeng",
    image="airbnb/spark:3.4",
    cmds=["spark-submit"],
    arguments=[
        "--master", "k8s://https://kubernetes.default.svc",
        "--conf", "spark.executor.instances=20",
        "/app/aggregate.py",
        "--date", "{{ ds }}",
    ],
    resources=k8s.V1ResourceRequirements(
        requests={"cpu": "1", "memory": "4Gi"},
        limits={"cpu": "4", "memory": "16Gi"},
    ),
    secrets=[
        {"key": "key", "name": "aws-credentials", "deploy_target": "AWS_ACCESS_KEY_ID"},
        {"key": "secret", "name": "aws-credentials", "deploy_target": "AWS_SECRET_ACCESS_KEY"},
    ],
    is_delete_operator_pod=True,
    in_cluster=True,
    get_logs=True,
)

apiVersion: airflow.lyft.com/v1
kind: AirflowCell
metadata:
  name: marketplace
spec:
  team: marketplace-data
  ownership: dataeng-marketplace@lyft.com

  scheduler:
    replicas: 2
    resources:
      requests: { cpu: 2, memory: 4Gi }
      limits:   { cpu: 4, memory: 8Gi }

  workers:
    minReplicas: 5
    maxReplicas: 80
    autoscale:
      metric: pending_tasks
      target: 50

  metadataDB:
    engine: aurora-postgres
    shard: cell-marketplace
    storage: 200Gi

  eventBus:
    kafka:
      topic: airflow.events.marketplace
      consumerGroup: airflow-events

  observability:
    openlineage: { enabled: true, endpoint: marquez.internal }
    cost:        { namespace: af-cell-marketplace }