File size: 4,959 Bytes
0383bfc
494a23b
0383bfc
494a23b
 
 
 
 
 
 
 
0383bfc
494a23b
 
 
 
 
 
 
 
 
 
0383bfc
494a23b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0383bfc
494a23b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0383bfc
494a23b
0383bfc
494a23b
 
 
 
 
 
0383bfc
 
494a23b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0383bfc
 
494a23b
 
0383bfc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
#!/bin/bash
set -euo pipefail

# ---- Helpers ----
log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }

# Graceful shutdown: forward signals to child processes
trap 'log "Caught signal, shutting down..."; kill -TERM $(jobs -p) 2>/dev/null; exit 0' SIGTERM SIGINT SIGQUIT

# ---- Load Configuration ----
log "Loading configuration..."
source /config.sh
validate_config

log "Configuration:"
log "  DB backend : $(echo "${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}" | sed -E 's|://[^@]*@|://***:***@|')"
log "  DAG repo   : ${DAG_REPO_URL} (branch: ${DAG_REPO_BRANCH})"
log "  Executor   : ${AIRFLOW__CORE__EXECUTOR}"
log "  Ports      : nginx=${NGINX_PORT}  webserver=${AIRFLOW_WEB_PORT}  refresh=${REFRESH_SERVICE_PORT}"

# ---- Ensure directories ----
mkdir -p "${DAGS_DIR}" "${LOGS_DIR}"

# ---- Wait for Database (if PostgreSQL) ----
if [[ "${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}" == postgresql* ]]; then
    DB_WAIT_TIMEOUT="${DB_WAIT_TIMEOUT:-60}"
    log "Waiting for PostgreSQL (timeout: ${DB_WAIT_TIMEOUT}s)..."
    elapsed=0
    while [ "${elapsed}" -lt "${DB_WAIT_TIMEOUT}" ]; do
        if python3 -c "
import sqlalchemy, sys
try:
    engine = sqlalchemy.create_engine('${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}', connect_args={'connect_timeout': 5})
    with engine.connect() as conn:
        conn.execute(sqlalchemy.text('SELECT 1'))
    sys.exit(0)
except Exception:
    sys.exit(1)
" 2>/dev/null; then
            log "PostgreSQL is ready."
            break
        fi
        sleep 3
        elapsed=$((elapsed + 3))
    done
    if [ "${elapsed}" -ge "${DB_WAIT_TIMEOUT}" ]; then
        log "[ERROR] PostgreSQL not reachable after ${DB_WAIT_TIMEOUT}s."
        log "Falling back to SQLite so the container still boots."
        export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN="sqlite:///${AIRFLOW_HOME}/airflow.db"
        export AIRFLOW__CORE__EXECUTOR="SequentialExecutor"
    fi
fi

# ---- Build authenticated repo URL (if PAT is set) ----
GIT_CLONE_URL="${DAG_REPO_URL}"
if [ -n "${DAG_REPO_TOKEN}" ]; then
    # Inject token: https://github.com/... → https://<token>@github.com/...
    GIT_CLONE_URL=$(echo "${DAG_REPO_URL}" | sed "s|https://|https://${DAG_REPO_TOKEN}@|")
    log "Using authenticated URL for DAG repo (PAT token set)."
fi

# ---- Sync DAG Repository ----
_create_sample_dag() {
    log "Creating sample DAG for testing..."
    cat > "${DAGS_DIR}/sample_dag.py" << 'SAMPLE_EOF'
"""
Sample DAG — created automatically because no DAG repository was configured.
Replace this by setting DAG_REPO_URL or adding your own DAGs to /opt/airflow/dags.
"""
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.bash import BashOperator

with DAG(
    dag_id="sample_hello_world",
    description="A sample DAG to verify Airflow is working",
    schedule=timedelta(hours=1),
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=["sample", "test"],
) as dag:
    hello = BashOperator(
        task_id="say_hello",
        bash_command='echo "Hello from Airflow! $(date)"',
    )
    check_env = BashOperator(
        task_id="check_environment",
        bash_command='echo "Python: $(python3 --version)" && echo "Airflow Home: $AIRFLOW_HOME"',
    )
    hello >> check_env
SAMPLE_EOF
    log "Sample DAG created at ${DAGS_DIR}/sample_dag.py"
}

log "Syncing DAG repository..."
if [ ! -d "${DAGS_DIR}/.git" ]; then
    # Fresh clone — clean any stale files (except sample_dag) first
    rm -rf "${DAGS_DIR:?}"/* 2>/dev/null || true
    if git clone --depth 1 --branch "${DAG_REPO_BRANCH}" "${GIT_CLONE_URL}" "${DAGS_DIR}" 2>/dev/null; then
        log "DAG repository cloned."
    else
        log "[WARN] Could not clone DAG repo (URL unreachable or private without token)."
        _create_sample_dag
    fi
else
    cd "${DAGS_DIR}"
    # If PAT changed, update the remote URL
    git remote set-url origin "${GIT_CLONE_URL}" 2>/dev/null || true
    if git fetch origin "${DAG_REPO_BRANCH}" 2>/dev/null && git reset --hard "origin/${DAG_REPO_BRANCH}" 2>/dev/null; then
        log "DAG repository updated."
    else
        log "[WARN] Could not update DAG repo. Using existing DAGs."
    fi
fi

# ---- Database Migrations ----
log "Running database migrations..."
airflow db migrate 2>&1 | tail -5
log "Migrations complete."

# ---- Create Admin User (idempotent) ----
log "Ensuring admin user exists..."
if ! airflow users list 2>/dev/null | grep -q "${AIRFLOW_ADMIN_USER}"; then
    airflow users create \
        --username "${AIRFLOW_ADMIN_USER}" \
        --firstname Admin \
        --lastname User \
        --role Admin \
        --email "${AIRFLOW_ADMIN_USER}@example.com" \
        --password "${AIRFLOW_ADMIN_PASSWORD}"
    log "Admin user '${AIRFLOW_ADMIN_USER}' created."
else
    log "Admin user '${AIRFLOW_ADMIN_USER}' already exists — skipping."
fi

# ---- Start Supervisor ----
log "Starting services via supervisord..."
exec /usr/bin/supervisord -c /etc/supervisord.conf