airflow / entrypoint.sh
Subham9126's picture
Upload 7 files
494a23b verified
Raw
History Blame Contribute Delete
4.96 kB
#!/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