1"""Dev-side database policy on top of a reachable Postgres server.
2
3The SQL lives in `plain.postgres.databases`. What lives here is everything that
4is dev's opinion rather than Postgres' mechanism: what metadata a dev database
5carries, and how to fork one without disrupting whoever is using the source.
6"""
7
8from __future__ import annotations
9
10import json
11import subprocess
12import time
13from dataclasses import dataclass
14from pathlib import Path
15from typing import TYPE_CHECKING, Any
16
17from .backends import Server
18from .identity import current_branch
19
20if TYPE_CHECKING:
21 from plain.postgres.database_url import DatabaseConfig
22
23
24@dataclass(frozen=True)
25class DevDatabase:
26 """A database in the project's cluster, with its dev metadata."""
27
28 name: str
29 checkout: str | None
30 branch: str | None
31 created_via: str | None
32 size_bytes: int
33 is_test: bool = False
34
35 @property
36 def checkout_exists(self) -> bool:
37 return self.checkout is not None and Path(self.checkout).exists()
38
39
40@dataclass(frozen=True)
41class Cluster:
42 """The project's Postgres server, as dev sees it."""
43
44 server: Server
45
46 def url(self, db_name: str) -> str:
47 return self.server.url(db_name)
48
49 @property
50 def config(self) -> DatabaseConfig:
51 from plain.postgres.database_url import parse_database_url
52
53 return parse_database_url(self.url("postgres"))
54
55 # -- lifecycle ---------------------------------------------------------
56
57 def database_exists(self, name: str) -> bool:
58 from plain.postgres.databases import database_exists
59
60 return database_exists(self.config, name=name)
61
62 def connection_count(self, name: str) -> int:
63 from plain.postgres.databases import connection_count
64
65 return connection_count(self.config, name=name)
66
67 def create_database(self, name: str, *, template: str | None = None) -> None:
68 from plain.postgres.databases import create_database
69
70 create_database(self.config, name=name, template=template)
71
72 def drop_database(self, name: str) -> None:
73 from plain.postgres.databases import drop_database
74
75 drop_database(self.config, name=name, force=True)
76
77 # -- metadata ----------------------------------------------------------
78
79 def set_metadata(self, name: str, metadata: dict[str, Any]) -> None:
80 """Stamp dev metadata onto the database itself.
81
82 Postgres stores this cluster-wide in `pg_shdescription`, which is why
83 there's no registry file: the metadata travels with the database, can't
84 drift from the real database list, and disappears when it's dropped.
85 """
86 from plain.postgres.databases import set_database_comment
87
88 set_database_comment(self.config, name=name, comment=json.dumps(metadata))
89
90 def get_metadata(self, name: str) -> dict[str, Any] | None:
91 from plain.postgres.databases import get_database_comment
92
93 return _decode_metadata(get_database_comment(self.config, name=name))
94
95 def record_created(
96 self,
97 name: str,
98 *,
99 checkout: str | None,
100 created_via: str,
101 project_root: Path,
102 ) -> None:
103 """Stamp a database with who made it, how, and on which branch.
104
105 The one place the metadata schema is spelled out — every creation path
106 (create, fork, reset, ensure, guard) records the same three facts.
107 """
108 self.set_metadata(
109 name,
110 {
111 "checkout": checkout,
112 "branch": current_branch(project_root),
113 "created_via": created_via,
114 },
115 )
116
117 def update_metadata(self, name: str, **changes: Any) -> None:
118 metadata = self.get_metadata(name) or {}
119 metadata.update(changes)
120 self.set_metadata(name, metadata)
121
122 def list_databases(self, project_name: str) -> list[DevDatabase]:
123 """Every database in this project, however it got its name.
124
125 Name prefix alone isn't enough: `plain db fork scratch` produces a
126 perfectly real database that doesn't start with the project name, and a
127 database you can't see is one you can't drop. So a database counts as
128 ours if it carries our metadata *or* it's named the way we name them.
129
130 Test databases (`test_{project}` / `test_{project}_{checkout}`) count
131 too, even though they carry no metadata. They're dropped on normal exit,
132 so one that exists at rest is usually debris from a crashed run — and
133 invisible debris is unreclaimable debris.
134
135 Going the other way — listing everything on the cluster — is wrong for
136 the local backend, where one server holds databases we never created.
137 """
138 from plain.postgres.databases import list_databases
139
140 databases = []
141 for info in list_databases(self.config):
142 metadata = _decode_metadata(info.comment) or {}
143 ours = "created_via" in metadata
144 named_like_ours = info.name == project_name or info.name.startswith(
145 f"{project_name}_"
146 )
147 is_test = info.name == f"test_{project_name}" or info.name.startswith(
148 f"test_{project_name}_"
149 )
150 if not (ours or named_like_ours or is_test):
151 continue
152 databases.append(
153 DevDatabase(
154 name=info.name,
155 checkout=metadata.get("checkout"),
156 branch=metadata.get("branch"),
157 created_via=metadata.get("created_via"),
158 size_bytes=info.size_bytes,
159 is_test=is_test,
160 )
161 )
162 return databases
163
164 # -- forking -----------------------------------------------------------
165
166 def fork_database(self, source: str, dest: str, *, force: bool = False) -> str:
167 """Copy `source` to `dest` with its data. Returns the mechanism used.
168
169 `CREATE DATABASE … TEMPLATE` is a file-level copy — near-instant at any
170 data size — but Postgres refuses it while anything else is connected to
171 the source. In the agentic case the source main is often running its own
172 `plain dev`, so we pick by whether it's actually busy rather than
173 forcing everyone off it:
174
175 - idle source → TEMPLATE
176 - busy source → `pg_dump | pg_restore`, which takes an MVCC snapshot
177 and never blocks the source
178 - `force` → terminate the source's connections and use TEMPLATE
179 once it drains; if it won't drain in time, fall back to dump-restore
180 """
181 from plain.postgres.databases import terminate_connections
182
183 busy = self.connection_count(source) > 0
184 if busy and force:
185 terminate_connections(self.config, name=source)
186 # pg_terminate_backend only signals; backends exit asynchronously,
187 # and a running `plain dev` pool may reconnect. Wait briefly for it
188 # to drain — if it doesn't, the dump-restore path below handles a
189 # busy source anyway.
190 deadline = time.monotonic() + 3
191 while time.monotonic() < deadline and self.connection_count(source) > 0:
192 time.sleep(0.1)
193 busy = self.connection_count(source) > 0
194
195 if not busy:
196 self.create_database(dest, template=source)
197 return "template"
198
199 self.create_database(dest)
200 try:
201 self._dump_restore(source, dest)
202 except Exception:
203 self.drop_database(dest) # don't strand a half-copied database
204 raise
205 return "dump-restore"
206
207 def _dump_restore(self, source: str, dest: str) -> None:
208 """Logical copy, streaming the dump straight into the restore.
209
210 Two processes joined by a real pipe, both exit codes checked: a
211 `pg_dump` that dies mid-stream would otherwise be masked by
212 `pg_restore`'s status, and `--exit-on-error` stops `pg_restore` from
213 finishing "successfully" over a truncated dump. Run where the server's
214 own binaries live so tool and server versions match — see
215 `Server.server_command`.
216 """
217 server = self.server
218 user = server.user
219 port = str(server.internal_port)
220 dump_argv, dump_env = server.server_command(
221 "pg_dump", "-U", user, "-h", "127.0.0.1", "-p", port, "-Fc", source
222 )
223 restore_argv, restore_env = server.server_command(
224 "pg_restore",
225 "-U",
226 user,
227 "-h",
228 "127.0.0.1",
229 "-p",
230 port,
231 "--no-owner",
232 "--exit-on-error",
233 "-d",
234 dest,
235 stdin=True,
236 )
237
238 # pg_dump's stdout is the binary custom-format dump, so it must not be
239 # decoded — only stderr is text here.
240 dump = subprocess.Popen(
241 dump_argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=dump_env
242 )
243 assert dump.stdout is not None
244 assert dump.stderr is not None
245 restore = subprocess.run(
246 restore_argv,
247 stdin=dump.stdout,
248 capture_output=True,
249 text=True,
250 env=restore_env,
251 )
252 dump.stdout.close() # let pg_dump see EOF/SIGPIPE if restore exited early
253 dump_stderr = dump.stderr.read().decode(errors="replace")
254 dump.stderr.close()
255 dump_rc = dump.wait()
256
257 if dump_rc != 0:
258 raise RuntimeError(
259 f"pg_dump of {source!r} failed (exit {dump_rc}): "
260 f"{dump_stderr.strip()[-2000:]}"
261 )
262 if restore.returncode != 0:
263 raise RuntimeError(
264 f"pg_restore into {dest!r} failed (exit {restore.returncode}): "
265 f"{restore.stderr.strip()[-2000:]}"
266 )
267
268
269def _decode_metadata(comment: str | None) -> dict[str, Any] | None:
270 if not comment:
271 return None
272 try:
273 decoded = json.loads(comment)
274 except ValueError:
275 return None
276 return decoded if isinstance(decoded, dict) else None