Spaces:
Running
Running
File size: 2,112 Bytes
2f1e7ff | 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 | """Peer-authenticated startup provisioning, without credential-bearing shell SQL."""
import os
import pwd
import sys
import psycopg2
from psycopg2 import sql
def provision():
postgres = pwd.getpwnam("postgres")
if os.geteuid() == 0:
os.initgroups(postgres.pw_name, postgres.pw_gid)
os.setgid(postgres.pw_gid)
os.setuid(postgres.pw_uid)
role = os.environ["POSTGRES_USERNAME"]
database = os.environ["POSTGRES_SYSTEM_DATABASE"]
conn = psycopg2.connect(host="/tmp", dbname="postgres", user="postgres")
try:
conn.autocommit = True
with conn.cursor() as cur:
# Suppress SQL text even on an unexpected error: ALTER ROLE contains
# a credential. This affects only this administrative connection.
cur.execute("SET log_statement='none'")
cur.execute("SET log_min_error_statement='panic'")
cur.execute("SELECT 1 FROM pg_roles WHERE rolname=%s", (role,))
if cur.fetchone() is None:
cur.execute(sql.SQL("CREATE ROLE {} LOGIN").format(sql.Identifier(role)))
cur.execute(sql.SQL("ALTER ROLE {} WITH LOGIN CREATEDB PASSWORD %s").format(sql.Identifier(role)),
(os.environ["POSTGRES_PASSWORD"],))
cur.execute("SELECT 1 FROM pg_database WHERE datname=%s", (database,))
if cur.fetchone() is None:
cur.execute(sql.SQL("CREATE DATABASE {} OWNER {}").format(sql.Identifier(database), sql.Identifier(role)))
cur.execute(sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {}").format(sql.Identifier(database), sql.Identifier(role)))
finally:
conn.close()
conn = psycopg2.connect(host="/tmp", dbname=database, user="postgres")
try:
conn.autocommit = True
with conn.cursor() as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
finally:
conn.close()
if __name__ == "__main__":
try:
provision()
except Exception:
print("POSTGRES_PROVISIONING_FAILED", file=sys.stderr)
raise SystemExit(1) from None
|