altool/airflow

1#!/bin/bash2set -euo pipefail3 4# ---- Helpers ----5log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }6 7# Graceful shutdown: forward signals to child processes8trap 'log "Caught signal, shutting down..."; kill -TERM $(jobs -p) 2>/dev/null; exit 0' SIGTERM SIGINT SIGQUIT9 10# ---- Load Configuration ----11log "Loading configuration..."12source /config.sh13validate_config14 15log "Configuration:"16log "  DB backend : $(echo "${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}" | sed -E 's|://[^@]*@|://***:***@|')"17log "  DAG repo   : ${DAG_REPO_URL} (branch: ${DAG_REPO_BRANCH})"18log "  Executor   : ${AIRFLOW__CORE__EXECUTOR}"19log "  Ports      : nginx=${NGINX_PORT}  webserver=${AIRFLOW_WEB_PORT}  refresh=${REFRESH_SERVICE_PORT}"20 21# ---- Ensure directories ----22mkdir -p "${DAGS_DIR}" "${LOGS_DIR}"23 24# ---- Wait for Database (if PostgreSQL) ----25if [[ "${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}" == postgresql* ]]; then26    DB_WAIT_TIMEOUT="${DB_WAIT_TIMEOUT:-60}"27    log "Waiting for PostgreSQL (timeout: ${DB_WAIT_TIMEOUT}s)..."28    elapsed=029    while [ "${elapsed}" -lt "${DB_WAIT_TIMEOUT}" ]; do30        if python3 -c "31import sqlalchemy, sys32try:33    engine = sqlalchemy.create_engine('${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}', connect_args={'connect_timeout': 5})34    with engine.connect() as conn:35        conn.execute(sqlalchemy.text('SELECT 1'))36    sys.exit(0)37except Exception:38    sys.exit(1)39" 2>/dev/null; then40            log "PostgreSQL is ready."41            break42        fi43        sleep 344        elapsed=$((elapsed + 3))45    done46    if [ "${elapsed}" -ge "${DB_WAIT_TIMEOUT}" ]; then47        log "[ERROR] PostgreSQL not reachable after ${DB_WAIT_TIMEOUT}s."48        log "Falling back to SQLite so the container still boots."49        export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN="sqlite:///${AIRFLOW_HOME}/airflow.db"50        export AIRFLOW__CORE__EXECUTOR="SequentialExecutor"51    fi52fi53 54# ---- Build authenticated repo URL (if PAT is set) ----55GIT_CLONE_URL="${DAG_REPO_URL}"56if [ -n "${DAG_REPO_TOKEN}" ]; then57    # Inject token: https://github.com/... → https://<token>@github.com/...58    GIT_CLONE_URL=$(echo "${DAG_REPO_URL}" | sed "s|https://|https://${DAG_REPO_TOKEN}@|")59    log "Using authenticated URL for DAG repo (PAT token set)."60fi61 62# ---- Sync DAG Repository ----63_create_sample_dag() {64    log "Creating sample DAG for testing..."65    cat > "${DAGS_DIR}/sample_dag.py" << 'SAMPLE_EOF'66"""67Sample DAG — created automatically because no DAG repository was configured.68Replace this by setting DAG_REPO_URL or adding your own DAGs to /opt/airflow/dags.69"""70from datetime import datetime, timedelta71from airflow import DAG72from airflow.operators.bash import BashOperator73 74with DAG(75    dag_id="sample_hello_world",76    description="A sample DAG to verify Airflow is working",77    schedule=timedelta(hours=1),78    start_date=datetime(2024, 1, 1),79    catchup=False,80    tags=["sample", "test"],81) as dag:82    hello = BashOperator(83        task_id="say_hello",84        bash_command='echo "Hello from Airflow! $(date)"',85    )86    check_env = BashOperator(87        task_id="check_environment",88        bash_command='echo "Python: $(python3 --version)" && echo "Airflow Home: $AIRFLOW_HOME"',89    )90    hello >> check_env91SAMPLE_EOF92    log "Sample DAG created at ${DAGS_DIR}/sample_dag.py"93}94 95log "Syncing DAG repository..."96if [ ! -d "${DAGS_DIR}/.git" ]; then97    # Fresh clone — clean any stale files (except sample_dag) first98    rm -rf "${DAGS_DIR:?}"/* 2>/dev/null || true99    if git clone --depth 1 --branch "${DAG_REPO_BRANCH}" "${GIT_CLONE_URL}" "${DAGS_DIR}" 2>/dev/null; then100        log "DAG repository cloned."101    else102        log "[WARN] Could not clone DAG repo (URL unreachable or private without token)."103        _create_sample_dag104    fi105else106    cd "${DAGS_DIR}"107    # If PAT changed, update the remote URL108    git remote set-url origin "${GIT_CLONE_URL}" 2>/dev/null || true109    if git fetch origin "${DAG_REPO_BRANCH}" 2>/dev/null && git reset --hard "origin/${DAG_REPO_BRANCH}" 2>/dev/null; then110        log "DAG repository updated."111    else112        log "[WARN] Could not update DAG repo. Using existing DAGs."113    fi114fi115 116# ---- Database Migrations ----117log "Running database migrations..."118airflow db migrate 2>&1 | tail -5119log "Migrations complete."120 121# ---- Create Admin User (idempotent) ----122log "Ensuring admin user exists..."123if ! airflow users list 2>/dev/null | grep -q "${AIRFLOW_ADMIN_USER}"; then124    airflow users create \125        --username "${AIRFLOW_ADMIN_USER}" \126        --firstname Admin \127        --lastname User \128        --role Admin \129        --email "${AIRFLOW_ADMIN_USER}@example.com" \130        --password "${AIRFLOW_ADMIN_PASSWORD}"131    log "Admin user '${AIRFLOW_ADMIN_USER}' created."132else133    log "Admin user '${AIRFLOW_ADMIN_USER}' already exists — skipping."134fi135 136# ---- Start Supervisor ----137log "Starting services via supervisord..."138exec /usr/bin/supervisord -c /etc/supervisord.conf