1#!/usr/bin/env python3
2import argparse
3import json
4import os
5from importlib import import_module
6from pathlib import Path
7import re
8import shlex
9import shutil
10import subprocess
11import sys
12import tempfile
13import time
14
15from release import SOURCES, excluded_services, tree_digest
16
17
18REPO = Path(__file__).resolve().parent.parent
19HOST = os.environ.get("STUDIO_DEPLOY_HOST", "root@127.0.0.1")
20PORT = os.environ.get("STUDIO_DEPLOY_PORT", "2222")
21REMOTE = Path("/opt/studio")
22NAME = re.compile(r"[a-z][a-z0-9-]*\Z")
23RELEASE_ID = re.compile(r"[0-9a-f]{16}\Z")
24
25
26def ssh(command, capture=False):
27 return subprocess.run(
28 ["ssh", "-p", PORT, "-o", "BatchMode=yes", HOST, command],
29 check=True, text=True, capture_output=capture,
30 )
31
32
33def run_logged(mode, target, command):
34 run = f"{int(time.time() * 1000)}-{mode}-{target or 'current'}"
35 path = f"/var/lib/studio/runs/{run}.log"
36 script = f"umask 077; mkdir -p /var/lib/studio/runs; set -o pipefail; {command} 2>&1 | tee {shlex.quote(path)}"
37 ssh("bash -c " + shlex.quote(script))
38 print(f"run={run}")
39
40
41def sync_manager(release):
42 source = REMOTE / "releases" / release / "tools"
43 ssh(f"cp {source}/release.py {source}/data.py {REMOTE}/")
44
45
46def upload(main=False):
47 excluded = set() if main else excluded_services(REPO)
48 with tempfile.TemporaryDirectory() as temporary:
49 snapshot = Path(temporary)
50 revision = None
51 if main:
52 output = subprocess.run(
53 ["jj", "--ignore-working-copy", "log", "-r", "main", "--no-graph", "-T",
54 'json(commit_id) ++ "\\n" ++ json(conflict) ++ "\\n" ++ json(description)'],
55 cwd=REPO, check=True, capture_output=True, text=True,
56 ).stdout.splitlines()
57 commit, conflicted, description = map(json.loads, output)
58 if conflicted or not description.strip():
59 raise ValueError("main needs a commit description and no conflicts before deployment")
60 revision = {"commit": commit, "description": description}
61 # Production must use main's exclusions too. A working staging
62 # change must not remove dependencies from the committed release.
63 excluded = set(json.loads(subprocess.run(
64 ["jj", "--ignore-working-copy", "file", "show", "-r", commit,
65 "config/excluded-services.json"],
66 cwd=REPO, check=True, capture_output=True, text=True,
67 ).stdout))
68 entries = subprocess.run(
69 ["jj", "--ignore-working-copy", "file", "list", "-r", commit, "-T",
70 '\"[\" ++ json(path) ++ \",\" ++ json(file_type) ++ \",\" ++ json(executable) ++ \"]\\n\"',
71 *SOURCES], cwd=REPO, check=True, capture_output=True, text=True,
72 ).stdout.splitlines()
73 for entry in entries:
74 name, kind, executable = json.loads(entry)
75 relative = Path(name)
76 if relative.parts[:1] == ("service",) and (
77 set(relative.parts[1:]) & excluded or
78 relative.suffix == ".pkl" and relative.stem in excluded
79 ):
80 continue
81 if kind != "file" or relative.is_absolute() or ".." in relative.parts:
82 raise ValueError(f"unsupported main release file: {name}")
83 destination = snapshot / relative
84 destination.parent.mkdir(parents=True, exist_ok=True)
85 destination.write_bytes(subprocess.run(
86 ["jj", "--ignore-working-copy", "file", "show", "-r", commit, name],
87 cwd=REPO, check=True, capture_output=True,
88 ).stdout)
89 destination.chmod(0o755 if executable else 0o644)
90 else:
91 subprocess.run(
92 ["rsync", "-a", "--exclude=.DS_Store", "--exclude=__pycache__", "--exclude=*.pyc",
93 "--exclude=.identities.lock", "--exclude=identities.pending", "--exclude=._*",
94 "--exclude=dashboard/node_modules", "--exclude=dashboard/dist",
95 "--exclude=dashboard/.cache", "--exclude=dashboard/data", "--exclude=dashboard/target",
96 *(f"--exclude=/service/{name}/" for name in sorted(excluded)),
97 *(f"--exclude=/service/*/{name}.pkl" for name in sorted(excluded)),
98 *(f"--exclude=/service/*/{name}/" for name in sorted(excluded)),
99 *(str(REPO / path) for path in SOURCES),
100 str(snapshot) + "/"], check=True,
101 )
102 for manifest in snapshot.glob("service/*/build-source.json"):
103 spec = json.loads(manifest.read_text())
104 source = (REPO / spec["source"]).resolve(strict=True)
105 if not source.is_dir() or not source.is_relative_to(REPO.parent):
106 raise ValueError(f"invalid build source: {manifest.parent.name}")
107 context = manifest.parent / "build"
108 if context.exists():
109 raise ValueError(f"build context already exists: {manifest.parent.name}")
110 context.mkdir()
111 for name in [*spec["include"], spec["dockerfile"]]:
112 relative = Path(name)
113 if relative.is_absolute() or ".." in relative.parts or not relative.parts:
114 raise ValueError(f"invalid build source path: {name}")
115 origin = source / relative
116 if not origin.resolve(strict=True).is_relative_to(source):
117 raise ValueError(f"build source escapes context: {name}")
118 destination = context / ("Dockerfile" if name == spec["dockerfile"] else name)
119 destination.parent.mkdir(parents=True, exist_ok=True)
120 if origin.is_dir():
121 shutil.copytree(origin, destination, symlinks=True)
122 else:
123 shutil.copy2(origin, destination)
124 if (snapshot / "dashboard/agent/source.json").exists():
125 import_module("build-agent").build(REPO, snapshot / "dashboard/agent/relay.mjs", snapshot / "dashboard/agent/source.json")
126 digest = tree_digest(snapshot)
127 release = digest[:16]
128 remote_release = REMOTE / "releases" / release
129 details = {"id": release, "digest": digest, "version": 2}
130 if revision:
131 details["main"] = revision
132 (snapshot / ".studio-release.json").write_text(json.dumps(details) + "\n")
133 if subprocess.run(
134 ["ssh", "-p", PORT, "-o", "BatchMode=yes", HOST, f"test -d {remote_release}"],
135 check=False,
136 ).returncode == 0:
137 ssh(f"python3 {remote_release}/tools/release.py verify {release}")
138 if revision:
139 script = "import pathlib,sys; p=pathlib.Path(" + repr(str(remote_release / ".studio-release.json")) + "); t=p.with_suffix('.pending'); t.write_text(sys.stdin.read()); t.chmod(0o444); t.replace(p)"
140 subprocess.run(["ssh", "-p", PORT, "-o", "BatchMode=yes", HOST,
141 "python3 -c " + shlex.quote(script)], input=json.dumps(details) + "\n", text=True, check=True)
142 return release
143 ssh(f"mkdir -p {remote_release}")
144 subprocess.run(
145 ["rsync", "-a", "-e", f"ssh -p {PORT} -o BatchMode=yes",
146 *(str(snapshot / path) for path in SOURCES), str(snapshot / ".studio-release.json"),
147 f"{HOST}:{remote_release}/"], check=True,
148 )
149 ssh(f"chmod -R a-w {remote_release}")
150 return release
151
152
153def main():
154 parser = argparse.ArgumentParser(description="Preview working changes and deploy main")
155 parser.add_argument("mode", choices=["stage", "publish", "prod", "rollback", "history", "backups", "data-restore", "bootstrap", "allocate", "destroy", "secrets"])
156 parser.add_argument("target", nargs="?")
157 parser.add_argument("--env", action="append", default=[])
158 parser.add_argument("--file", type=Path)
159 parser.add_argument("--key", action="append", default=[])
160 parser.add_argument("--backup")
161 parser.add_argument("--discard-writes", action="store_true")
162 args = parser.parse_args()
163 if args.mode in {"stage", "allocate", "destroy", "secrets", "data-restore"} and (not args.target or not NAME.fullmatch(args.target)):
164 parser.error(f"{args.mode} requires a service or stage ID")
165 if args.mode == "rollback" and args.target and not RELEASE_ID.fullmatch(args.target):
166 parser.error("rollback target must be a release ID")
167 if args.mode in {"publish", "prod", "history", "backups", "bootstrap"} and args.target:
168 parser.error(f"{args.mode} takes no target")
169 if args.env and args.mode != "stage":
170 parser.error("--env is available only for stage")
171 if bool(args.file) != (args.mode == "secrets"):
172 parser.error("--file is required only for secrets")
173 if args.key and args.mode != "secrets":
174 parser.error("--key is available only for secrets")
175 if len(set(args.key)) != len(args.key) or any(not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key) for key in args.key):
176 parser.error("--key values must be unique secret names")
177 if args.mode == "data-restore":
178 if not args.backup or not re.fullmatch(r"\d{8}T\d{6}Z-[0-9a-f]{6}", args.backup) or not args.discard_writes:
179 parser.error("data-restore requires --backup ID and --discard-writes")
180 elif args.backup or args.discard_writes:
181 parser.error("--backup and --discard-writes are available only for data-restore")
182 if args.mode == "secrets":
183 if str(args.file) == "-":
184 values = sys.stdin.read()
185 else:
186 if args.file.stat().st_mode & 0o077:
187 raise ValueError("secret file must be readable only by its owner")
188 values = args.file.read_text()
189 if not values.strip():
190 raise ValueError("secret input is empty")
191 manager = f"python3 {REMOTE}/release.py"
192 if args.mode == "stage":
193 release = upload()
194 overrides = "".join(f" --env {shlex.quote(value)}" for value in args.env)
195 run_logged("stage", args.target, f"python3 {REMOTE}/releases/{release}/tools/studio.py stage {args.target}{overrides}")
196 elif args.mode == "secrets":
197 release = upload()
198 keys = "".join(f" --key {shlex.quote(key)}" for key in args.key)
199 subprocess.run(
200 ["ssh", "-p", PORT, "-o", "BatchMode=yes", HOST,
201 f"python3 {REMOTE}/releases/{release}/tools/studio.py secrets {args.target}{keys}"],
202 input=values, text=True, check=True,
203 )
204 elif args.mode == "bootstrap":
205 release = upload(main=True)
206 sync_manager(release)
207 ssh(f"{manager} publish {release}")
208 ssh(f"python3 {REMOTE}/releases/{release}/tools/studio.py pool")
209 ssh(f"python3 {REMOTE}/releases/{release}/tools/studio.py bootstrap")
210 ssh(f"python3 {REMOTE}/releases/{release}/tools/release.py bootstrap {release}")
211 elif args.mode == "allocate":
212 release = upload()
213 ssh(f"python3 {REMOTE}/releases/{release}/tools/studio.py allocate {args.target}")
214 elif args.mode in {"publish", "prod"}:
215 release = upload(main=True)
216 sync_manager(release)
217 ssh(f"{manager} publish {release}")
218 if args.mode == "prod":
219 run_logged("prod", release, f"{manager} deploy {release}")
220 elif args.mode == "rollback":
221 run_logged("rollback", args.target, f"{manager} rollback{(' ' + args.target) if args.target else ''}")
222 elif args.mode == "history":
223 ssh(f"{manager} history")
224 elif args.mode == "backups":
225 release = upload()
226 ssh(f"python3 {REMOTE}/releases/{release}/tools/data.py list")
227 elif args.mode == "data-restore":
228 release = upload()
229 run_logged("data-restore", args.target, f"python3 {REMOTE}/releases/{release}/tools/data.py restore {args.backup} {args.target} --discard-writes")
230 else:
231 metadata = json.loads(ssh(f"cat /var/lib/studio/stages/{args.target}.json", capture=True).stdout)
232 release = metadata.get("release")
233 if not isinstance(release, str) or not RELEASE_ID.fullmatch(release):
234 raise ValueError("stage has no successful release; inspect it before removal")
235 run_logged("destroy", args.target, f"python3 {REMOTE}/releases/{release}/tools/studio.py destroy {args.target}")
236
237
238if __name__ == "__main__":
239 try:
240 main()
241 except (OSError, ValueError, subprocess.CalledProcessError) as error:
242 print(f"Home server deployment stopped: {error}", file=sys.stderr)
243 sys.exit(error.returncode if isinstance(error, subprocess.CalledProcessError) else 1)