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