docker-compose.yml

# Adapted from Airflow's docker-compose.yml file:
# https://airflow.apache.org/docs/apache-airflow/2.8.1/docker-compose.yaml
networks:
  default:
    name: airflow-config_default
  cumulus:
    name: cumulus-api_default
    external: true
x-airflow-common:
  # In order to add custom dependencies or upgrade provider packages you can use your extended image.
  # Comment the image line, place your Dockerfile in the directory where you placed the docker-compose.yaml
  # and uncomment the "build" line below, Then run `docker-compose build` to build the images.
  &airflow-common
  #
  # image: apache/airflow:2.8.1-python3.8-usace
  image: airflow-common:2.10.5-python3.12-usace
  build:
    context: _docker/airflow
    args:
      - AIRFLOW_VERSION=2.10.5
      - PYTHON_BASE_IMAGE=python:3.12-slim-bookworm
      - ADDITIONAL_DEV_APT_DEPS=libcurl4-openssl-dev libssl-dev
      - ADDITIONAL_PYTHON_DEPS=celery[sqs] celery[redis] dataretrieval==1.0.12 cwms-python==1.0.3 cwms-cli==0.3.2 shef-parser==1.9.0
      - INSTALL_MYSQL_CLIENT=false
      - INSTALL_MSSQL_CLIENT=false
  environment:
    #
    &airflow-common-env
    #######################################
    # AWS LOCAL MOCK
    #######################################
    AWS_ACCESS_KEY_ID: AKIAIOSFODNN7EXAMPLE
    AWS_SECRET_ACCESS_KEY: wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY
    AWS_DEFAULT_REGION: us-east-1 # aws default env var
    AWS_ENDPOINT_URL_SQS: http://airflow-elasticmq:9324 # aws default env var
    #######################################
    # VARIABLES DEFINED BY AIRFLOW SOFTWARE
    #######################################
    AIRFLOW__CORE__EXECUTOR: CeleryExecutor
    AIRFLOW__CORE__FERNET_KEY: u1rsP4UVZf5_uREVUhAjqYYygN328ASxeNYdy0yd6_Q=
    AIRFLOW__CORE__DAGS_ARE_PAUSED_AT_CREATION: 'true'
    # Prevent the 30+ connections from being loaded
    AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS: False
    #--> AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://airflow:airflow@pgbouncer:5432/airflow
    AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://
    AIRFLOW__DATABASE__SQL_ALCHEMY_POOL_PRE_PING: 'True'
    PGUSER: airflow
    PGPASSWORD: airflow
    PGHOST: pgbouncer
    PGPORT: 5432
    PGDATABASE: airflow
    # ----------
    # CELERY
    # ----------
    # AIRFLOW__CELERY__RESULT_BACKEND: db+postgresql://airflow:airflow@pgbouncer/airflow
    AIRFLOW__CELERY__RESULT_BACKEND: db+postgresql://
    # DB_USER_NAME: airflow
    # DB_USER_PASSWORD: airflow
    DB_HOST: pgbouncer
    DB_PORT: 5432
    # AIRFLOW__CELERY__BROKER_URL: sqs://AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI%2FK7MDENG%2FbPxRfiCYEXAMPLEKEY@elasticmq
    AIRFLOW__CELERY__BROKER_URL: redis://:@redis:6379/0
    AIRFLOW__CELERY__CELERY_CONFIG_OPTIONS: celery_config.CELERY_CONFIG
    # default queue; specified in celery_config.py
    AIRFLOW__OPERATORS__DEFAULT_QUEUE: airflow
    # ----------
    # SCHEDULER
    # ----------
    # How often (in seconds) to scan the DAGs directory for new files. Default=300 (5 minutes).
    AIRFLOW__SCHEDULER__DAG_DIR_LIST_INTERVAL: 60
    # Number of seconds after which a DAG file is parsed.
    # The DAG file is parsed everymin_file_process_interval number of seconds.
    # Updates to DAGs are reflected after this interval.
    # Keeping this number low will increase CPU usage. Default=30
    AIRFLOW__SCHEDULER__MIN_FILE_PROCESS_INTERVAL: 30
    # ----------
    # LOGGING
    # ----------
    AIRFLOW__LOGGING__LOGGING_LEVEL: INFO
    AIRFLOW__LOGGING__REMOTE_LOGGING: TRUE
    AIRFLOW__LOGGING__REMOTE_LOG_CONN_ID: LOCAL_MINIO
    AIRFLOW__LOGGING__REMOTE_BASE_LOG_FOLDER: s3://cwbi-airflow-dev/airflow/logs
    # Whether the local log files for S3 remote logging should be deleted after they are uploaded to the remote location.
    AIRFLOW__LOGGING__DELETE_LOCAL_LOGS: False
    # ----------
    # METRICS
    # ----------
    AIRFLOW__METRICS__STATSD_ON: True
    AIRFLOW__METRICS__STATSD_HOST: airflow-statsd-exporter
    AIRFLOW__METRICS__STATSD_PORT: 9102
    AIRFLOW__METRICS__STATSD_PREFIX: airflow
    # ----------
    # WEBSERVER
    # ----------
    AIRFLOW__WEBSERVER__INSTANCE_NAME: 'Local Dev'
    AIRFLOW__WEBSERVER__NAVBAR_COLOR: '#044e14'
    AIRFLOW__WEBSERVER__NAVBAR_LOGO_TEXT_COLOR: '#fff'
    AIRFLOW__WEBSERVER__NAVBAR_TEXT_COLOR: '#fff'

    ######################
    # VARIABLES WE CREATED
    ######################
    SQS_BASE_URL: http://airflow-elasticmq:9324/queue
    # SQS_IS_ELASTICMQ: true
    # ----------
    # CELERY QUEUES
    # ----------
    CELERY_QUEUE_URL_AIRFLOW: http://airflow-elasticmq:9324/queue/airflow
    CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP: 'True'
    # ----------
    # CONNECTIONS
    # ----------
    AIRFLOW_CONN_SQS: aws://AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI%2FK7MDENG%2FbPxRfiCYEXAMPLEKEY@?region_name=elasticmq&host=http%3A%2F%2Fairflow-elasticmq%3A9324
    # Host is deprecated in connections - use "endpoint_url"
    # AIRFLOW_CONN_LOCAL_MINIO: aws://AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI%2FK7MDENG%2FbPxRfiCYEXAMPLEKEY@?host=http%3A%2F%2Fminio%3A9000&region=us-east-1
    AIRFLOW_CONN_LOCAL_MINIO: aws://AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI%2FK7MDENG%2FbPxRfiCYEXAMPLEKEY@?endpoint_url=http%3A%2F%2Fminio%3A9000&region=us-east-1
    AIRFLOW_CONN_CUMULUS: http://nologin:appkey@api
    # AWS_DEFAULT - Used for localstack mocking
    AIRFLOW_CONN_AWS_DEFAULT: aws://AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI%2FK7MDENG%2FbPxRfiCYEXAMPLEKEY@localstack:4566?region_name=us-east-1
    # ----------
    # VARIABLES
    # ----------
    AIRFLOW_VAR_CWMSDATA: https://cwms-data.usace.army.mil/cwms-data
    AIRFLOW_VAR_S3_BUCKET: castle-data-develop
    AIRFLOW_VAR_CUMULUS_S3_BUCKET: castle-data-develop
    AIRFLOW_VAR_WMES_SHEF_QUEUE_NAME: wmes-shef-service
    # AWS_ENDPOINT_URL_SQS: http://airflow-elasticmq:9324
    # AWS_DEFAULT_REGION: us-east-1
    # How Downloads are Handled (Download or Fire Lambda Function to Handle Download
    DOWNLOAD_OPERATOR_USE_CONNECTION: LOCAL_MINIO
    DOWNLOAD_OPERATOR_USE_LAMBDA: FALSE
    DOWNLOAD_OPERATOR_USE_LAMBDA_NAME: ''
    DOWNLOAD_OPERATOR_USE_LAMBDA_REGION: us-east-1
    # ----------
    # TELEMETRY (back to airflow devs)
    # ----------
    SCARF_ANALYTICS: False
    ###############################
    # Things I'm unsure if we need
    ###############################
    # Use simple http server on scheduler for health checks
    # See https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/logging-monitoring/check-health.html#scheduler-health-check-server
    AIRFLOW__SCHEDULER__ENABLE_HEALTH_CHECK: 'true'
    AIRFLOW__API__AUTH_BACKENDS: 'airflow.api.auth.backend.basic_auth,airflow.api.auth.backend.session'
    ###############################
    # AWS Only things
    ###############################
    # https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/secrets/secrets_manager/index.html#airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
    # AIRFLOW__SECRETS__BACKEND: airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
    # AIRFLOW__SECRETS__BACKEND_KWARGS: '{"connections_prefix": "airflow/connections", "region_name": "us-east-1"}'
    # AIRFLOW__SECRETS__USE_CACHE: 'True'
    # AIRFLOW__SECRETS__CACHE_TTL_SECONDS: 900 # Default 900
  env_file:
    - .env
  volumes:
    # volumes are attached to any container with <<: *airflow-common
    - ./dags:/opt/airflow/dags
    - ./logs:/opt/airflow/logs
    - ./plugins:/opt/airflow/plugins
    - /var/run/docker.sock:/var/run/docker.sock
  user: '${AIRFLOW_UID:-50000}:0'
  depends_on: &airflow-common-depends-on
    airflowdb:
      condition: service_healthy
    pgbouncer:
      condition: service_started
    redis:
      condition: service_healthy

services:
  #############################################################################
  # Postgres Database
  #############################################################################
  airflowdb:
    image: postgres:16
    container_name: airflow-db
    environment:
      POSTGRES_USER: airflow
      POSTGRES_PASSWORD: airflow
      POSTGRES_DB: airflow
    healthcheck:
      test: ['CMD', 'pg_isready', '-U', 'airflow']
      interval: 5s
      retries: 5
    ports:
      - '50432:5432' # host port changed to prevent clash with other app's DB
  #############################################################################
  # Postgres Database
  #############################################################################
  pgbouncer:
    image: pgbouncer:latest
    container_name: airflow-pgbouncer
    build:
      context: _docker/pgbouncer
      args:
        - PGBOUNCER_VERSION=1.24.1
        - PGBOUNCER_DOWNLOAD_URL=https://github.com/pgbouncer/pgbouncer/releases/download/pgbouncer_1_24_1/pgbouncer-1.24.1.tar.gz
    environment:
      DB_HOST: airflowdb
      DB_USER: airflow
      DB_PASSWORD: airflow
      DB_PORT: 5432
      LISTEN_PORT: 5432
      VERBOSE: 0
      QUIET: 0
      POOL_MODE: transaction
  #############################################################################
  # REDIS
  #############################################################################
  redis:
    image: redis:latest
    container_name: airflow-redis
    expose:
      - 6379
    healthcheck:
      test: ['CMD', 'redis-cli', 'ping']
      interval: 10s
      timeout: 30s
      retries: 50
      start_period: 30s
    restart: always
  #############################################################################
  # elasticmq
  #############################################################################
  airflow-elasticmq:
    container_name: airflow-elasticmq
    build:
      context: ./_docker/elasticmq
    ports:
      - '19324:9324'
      - '19325:9325'
    networks:
      - default
    restart: always
  #############################################################################
  # Airflow Init
  #############################################################################
  airflow_init:
    <<: *airflow-common
    container_name: airflow-init
    depends_on:
      - airflowdb
    volumes:
      - ./_utils/variables.json:/opt/airflow/variables.json
    entrypoint: |-
      /bin/sh -c "
      /entrypoint db migrate;
      /entrypoint users create --role Admin --username airflow --email airflow@airflow.com --firstname airflow --lastname airflow --password airflow;
      /entrypoint variables import /opt/airflow/variables.json;
      exit 0;
      "
  #############################################################################
  # Airflow Scheduler
  #############################################################################
  airflow_scheduler:
    <<: *airflow-common
    container_name: airflow-scheduler
    command: scheduler
    healthcheck:
      # test:
      #   [
      #     'CMD-SHELL',
      #     'airflow jobs check --job-type SchedulerJob --hostname "$${HOSTNAME}"',
      #   ]
      test: ['CMD', 'curl', '--fail', 'http://localhost:8974/health']
      interval: 30s
      timeout: 10s
      retries: 5
      start_period: 30s
    restart: always
    depends_on:
      <<: *airflow-common-depends-on
      airflow_init:
        condition: service_completed_successfully
      airflow-elasticmq:
        condition: service_started
    networks:
      - default
      - cumulus
  #############################################################################
  # Airflow Celery Worker
  #############################################################################
  airflow_worker:
    <<: *airflow-common
    container_name: airflow-worker
    command: celery worker
    healthcheck:
      # test:
      #   - 'CMD-SHELL'
      #   - 'celery --app airflow.executors.celery_executor.app inspect ping -d "celery@$${HOSTNAME}"'
      test:
        - 'CMD-SHELL'
        - 'celery --app airflow.providers.celery.executors.celery_executor.app inspect ping -d "celery@$${HOSTNAME}" || celery --app airflow.executors.celery_executor.app inspect ping -d "celery@$${HOSTNAME}"'
      interval: 10s
      timeout: 10s
      retries: 5
      start_period: 30s
    environment:
      <<: *airflow-common-env
      # Required to handle warm shutdown of the celery workers properly
      # See https://airflow.apache.org/docs/docker-stack/entrypoint.html#signal-propagation
      DUMB_INIT_SETSID: '0'
    restart: always
    depends_on:
      <<: *airflow-common-depends-on
      airflow_init:
        condition: service_completed_successfully
      airflow-elasticmq:
        condition: service_started
    networks:
      - default
      - cumulus
  #############################################################################
  # Airflow Triggerer
  #############################################################################
  airflow_triggerer:
    <<: *airflow-common
    container_name: airflow-triggerer
    command: triggerer
    healthcheck:
      test:
        [
          'CMD-SHELL',
          'airflow jobs check --job-type TriggererJob --hostname "$${HOSTNAME}"',
        ]
      interval: 30s
      timeout: 10s
      retries: 5
      start_period: 30s
    restart: always
    depends_on:
      <<: *airflow-common-depends-on
      airflow_init:
        condition: service_completed_successfully
  #############################################################################
  # Airflow User Interface
  #############################################################################
  airflow_ui:
    <<: *airflow-common
    container_name: airflow-ui
    # image: apache/airflow:2.4.3-python3.8-usace
    command: webserver
    ports:
      - '8000:8080'
    environment:
      <<: *airflow-common-env
      AIRFLOW__WEBSERVER__ACCESS_LOGFILE: '/dev/null' # <-- stop spewing http logs to console
    healthcheck:
      test: ['CMD', 'curl', '--fail', 'http://localhost:8080/health']
      interval: 10s
      timeout: 10s
      retries: 5
    restart: always
    depends_on:
      <<: *airflow-common-depends-on
      airflow_init:
        condition: service_completed_successfully
    networks:
      - default
      - cumulus
  #############################################################################
  # Celery Dashboard - local only
  #############################################################################
  airflow_flower:
    <<: *airflow-common
    container_name: airflow-flower
    command: celery flower
    # profiles:
    #   - flower
    ports:
      - '5555:5555'
    healthcheck:
      test: ['CMD', 'curl', '--fail', 'http://localhost:5555/']
      interval: 30s
      timeout: 10s
      retries: 5
      start_period: 30s
    restart: always
    depends_on:
      <<: *airflow-common-depends-on
      airflow_init:
        condition: service_completed_successfully