1#!/usr/bin/env python3
2import json
3import os
4import re
5import secrets
6import subprocess
7import sys
8import urllib.request
9
10
11data = json.load(sys.stdin)
12request = data["request"]
13if request["kind"] == "admin":
14 if data.get("operation") != "delete":
15 print(json.dumps({"username": "postgres", "password": data["providerSecrets"]["password"]}))
16 sys.exit(0)
17if request["kind"] != "database":
18 raise ValueError("unsupported Postgres input")
19source_database = request["name"]
20if not re.fullmatch(r"[a-z][a-z0-9_]{0,55}", source_database):
21 raise ValueError("invalid database name")
22extensions = request.get("extensions", [])
23if any(not re.fullmatch(r"[a-z][a-z0-9_]*", extension) for extension in extensions):
24 raise ValueError("invalid extension name")
25stage_id = data.get("stageId")
26database = source_database[:42] + "_s_" + stage_id[-8:] if stage_id else source_database
27username = "svc_" + database
28existing = data.get("existing") or {}
29password = existing.get("password") or secrets.token_urlsafe(32)
30if data.get("operation") == "delete" and not stage_id:
31 raise ValueError("refusing to delete a production database")
32if existing.get("name") and existing["name"] != database:
33 raise ValueError("database name changed; migrate the existing database before deployment")
34
35nomad_request = urllib.request.Request(
36 "http://127.0.0.1:4646/v1/job/postgres/allocations",
37 headers={"X-Nomad-Token": os.environ["NOMAD_TOKEN"]},
38)
39with urllib.request.urlopen(nomad_request, timeout=5) as response:
40 allocations = json.load(response)
41running = [item["ID"] for item in allocations if item["ClientStatus"] == "running" and item["DesiredStatus"] == "run"]
42if len(running) != 1:
43 raise ValueError("Postgres needs exactly one running allocation")
44result = subprocess.run(
45 ["podman", "--url", "unix:///run/podman/podman.sock", "ps", "--format", "{{.ID}} {{.Names}}"],
46 text=True, capture_output=True, check=True,
47)
48containers = [parts[0] for line in result.stdout.splitlines() if len(parts := line.split()) == 2 and parts[1].endswith(running[0])]
49if len(containers) != 1:
50 raise ValueError("Postgres allocation container is unavailable")
51container = containers[0]
52podman = ["podman", "--url", "unix:///run/podman/podman.sock", "exec"]
53
54
55def sql(statement, db="postgres"):
56 result = subprocess.run(
57 [*podman, "-i", container,
58 "psql", "-U", "postgres", "-d", db, "-tA", "-v", "ON_ERROR_STOP=1"],
59 input=statement + "\n", text=True, capture_output=True, check=True,
60 )
61 return result.stdout.strip()
62
63
64if data.get("operation") == "delete":
65 sql(f"DROP DATABASE IF EXISTS {database} WITH (FORCE);")
66 sql(f"DROP ROLE IF EXISTS {username};")
67 sys.exit(0)
68
69quoted_password = password.replace("'", "''")
70role = sql(f"SELECT rolname FROM pg_roles WHERE rolname = '{username}';")
71if not role:
72 sql(f"CREATE ROLE {username} LOGIN PASSWORD '{quoted_password}';")
73elif not existing.get("password"):
74 sql(f"ALTER ROLE {username} PASSWORD '{quoted_password}';")
75owner = sql(f"SELECT pg_get_userbyid(datdba) FROM pg_database WHERE datname = '{database}';")
76if not owner:
77 sql(f"CREATE DATABASE {database} OWNER {username};")
78 for extension in extensions:
79 sql(f"CREATE EXTENSION IF NOT EXISTS {extension};", database)
80 source_container = data.get("sourceContainer") or container
81 source_exists = stage_id and subprocess.run(
82 [*podman, "-i", source_container, "psql", "-U", "postgres", "-d", "postgres", "-tA", "-v", "ON_ERROR_STOP=1"],
83 input=f"SELECT 1 FROM pg_database WHERE datname = '{source_database}';\n",
84 text=True, capture_output=True, check=True,
85 ).stdout.strip()
86 if source_exists:
87 dump = subprocess.Popen(
88 [*podman, source_container, "pg_dump", "-U", "postgres", "-Fc", "--no-owner", "--no-acl",
89 *(f"--exclude-extension={extension}" for extension in extensions), source_database],
90 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
91 )
92 restore = subprocess.Popen(
93 [*podman, "-i", container, "pg_restore", "-U", "postgres", "--no-owner", "--no-acl",
94 "--role=" + username, "-d", database],
95 stdin=dump.stdout, stderr=subprocess.PIPE,
96 )
97 dump.stdout.close()
98 restore_error = restore.communicate()[1]
99 dump_error = dump.stderr.read()
100 if dump.wait() or restore.returncode:
101 raise RuntimeError((dump_error + restore_error).decode(errors="replace"))
102elif owner != username:
103 raise ValueError(f"database {database} is owned by {owner}")
104
105for extension in extensions:
106 sql(f"CREATE EXTENSION IF NOT EXISTS {extension};", database)
107
108print(json.dumps({"name": database, "username": username, "password": password}))