Spaces:
Running
Running
| """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 | |