1#!/usr/bin/env python3
2import argparse
3from contextlib import contextmanager
4from datetime import datetime, timezone
5import hashlib
6import json
7import os
8from pathlib import Path
9import re
10import secrets
11import shutil
12import subprocess
13import time
14
15
16ROOT = Path("/opt/studio")
17STATE = Path("/var/lib/studio")
18BACKUPS = STATE / "backups"
19BACKUP_ID = re.compile(r"\d{8}T\d{6}Z-[0-9a-f]{6}\Z")
20SERVICE_ID = re.compile(r"[a-z][a-z0-9-]*\Z")
21RELEASE_ID = re.compile(r"[0-9a-f]{16}\Z")
22PODMAN = ["podman", "--url", "unix:///run/podman/podman.sock"]
23
24
25def run(*args, capture=False, **kwargs):
26 return subprocess.run(args, check=True, capture_output=capture, **kwargs)
27
28
29def current_release():
30 current = ROOT / "current"
31 if not current.is_symlink():
32 raise ValueError("No active home server release")
33 release = current.resolve()
34 if release.parent != ROOT / "releases" or not RELEASE_ID.fullmatch(release.name):
35 raise ValueError("Active home server release is invalid")
36 return release
37
38
39def postgres_container():
40 allocations = json.loads(run("nomad", "job", "allocs", "-json", "postgres", capture=True, text=True).stdout)
41 running = [a["ID"] for a in allocations if a["ClientStatus"] == "running" and a["DesiredStatus"] == "run"]
42 if len(running) != 1:
43 raise ValueError("Postgres needs one running allocation")
44 containers = run(*PODMAN, "ps", "--format", "{{.ID}} {{.Names}}", capture=True, text=True).stdout.splitlines()
45 matches = [parts[0] for line in containers if len(parts := line.split()) == 2 and parts[1].endswith(running[0])]
46 if len(matches) != 1:
47 raise ValueError("Postgres container is unavailable")
48 return matches[0]
49
50
51def database_inputs(service):
52 prefix = f"nomad/jobs/{service}/inputs/"
53 listed = json.loads(run("nomad", "var", "list", "-out=json", prefix, capture=True, text=True).stdout)
54 databases = {}
55 for item in listed:
56 path = item["Path"]
57 alias = path.removeprefix(prefix)
58 if not path.startswith(prefix) or not SERVICE_ID.fullmatch(alias):
59 raise ValueError(f"Invalid input path for {service}: {path}")
60 values = json.loads(run("nomad", "var", "get", "-out=json", path, capture=True, text=True).stdout)["Items"]
61 if "name" not in values or "username" not in values:
62 continue
63 name, owner = values["name"], values["username"]
64 if not owner.startswith("svc_"):
65 continue
66 if not re.fullmatch(r"[a-z][a-z0-9_]*", name) or owner != "svc_" + name:
67 raise ValueError(f"Invalid database allocation for {service}.{alias}")
68 if any(database == name for database, _ in databases.values()):
69 raise ValueError(f"Database allocated twice for {service}: {name}")
70 databases[alias] = (name, owner)
71 return databases
72
73
74def database_dump(container, database, destination, extensions):
75 with destination.open("wb") as output:
76 run(*PODMAN, "exec", container, "pg_dump", "-U", "postgres", "-Fc", "--no-owner", "--no-acl",
77 *(f"--exclude-extension={extension}" for extension in extensions), database, stdout=output)
78 if not destination.stat().st_size:
79 raise ValueError(f"Empty database backup: {database}")
80
81
82def checksum(path):
83 digest = hashlib.sha256()
84 with path.open("rb") as file:
85 for chunk in iter(lambda: file.read(1024 * 1024), b""):
86 digest.update(chunk)
87 return digest.hexdigest()
88
89
90def dataset_for(service):
91 root = Path("/srv/prod") / service
92 result = subprocess.run(["findmnt", "-n", "-o", "SOURCE,FSTYPE", "--mountpoint", str(root)], capture_output=True, text=True)
93 if result.returncode:
94 return None
95 source, fstype = result.stdout.split()
96 if fstype != "zfs" or not source.endswith("/prod/" + service):
97 raise ValueError(f"Unexpected dataset at {root}: {source}")
98 return source
99
100
101@contextmanager
102def cloned_postgres(snapshot, clone, mountpoint):
103 live_container = postgres_container()
104 inspect = json.loads(run(*PODMAN, "inspect", live_container, capture=True, text=True).stdout)[0]
105 mounts = [mount for mount in inspect["Mounts"] if mount["Destination"] == "/var/lib/postgresql"]
106 if len(mounts) != 1 or mounts[0]["Type"] != "bind":
107 raise ValueError("Postgres data mount is unavailable")
108 data_path = Path(mounts[0]["Source"])
109 if not data_path.is_relative_to("/srv/prod/postgres"):
110 raise ValueError("Postgres data is outside its dataset")
111 user = inspect["Config"]["User"]
112 if not re.fullmatch(r"[0-9]+:[0-9]+", user):
113 raise ValueError("Postgres container user is invalid")
114 pgdata = [env for env in inspect["Config"]["Env"] if env.startswith("PGDATA=")]
115 if len(pgdata) != 1 or not Path(pgdata[0][7:]).is_relative_to("/var/lib/postgresql"):
116 raise ValueError("Postgres data directory is invalid")
117 image = inspect["Image"]
118 if not re.fullmatch(r"[0-9a-f]{64}", image):
119 raise ValueError("Postgres image ID is invalid")
120 probe = "studio-backup-" + clone.rsplit("/", 1)[1]
121 created = False
122 try:
123 run("zfs", "clone", "-o", f"mountpoint={mountpoint}", snapshot, clone)
124 created = True
125 run(*PODMAN, "run", "-d", "--name", probe, "--network", "none", "--user", user,
126 "-e", pgdata[0], "-e", "POSTGRES_PASSWORD=backup-probe",
127 "-v", f"{Path(mountpoint) / data_path.relative_to('/srv/prod/postgres')}:/var/lib/postgresql", image,
128 stdout=subprocess.DEVNULL)
129 for _ in range(60):
130 ready = subprocess.run([*PODMAN, "exec", probe, "pg_isready", "-U", "postgres"], capture_output=True)
131 if ready.returncode == 0:
132 break
133 time.sleep(1)
134 else:
135 raise ValueError("Snapshot PostgreSQL did not become ready")
136 yield probe
137 finally:
138 subprocess.run([*PODMAN, "rm", "-f", probe], check=False, stdout=subprocess.DEVNULL)
139 if created:
140 run("zfs", "destroy", clone)
141
142
143def backup(from_release, to_release):
144 if not RELEASE_ID.fullmatch(from_release) or not RELEASE_ID.fullmatch(to_release):
145 raise ValueError("Invalid release ID")
146 if current_release().name != from_release:
147 raise ValueError("Active release changed before backup")
148 names = json.loads((STATE / "managed-jobs.json").read_text())
149 if not names or any(not SERVICE_ID.fullmatch(name) for name in names):
150 raise ValueError("Managed jobs list is invalid")
151 backup_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-" + secrets.token_hex(3)
152 directory = BACKUPS / backup_id
153 manifest = {"id": backup_id, "fromRelease": from_release, "toRelease": to_release,
154 "time": int(time.time()), "services": {}}
155 snapshot = "studio-" + backup_id
156 datasets = set()
157 allocations = {}
158 for service in sorted(names):
159 entry = {}
160 dataset = dataset_for(service)
161 if dataset:
162 datasets.add(dataset)
163 entry["dataset"] = dataset
164 entry["snapshot"] = snapshot
165 allocated = database_inputs(service)
166 if allocated:
167 entry["databases"] = {}
168 allocations[service] = allocated
169 if entry:
170 manifest["services"][service] = entry
171 if allocations:
172 postgres_dataset = dataset_for("postgres")
173 if not postgres_dataset:
174 raise ValueError("Postgres needs a ZFS dataset for consistent backups")
175 datasets.add(postgres_dataset)
176 directory.mkdir(parents=True, mode=0o700)
177 snapshots = [f"{dataset}@{snapshot}" for dataset in sorted(datasets)]
178 created = False
179 try:
180 if snapshots:
181 run("zfs", "snapshot", *snapshots)
182 created = True
183 nomad_snapshot = directory / "nomad.snap"
184 run("nomad", "operator", "snapshot", "save", str(nomad_snapshot))
185 manifest["nomadSnapshot"] = {"file": nomad_snapshot.name, "sha256": checksum(nomad_snapshot)}
186 if allocations:
187 clone = postgres_dataset.rsplit("/prod/", 1)[0] + "/staging/" + backup_id + "-postgres"
188 mountpoint = Path("/srv/staging") / (backup_id + "-postgres")
189 with cloned_postgres(f"{postgres_dataset}@{snapshot}", clone, mountpoint) as probe:
190 for service, allocated in allocations.items():
191 entry = manifest["services"][service]
192 for alias, (database, owner) in sorted(allocated.items()):
193 result = run(*PODMAN, "exec", probe, "psql", "-U", "postgres", "-d", database,
194 "-At", "-c", "SELECT extname FROM pg_extension WHERE extname <> 'plpgsql' ORDER BY extname",
195 capture=True, text=True)
196 extensions = result.stdout.splitlines()
197 if any(not re.fullmatch(r"[a-z][a-z0-9_]*", extension) for extension in extensions):
198 raise ValueError(f"Invalid extension in {database}")
199 dump = directory / f"{service}-{alias}.dump"
200 database_dump(probe, database, dump, extensions)
201 entry["databases"][alias] = {"name": database, "owner": owner, "extensions": extensions,
202 "dump": dump.name, "sha256": checksum(dump)}
203 pending = directory / "manifest.pending"
204 pending.write_text(json.dumps(manifest, indent=2) + "\n")
205 pending.replace(directory / "manifest.json")
206 except Exception:
207 if created:
208 for target in snapshots:
209 subprocess.run(["zfs", "destroy", target], check=False)
210 shutil.rmtree(directory)
211 raise
212 print(f"backup={backup_id} services={len(manifest['services'])}")
213 return backup_id
214
215
216def list_backups():
217 for path in sorted(BACKUPS.glob("*/manifest.json")):
218 manifest = json.loads(path.read_text())
219 print(f"{manifest['id']} {manifest['fromRelease']} -> {manifest['toRelease']} {', '.join(sorted(manifest['services']))}")
220
221
222def restore(backup_id, service):
223 if not BACKUP_ID.fullmatch(backup_id) or not SERVICE_ID.fullmatch(service):
224 raise ValueError("Invalid backup or service ID")
225 directory = BACKUPS / backup_id
226 manifest = json.loads((directory / "manifest.json").read_text())
227 if manifest["id"] != backup_id or service not in manifest["services"]:
228 raise ValueError("Backup does not contain that service")
229 if current_release().name != manifest["fromRelease"]:
230 raise ValueError(f"Roll back code to {manifest['fromRelease']} before restoring data")
231 if service == "postgres":
232 raise ValueError("Postgres serves multiple services; restore a specific database owner")
233 entry = manifest["services"][service]
234 dataset = entry.get("dataset")
235 if dataset and dataset_for(service) != dataset:
236 raise ValueError("Service dataset changed since backup")
237 root = Path("/srv/prod") / service
238 source = root / ".zfs/snapshot" / entry["snapshot"] if dataset else None
239 if source and not source.is_dir():
240 raise ValueError("ZFS backup snapshot is unavailable")
241 databases = entry.get("databases", {"database": entry["database"]} if "database" in entry else {})
242 allocated = database_inputs(service)
243 for alias, database in databases.items():
244 if allocated.get(alias) != (database["name"], database["owner"]):
245 raise ValueError(f"Service database changed since backup: {alias}")
246 dump = directory / database["dump"]
247 if not dump.is_file() or checksum(dump) != database["sha256"]:
248 raise ValueError("Database dump is missing or changed")
249 run("nomad", "job", "stop", "-yes", service)
250 for _ in range(30):
251 allocations = json.loads(run("nomad", "job", "allocs", "-json", service, capture=True, text=True).stdout)
252 if not any(a["ClientStatus"] == "running" for a in allocations):
253 break
254 time.sleep(2)
255 else:
256 raise ValueError(f"{service} did not stop; data remains unchanged")
257 safety = "before-restore-" + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-" + secrets.token_hex(3)
258 if dataset:
259 run("zfs", "snapshot", f"{dataset}@{safety}")
260 container = postgres_container() if databases else None
261 for alias, database in databases.items():
262 database_dump(container, database["name"], directory / f"{service}-{alias}-{safety}.dump", database["extensions"])
263 if dataset:
264 run("rsync", "-aHAX", "--numeric-ids", "--one-file-system", "--delete", "--exclude=/.zfs/",
265 str(source) + "/", str(root) + "/")
266 for database in databases.values():
267 name, owner = database["name"], database["owner"]
268 run(*PODMAN, "exec", "-i", container, "psql", "-U", "postgres", "-d", "postgres", "-v", "ON_ERROR_STOP=1",
269 input=f"DROP DATABASE {name} WITH (FORCE);\nCREATE DATABASE {name} OWNER {owner};\n", text=True)
270 for extension in database["extensions"]:
271 run(*PODMAN, "exec", "-i", container, "psql", "-U", "postgres", "-d", name, "-v", "ON_ERROR_STOP=1",
272 input=f"CREATE EXTENSION {extension};\n", text=True)
273 with (directory / database["dump"]).open("rb") as input_file:
274 run(*PODMAN, "exec", "-i", container, "pg_restore", "-U", "postgres", "-d", name,
275 "--no-owner", "--no-acl", "--role=" + owner, stdin=input_file)
276 script = current_release() / "tools/studio.py"
277 run("python3", str(script), "deploy", service)
278 run("python3", str(script), "check", service)
279 print(f"restored={service} backup={backup_id} safety={safety}")
280
281
282def main():
283 parser = argparse.ArgumentParser(description="Back up or restore service data on the home server")
284 parser.add_argument("mode", choices=["backup", "list", "restore"])
285 parser.add_argument("first", nargs="?")
286 parser.add_argument("second", nargs="?")
287 parser.add_argument("--discard-writes", action="store_true")
288 args = parser.parse_args()
289 os.umask(0o077)
290 os.environ["NOMAD_TOKEN"] = (STATE / "nomad.token").read_text().strip()
291 if args.mode == "backup" and args.first and args.second and not args.discard_writes:
292 backup(args.first, args.second)
293 elif args.mode == "list" and not args.first and not args.second and not args.discard_writes:
294 list_backups()
295 elif args.mode == "restore" and args.first and args.second and args.discard_writes:
296 restore(args.first, args.second)
297 else:
298 parser.error("Expected backup FROM TO, list, or restore BACKUP SERVICE --discard-writes")
299
300
301if __name__ == "__main__":
302 main()