Spaces:
Running
Running
File size: 2,300 Bytes
e6ef7ae 3b9f745 e6ef7ae 684d2f7 e6ef7ae 684d2f7 4eab157 684d2f7 4eab157 684d2f7 6ef084a e6ef7ae 684d2f7 e6ef7ae 684d2f7 e6ef7ae 684d2f7 e6ef7ae |
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 |
#!/bin/bash
set -e
# Load environment variables from .env file if it exists
# This is useful for local development
if [ -f .env ]; then
export $(grep -v '^#' .env | xargs)
fi
# Dynamically extract database connection details from environment variables
# airflow_db
AF_USER=$(echo "$AIRFLOW__DATABASE__SQL_ALCHEMY_CONN" | sed -E 's|postgresql\+psycopg2://([^:]+):.*|\1|')
AF_PASS=$(echo "$AIRFLOW__DATABASE__SQL_ALCHEMY_CONN" | sed -E 's|postgresql\+psycopg2://[^:]+:([^@]+)@.*|\1|')
AF_HOST=$(echo "$AIRFLOW__DATABASE__SQL_ALCHEMY_CONN" | sed -E 's|.*@([^:/]+):([0-9]+)/.*|\1|')
AF_PORT=$(echo "$AIRFLOW__DATABASE__SQL_ALCHEMY_CONN" | sed -E 's|.*@[^:/]+:([0-9]+)/.*|\1|')
AF_DB=$(echo "$AIRFLOW__DATABASE__SQL_ALCHEMY_CONN" | sed -E 's|.*/([^?]+).*|\1|')
# transaction_db
TX_USER=$(echo "$DATABASE_URL" | sed -E 's|postgresql\+psycopg2://([^:]+):.*|\1|')
TX_PASS=$(echo "$DATABASE_URL" | sed -E 's|postgresql\+psycopg2://[^:]+:([^@]+)@.*|\1|')
TX_DB=$(echo "$DATABASE_URL" | sed -E 's|.*/([^?]+).*|\1|')
# Generate pgbouncer.ini
cat <<EOF > /etc/pgbouncer/pgbouncer.ini
[databases]
airflow_db = host=${AF_HOST} port=${AF_PORT} dbname=${AF_DB} user=${AF_USER} password=${AF_PASS}
transaction_db = host=${AF_HOST} port=${AF_PORT} dbname=${TX_DB} user=${TX_USER} password=${TX_PASS}
[pgbouncer]
listen_addr = 127.0.0.1
listen_port = 6432
auth_type = trust
auth_file = /etc/pgbouncer/userlist.txt
pool_mode = transaction
server_tls_sslmode = require
ignore_startup_parameters = extra_float_digits
log_connections = 0
log_disconnections = 0
default_pool_size = 5
EOF
# Generate userlist.txt
cat <<EOF > /etc/pgbouncer/userlist.txt
"${AF_USER}" "${AF_PASS}"
"${TX_USER}" "${TX_PASS}"
EOF
# Launch PgBouncer
pgbouncer /etc/pgbouncer/pgbouncer.ini &
# Dynamically modify the URL to PgBouncer
export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=postgresql+psycopg2://${AF_USER}:${AF_PASS}@127.0.0.1:6432/airflow_db
# Init DB if needed
airflow db migrate
airflow connections create-default-connections
# Create a superuser if it doesn't exist
airflow users create \
--username admin \
--firstname Admin \
--lastname User \
--role Admin \
--email admin@example.com \
--password admin || true
# Launch the scheduler
airflow scheduler &
# Launch the webserver
exec airflow webserver --port 8088
|