|
| 1 | +"""DEV-only read-only data checks and a protected, validated database backup.""" |
| 2 | + |
| 3 | +import hashlib |
| 4 | +import json |
| 5 | +import os |
| 6 | +from pathlib import Path |
| 7 | +import re |
| 8 | +import shutil |
| 9 | +import subprocess |
| 10 | +import sys |
| 11 | +from datetime import datetime, timezone |
| 12 | +from uuid import uuid4 |
| 13 | + |
| 14 | +REPO = Path("/root/api") |
| 15 | +BACKUPS = Path("/root/procollab-dependency-backups") |
| 16 | +KEY = re.compile(r"dependency-[0-9]{8}T[0-9]{6}Z-[a-f0-9]{8}") |
| 17 | + |
| 18 | +PROBE = r""" |
| 19 | +import hashlib, json, os |
| 20 | +os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'procollab.settings') |
| 21 | +import django |
| 22 | +django.setup() |
| 23 | +from django.db import connection, transaction |
| 24 | +payload = {'django': django.get_version(), 'tables': {}} |
| 25 | +with transaction.atomic(): |
| 26 | + with connection.cursor() as cursor: |
| 27 | + cursor.execute('SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY') |
| 28 | + cursor.execute("SET LOCAL statement_timeout = '60s'") |
| 29 | + cursor.execute("SET LOCAL lock_timeout = '10s'") |
| 30 | + cursor.execute( |
| 31 | + "SELECT current_database(), current_setting('server_version_num')::int, " |
| 32 | + "pg_database_size(current_database())" |
| 33 | + ) |
| 34 | + name, version, size = cursor.fetchone() |
| 35 | + if name != 'procollab_db_dev' or not 140000 <= version < 190000: |
| 36 | + raise RuntimeError('Unexpected DEV database or unsupported PostgreSQL version') |
| 37 | + payload.update(server_major=version // 10000, database_bytes=size) |
| 38 | + cursor.execute("SELECT tablename FROM pg_tables WHERE schemaname='public' ORDER BY tablename") |
| 39 | + tables = [ |
| 40 | + row[0] for row in cursor.fetchall() |
| 41 | + if row[0].startswith( |
| 42 | + ('users_', 'projects_', 'partner_programs_', 'project_rates_', 'files_') |
| 43 | + ) |
| 44 | + ] |
| 45 | + for table in tables: |
| 46 | + cursor.execute( |
| 47 | + 'SELECT a.attname FROM pg_index i JOIN pg_attribute a ' |
| 48 | + 'ON a.attrelid=i.indrelid AND a.attnum=ANY(i.indkey) ' |
| 49 | + 'WHERE i.indrelid=to_regclass(%s) AND i.indisprimary ORDER BY a.attnum', |
| 50 | + [table], |
| 51 | + ) |
| 52 | + primary = [row[0] for row in cursor.fetchall()] |
| 53 | + if not primary: |
| 54 | + raise RuntimeError('Missing primary key in a protected table') |
| 55 | + identity = 'jsonb_build_array(' + ','.join( |
| 56 | + 't.' + connection.ops.quote_name(key) for key in primary |
| 57 | + ) + ')::text' |
| 58 | + record = "(to_jsonb(t) - 'last_activity' - 'last_login' - 'datetime_updated')::text" |
| 59 | + hash_sql = lambda expression: "encode(sha256(convert_to(" + expression + ", 'UTF8')), 'hex')" |
| 60 | + cursor.execute( |
| 61 | + 'SELECT ' + hash_sql(identity) + ',' + hash_sql(record) |
| 62 | + + ' FROM ' + connection.ops.quote_name(table) + ' t' |
| 63 | + ) |
| 64 | + records = {} |
| 65 | + while True: |
| 66 | + batch = cursor.fetchmany(500) |
| 67 | + if not batch: |
| 68 | + break |
| 69 | + for identity, record_hash in batch: |
| 70 | + records[identity] = record_hash |
| 71 | + payload['tables'][table] = records |
| 72 | +print(json.dumps(payload, ensure_ascii=False)) |
| 73 | +""" |
| 74 | + |
| 75 | + |
| 76 | +def emit(message, *, file=sys.stdout, flush=True): |
| 77 | + file.write(message + "\n") |
| 78 | + if flush: |
| 79 | + file.flush() |
| 80 | + |
| 81 | + |
| 82 | +def command(args, **kwargs): |
| 83 | + try: |
| 84 | + return subprocess.run(args, check=True, stderr=subprocess.PIPE, **kwargs) |
| 85 | + except (subprocess.CalledProcessError, subprocess.TimeoutExpired): |
| 86 | + raise RuntimeError( |
| 87 | + "A DEV backup/check command failed; no credentials are printed" |
| 88 | + ) from None |
| 89 | + |
| 90 | + |
| 91 | +def output(args, **kwargs): |
| 92 | + return command(args, stdout=subprocess.PIPE, **kwargs).stdout |
| 93 | + |
| 94 | + |
| 95 | +def inspect_container(container): |
| 96 | + value = json.loads(output(["docker", "inspect", container]))[0] |
| 97 | + labels = value["Config"].get("Labels", {}) |
| 98 | + paths = labels.get("com.docker.compose.project.config_files", "").split(",") |
| 99 | + if str(REPO / "docker-compose.dev-ci.yml") not in paths: |
| 100 | + raise RuntimeError("Container does not belong to the DEV compose configuration") |
| 101 | + return value |
| 102 | + |
| 103 | + |
| 104 | +def snapshot(web): |
| 105 | + return json.loads( |
| 106 | + output( |
| 107 | + ["docker", "exec", "-i", web, "python", "-"], |
| 108 | + input=PROBE.encode(), |
| 109 | + timeout=180, |
| 110 | + ) |
| 111 | + ) |
| 112 | + |
| 113 | + |
| 114 | +def digest(path): |
| 115 | + value = hashlib.sha256() |
| 116 | + with path.open("rb") as handle: |
| 117 | + for chunk in iter(lambda: handle.read(1024 * 1024), b""): |
| 118 | + value.update(chunk) |
| 119 | + return value.hexdigest() |
| 120 | + |
| 121 | + |
| 122 | +def main(): |
| 123 | + phase = sys.argv[1] |
| 124 | + key = sys.argv[2] if len(sys.argv) > 2 else "" |
| 125 | + if phase not in ("backup", "verify") or REPO.resolve() != REPO: |
| 126 | + raise RuntimeError("Invalid operation or repository location") |
| 127 | + os.umask(0o077) |
| 128 | + os.chdir(REPO) |
| 129 | + web = output( |
| 130 | + ["docker", "compose", "-f", "docker-compose.dev-ci.yml", "ps", "-q", "web"], |
| 131 | + text=True, |
| 132 | + ).strip() |
| 133 | + worker = output( |
| 134 | + ["docker", "compose", "-f", "docker-compose.dev-ci.yml", "ps", "-q", "celerys"], |
| 135 | + text=True, |
| 136 | + ).strip() |
| 137 | + if not re.fullmatch(r"[a-f0-9]{12,64}", web) or not re.fullmatch( |
| 138 | + r"[a-f0-9]{12,64}", worker |
| 139 | + ): |
| 140 | + raise RuntimeError("Expected exactly one running DEV web and worker") |
| 141 | + runtime = inspect_container(web) |
| 142 | + worker_runtime = inspect_container(worker) |
| 143 | + if runtime["Image"] != worker_runtime["Image"]: |
| 144 | + raise RuntimeError("DEV web and worker use different images") |
| 145 | + current = snapshot(web) |
| 146 | + emit("POSTGRESQL_MAJOR=" + str(current["server_major"]), flush=True) |
| 147 | + emit("DJANGO_RUNTIME=" + current["django"], flush=True) |
| 148 | + |
| 149 | + if phase == "backup": |
| 150 | + required = max(2 * 1024**3, current["database_bytes"] * 2 + 512 * 1024**2) |
| 151 | + if shutil.disk_usage("/root").free < required: |
| 152 | + raise RuntimeError( |
| 153 | + "Insufficient free space for a protected DEV backup and rollout" |
| 154 | + ) |
| 155 | + key = ( |
| 156 | + "dependency-" |
| 157 | + + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") |
| 158 | + + "-" |
| 159 | + + uuid4().hex[:8] |
| 160 | + ) |
| 161 | + BACKUPS.mkdir(mode=0o700, exist_ok=True) |
| 162 | + folder = BACKUPS / key |
| 163 | + folder.mkdir(mode=0o700) |
| 164 | + dump = folder / "database.dump" |
| 165 | + environment_script = ( |
| 166 | + "import json,os; print(json.dumps({k:os.environ[k] for k in " |
| 167 | + '("DATABASE_HOST","DATABASE_PORT","DATABASE_NAME",' |
| 168 | + '"DATABASE_USER","DATABASE_PASSWORD")}))' |
| 169 | + ) |
| 170 | + config = json.loads( |
| 171 | + output(["docker", "exec", web, "python", "-c", environment_script]) |
| 172 | + ) |
| 173 | + if config["DATABASE_NAME"] != "procollab_db_dev": |
| 174 | + raise RuntimeError("Unexpected DEV database configuration") |
| 175 | + environment = dict(os.environ) |
| 176 | + for destination, source in ( |
| 177 | + ("PGHOST", "DATABASE_HOST"), |
| 178 | + ("PGPORT", "DATABASE_PORT"), |
| 179 | + ("PGDATABASE", "DATABASE_NAME"), |
| 180 | + ("PGUSER", "DATABASE_USER"), |
| 181 | + ("PGPASSWORD", "DATABASE_PASSWORD"), |
| 182 | + ): |
| 183 | + environment[destination] = config[source] |
| 184 | + client = "postgres:" + str(current["server_major"]) + "-alpine" |
| 185 | + container_name = "dependency-backup-" + uuid4().hex[:12] |
| 186 | + try: |
| 187 | + with dump.open("xb") as handle: |
| 188 | + command( |
| 189 | + [ |
| 190 | + "docker", |
| 191 | + "run", |
| 192 | + "--rm", |
| 193 | + "--name", |
| 194 | + container_name, |
| 195 | + "--network", |
| 196 | + "container:" + web, |
| 197 | + "--env", |
| 198 | + "PGHOST", |
| 199 | + "--env", |
| 200 | + "PGPORT", |
| 201 | + "--env", |
| 202 | + "PGDATABASE", |
| 203 | + "--env", |
| 204 | + "PGUSER", |
| 205 | + "--env", |
| 206 | + "PGPASSWORD", |
| 207 | + client, |
| 208 | + "pg_dump", |
| 209 | + "--format=custom", |
| 210 | + "--no-owner", |
| 211 | + "--no-acl", |
| 212 | + "--lock-wait-timeout=30s", |
| 213 | + ], |
| 214 | + env=environment, |
| 215 | + stdout=handle, |
| 216 | + timeout=600, |
| 217 | + ) |
| 218 | + finally: |
| 219 | + subprocess.run( |
| 220 | + ["docker", "rm", "-f", container_name], |
| 221 | + stdout=subprocess.DEVNULL, |
| 222 | + stderr=subprocess.DEVNULL, |
| 223 | + ) |
| 224 | + with dump.open("rb") as handle: |
| 225 | + if handle.read(5) != b"PGDMP": |
| 226 | + raise RuntimeError("Invalid backup archive") |
| 227 | + handle.seek(0) |
| 228 | + command( |
| 229 | + ["docker", "run", "--rm", "-i", client, "pg_restore", "--file=/dev/null"], |
| 230 | + stdin=handle, |
| 231 | + stdout=subprocess.DEVNULL, |
| 232 | + timeout=600, |
| 233 | + ) |
| 234 | + rollback = "procollab-dependency-rollback:" + key |
| 235 | + command(["docker", "tag", runtime["Image"], rollback], stdout=subprocess.DEVNULL) |
| 236 | + shutil.copyfile(REPO / ".env", folder / "runtime.env") |
| 237 | + os.chmod(folder / "runtime.env", 0o600) |
| 238 | + metadata = { |
| 239 | + "snapshot": current, |
| 240 | + "image_id": runtime["Image"], |
| 241 | + "image": runtime["Config"]["Image"], |
| 242 | + "rollback_tag": rollback, |
| 243 | + "source_sha": output(["git", "rev-parse", "HEAD"], text=True).strip(), |
| 244 | + "dump_sha256": digest(dump), |
| 245 | + "dump_bytes": dump.stat().st_size, |
| 246 | + } |
| 247 | + (folder / "manifest.json").write_text(json.dumps(metadata), encoding="utf8") |
| 248 | + emit("BACKUP_KEY=" + key) |
| 249 | + emit("BACKUP_BYTES=" + str(metadata["dump_bytes"])) |
| 250 | + emit("BACKUP_SHA256=" + metadata["dump_sha256"]) |
| 251 | + emit("ROLLBACK_IMAGE=" + rollback) |
| 252 | + else: |
| 253 | + if KEY.fullmatch(key) is None: |
| 254 | + raise RuntimeError("Invalid backup key") |
| 255 | + folder = BACKUPS / key |
| 256 | + if folder.resolve().parent != BACKUPS.resolve(): |
| 257 | + raise RuntimeError("Invalid backup location") |
| 258 | + before = json.loads((folder / "manifest.json").read_text(encoding="utf8")) |
| 259 | + if digest(folder / "database.dump") != before["dump_sha256"]: |
| 260 | + raise RuntimeError("Backup digest mismatch") |
| 261 | + removed, changed = {}, {} |
| 262 | + for table, records in before["snapshot"]["tables"].items(): |
| 263 | + after = current["tables"].get(table, {}) |
| 264 | + missing = set(records) - set(after) |
| 265 | + different = [ |
| 266 | + identity |
| 267 | + for identity in set(records) & set(after) |
| 268 | + if records[identity] != after[identity] |
| 269 | + ] |
| 270 | + if missing: |
| 271 | + removed[table] = len(missing) |
| 272 | + if different: |
| 273 | + changed[table] = len(different) |
| 274 | + emit("REMOVED_EXISTING_RECORDS=" + json.dumps(removed)) |
| 275 | + emit("CHANGED_EXISTING_RECORDS=" + json.dumps(changed)) |
| 276 | + if removed or changed: |
| 277 | + raise RuntimeError("Existing DEV records changed; inspect before continuing") |
| 278 | + emit("EXISTING_DEV_DATA_PRESERVED=true") |
| 279 | + emit("CURRENT_IMAGE=" + runtime["Config"]["Image"]) |
| 280 | + emit( |
| 281 | + "PROTECTED_TABLE_COUNTS=" |
| 282 | + + json.dumps({name: len(rows) for name, rows in current["tables"].items()}) |
| 283 | + ) |
| 284 | + |
| 285 | + |
| 286 | +if __name__ == "__main__": |
| 287 | + try: |
| 288 | + main() |
| 289 | + except Exception as error: |
| 290 | + emit("DEV_DEPENDENCY_PREFLIGHT_FAILED=" + type(error).__name__, file=sys.stderr) |
| 291 | + if isinstance(error, RuntimeError): |
| 292 | + emit(str(error), file=sys.stderr) |
| 293 | + raise SystemExit(1) |
0 commit comments