1"""Where the managed Postgres server actually comes from.
2
3Two backends, chosen in this order when `backend = "auto"`:
4
51. **Docker** — a container per project. Preferred on a workstation: the PG
6 version is pinnable, data is isolated per project, and nothing contends for
7 port 5432.
82. **Local** — a Postgres already listening on 127.0.0.1:5432. This is the
9 path that makes remote/agent sandboxes work, where a Docker daemon usually
10 isn't available but a system Postgres often is. Databases stay namespaced by
11 project, so sharing one cluster is safe even though it's less isolated.
12
13If neither is available we return `None` and say why. Never start a Docker
14daemon, and never install Postgres.
15"""
16
17from __future__ import annotations
18
19import json
20import os
21import re
22import shutil
23import socket
24import subprocess
25import time
26from dataclasses import dataclass
27
28import click
29
30from .identity import DEV_PASSWORD, DEV_USER
31
32LOCAL_PORT = 5432
33
34
35@dataclass(frozen=True)
36class Server:
37 """A reachable Postgres server we can create databases on."""
38
39 host: str
40 port: int
41 user: str
42 password: str
43 backend: str # "docker" | "local"
44 container: str | None = None # set for the docker backend
45
46 def url(self, db_name: str) -> str:
47 return (
48 f"postgres://{self.user}:{self.password}@{self.host}:{self.port}/{db_name}"
49 )
50
51 @property
52 def internal_port(self) -> int:
53 """The port Postgres listens on *from where its binaries run*.
54
55 Inside a container that's always 5432 regardless of what Docker
56 published on the host; for a local server it's the host port.
57 """
58 return 5432 if self.container else self.port
59
60 def server_command(
61 self, *args: str, stdin: bool = False
62 ) -> tuple[list[str], dict[str, str]]:
63 """Argv + env to run a Postgres binary *where the server's binaries live*.
64
65 For Docker that's inside the container (pass `stdin=True` when the
66 command reads stdin, so `docker exec -i` attaches it), which guarantees
67 `pg_dump`/`pg_restore` match the server version — the host may not have
68 them at all. For a local server they're on the host by definition.
69 """
70 if self.container:
71 argv = [
72 "docker",
73 "exec",
74 *(["-i"] if stdin else []),
75 "-e",
76 f"PGPASSWORD={self.password}",
77 self.container,
78 *args,
79 ]
80 return argv, dict(os.environ)
81 return [*args], {**os.environ, "PGPASSWORD": self.password}
82
83
84def port_is_open(port: int, *, host: str = "127.0.0.1", timeout: float = 0.25) -> bool:
85 """Is something listening? ~0.2ms locally, so it's free to call on any command."""
86 with socket.socket() as sock:
87 sock.settimeout(timeout)
88 try:
89 sock.connect((host, port))
90 except OSError:
91 return False
92 return True
93
94
95# --------------------------------------------------------------------------
96# Docker
97# --------------------------------------------------------------------------
98
99
100def _docker(*args: str, check: bool = True) -> subprocess.CompletedProcess[str]:
101 return subprocess.run(
102 ["docker", *args], capture_output=True, text=True, check=check
103 )
104
105
106def docker_available() -> bool:
107 """Is the Docker CLI present *and* the daemon responding?"""
108 if not shutil.which("docker"):
109 return False
110 try:
111 return _docker("info", check=False).returncode == 0
112 except OSError:
113 return False
114
115
116@dataclass(frozen=True)
117class ContainerState:
118 """What one `docker inspect` tells us about an existing container."""
119
120 running: bool
121 image: str
122 bound_host_ip: str # "" means all interfaces
123 host_port: int | None # the assigned host port for 5432/tcp; None if not running
124 max_connections: int | None # from the `-c max_connections=N` launch arg
125
126
127def _inspect_container(name: str) -> ContainerState | None:
128 """Everything we ask about a container, in a single `docker` call.
129
130 Each `docker` invocation costs ~100-300ms under Docker Desktop, and the
131 questions (exists? running? which image? bound where? which port?) always
132 travel together — separate probes made opening a cluster the slowest part
133 of every `plain db` command.
134 """
135 result = _docker("inspect", name, check=False)
136 if result.returncode != 0:
137 return None # no such container
138
139 data = json.loads(result.stdout)[0]
140 bindings = (data["HostConfig"].get("PortBindings") or {}).get("5432/tcp") or []
141
142 # The *assigned* host port lives here once the container is running;
143 # HostConfig.PortBindings only holds the requested "0" (pick one for me).
144 published = (data["NetworkSettings"].get("Ports") or {}).get("5432/tcp") or []
145 host_port = (
146 int(published[0]["HostPort"])
147 if published and published[0].get("HostPort")
148 else None
149 )
150
151 # max_connections is fixed at creation via `-c max_connections=N`.
152 cmd = data["Config"].get("Cmd") or []
153 max_connections = None
154 for i, arg in enumerate(cmd):
155 if (
156 arg == "-c"
157 and i + 1 < len(cmd)
158 and cmd[i + 1].startswith("max_connections=")
159 ):
160 max_connections = int(cmd[i + 1].split("=", 1)[1])
161 break
162
163 return ContainerState(
164 running=data["State"]["Running"],
165 image=data["Config"]["Image"],
166 bound_host_ip=bindings[0].get("HostIp", "") if bindings else "",
167 host_port=host_port,
168 max_connections=max_connections,
169 )
170
171
172def _container_host_port(name: str) -> int:
173 lines = _docker("port", name, "5432/tcp").stdout.strip().splitlines()
174 if not lines:
175 raise RuntimeError(
176 f"Container {name!r} has no published port for 5432/tcp. "
177 f"It wasn't created by plain-dev in its current form — recreate it: "
178 f"docker rm -f {name}"
179 )
180 return int(lines[0].rsplit(":", 1)[1])
181
182
183def _warn_if_container_drifted(
184 container: str,
185 state: ContainerState,
186 *,
187 image: str,
188 max_connections: int,
189) -> None:
190 """Flag an existing container whose creation-time settings no longer match.
191
192 Image, the localhost binding, and `max_connections` are all fixed when the
193 container is created — changing them in code (or picking up a newer
194 plain-dev) doesn't touch a container that already exists. We collect every
195 drift and print the one shared remedy: recreate it. That's cheap, because
196 the data lives on a separate volume and survives — except across a Postgres
197 *major* version bump, where a dump and reload is needed instead.
198 """
199 lines: list[str] = []
200 major_version_caution = False
201
202 if state.image and state.image != image:
203 lines.append(
204 f" Image is set to {image!r} but the running container was built "
205 f"from {state.image!r}."
206 )
207 major_version_caution = True
208
209 if state.bound_host_ip in ("", "0.0.0.0"):
210 lines.append(
211 " It's published on all network interfaces, but the dev "
212 "credentials are fixed — anyone on the same network can log in."
213 )
214
215 if state.max_connections is not None and state.max_connections != max_connections:
216 lines.append(
217 f" max_connections is {state.max_connections}, code wants "
218 f"{max_connections}."
219 )
220
221 if not lines:
222 return
223
224 click.secho(
225 f"Postgres container {container!r} has drifted from its configuration:",
226 fg="yellow",
227 err=True,
228 )
229 for line in lines:
230 click.echo(line, err=True)
231 click.echo(
232 f" Recreate it to pick these up; your data is on a separate volume and "
233 f"survives:\n"
234 f" docker rm -f {container}",
235 err=True,
236 )
237 if major_version_caution:
238 click.echo(
239 " If this is a different Postgres major version, back up first — a "
240 "new major can't read the old data directory:\n"
241 ' pg_dump -Fc "$(plain db url)" > backup.dump',
242 err=True,
243 )
244
245
246def _create_container(*args: str) -> None:
247 """`docker run`, tolerating another process having just done the same.
248
249 Several worktrees of one project share a container, so two of them starting
250 together will both find it missing and both try to create it. The loser gets
251 a name conflict — which means the container it wanted now exists, so there is
252 nothing to report.
253 """
254 try:
255 _docker(*args)
256 except subprocess.CalledProcessError as e:
257 if "already in use" not in (e.stderr or ""):
258 raise
259
260
261def start_docker_server(
262 *, container: str, volume: str, image: str, max_connections: int
263) -> Server:
264 """Ensure the project's container is running and return how to reach it.
265
266 Published as `-p 127.0.0.1:0:5432` so Docker picks a free host port — no
267 port registry, no contention between projects — and binds it to localhost
268 only. The credentials are fixed dev credentials (`postgres`/`postgres`), so
269 the cluster must never be reachable from another machine on the network.
270 The container is the durable handle; the port is read back from it.
271 """
272 state = _inspect_container(container)
273 if state is None:
274 _create_container(
275 "run",
276 "-d",
277 "--name",
278 container,
279 # Deliberately no restart policy. A stopped container is started on
280 # demand (~2s) by the branch below, so auto-starting every project's
281 # Postgres at every boot would only buy those couple of seconds — at
282 # a cost of ~76MB of RAM per project, permanently, for projects you
283 # may not have touched in months.
284 "-e",
285 f"POSTGRES_USER={DEV_USER}",
286 "-e",
287 f"POSTGRES_PASSWORD={DEV_PASSWORD}",
288 "-v",
289 f"{volume}:/var/lib/postgresql/data",
290 "-p",
291 "127.0.0.1:0:5432",
292 image,
293 # One cluster serves every worktree's pool plus xdist workers, so
294 # PG's default of 100 runs out fast.
295 "-c",
296 f"max_connections={max_connections}",
297 )
298 port = _container_host_port(container)
299 else:
300 _warn_if_container_drifted(
301 container, state, image=image, max_connections=max_connections
302 )
303 if state.running and state.host_port is not None:
304 port = state.host_port # already parsed from the inspect call
305 else:
306 if not state.running:
307 _docker("start", container)
308 # Ports is only populated once running, so read it back now.
309 port = _container_host_port(container)
310
311 wait_until_ready(port)
312 return Server(
313 host="127.0.0.1",
314 port=port,
315 user=DEV_USER,
316 password=DEV_PASSWORD,
317 backend="docker",
318 container=container,
319 )
320
321
322def stop_container(container: str) -> bool:
323 """Stop the container. Returns False if it wasn't running."""
324 state = _inspect_container(container)
325 if state is None or not state.running:
326 return False
327 _docker("stop", container, check=False)
328 return True
329
330
331def remove_container(container: str, *, volume: str | None = None) -> None:
332 """Remove the container, and optionally the volume holding its data."""
333 _docker("rm", "-f", container, check=False)
334 if volume:
335 _docker("volume", "rm", volume, check=False)
336
337
338@dataclass(frozen=True)
339class ManagedContainer:
340 """A Postgres container plain-dev created, on this machine."""
341
342 name: str
343 running: bool
344 image: str
345 size: str
346
347
348# The names `cluster_name()` produces: `plain-postgres-{project}-{8 hex}`.
349# Docker's `--filter name=` is a substring match, so it happily returns
350# containers we never created — a hand-rolled `plain-postgres-dev` from some
351# older script matches too. Listing someone else's container as ours invites
352# `plain db server remove` on it, so the shape is checked here.
353MANAGED_CONTAINER_RE = re.compile(r"^plain-postgres-.+-[0-9a-f]{8}$")
354
355
356def list_managed_containers() -> list[ManagedContainer]:
357 """Every plain-dev Postgres container here, not just this project's.
358
359 Containers are per project, so they accumulate as you work on more of them.
360 Nothing removes one automatically — a container we can't reach the project
361 for might still hold the only copy of something — so this exists to make the
362 pile visible rather than to guess at it.
363 """
364 result = _docker(
365 "ps",
366 "-a",
367 "--filter",
368 "name=plain-postgres-",
369 "--format",
370 "{{.Names}}\t{{.State}}\t{{.Image}}\t{{.Size}}",
371 check=False,
372 )
373 containers = []
374 for line in result.stdout.strip().splitlines():
375 parts = line.split("\t")
376 if len(parts) != 4:
377 continue
378 name, state, image, size = parts
379 if not MANAGED_CONTAINER_RE.match(name):
380 continue
381 containers.append(
382 ManagedContainer(
383 name=name, running=state == "running", image=image, size=size
384 )
385 )
386 return sorted(containers, key=lambda c: c.name)
387
388
389# --------------------------------------------------------------------------
390# Local
391# --------------------------------------------------------------------------
392
393
394def local_available() -> bool:
395 """Is there a local Postgres we can actually log into?
396
397 An open port is not enough. Homebrew and Postgres.app — the two most common
398 ways to have Postgres on a Mac — set up a superuser named after your OS
399 account with trust auth and no `postgres` role at all. Treating "something
400 is listening" as availability picked that server and then failed on every
401 connection, with nothing saying why. So the probe is the real question:
402 can we authenticate as the user we're going to use?
403 """
404 if not port_is_open(LOCAL_PORT):
405 return False
406
407 import psycopg
408
409 try:
410 psycopg.connect(
411 host="127.0.0.1",
412 port=LOCAL_PORT,
413 user=DEV_USER,
414 password=DEV_PASSWORD,
415 dbname="postgres",
416 connect_timeout=2,
417 ).close()
418 except psycopg.OperationalError:
419 return False
420 return True
421
422
423def local_rejected_us() -> bool:
424 """Is a local Postgres listening, but not one we can use?
425
426 Distinguishes "you have no Postgres" from "you have one and we can't log
427 into it" — different problems with different fixes, and only the second is
428 worth telling someone their existing server is the reason.
429 """
430 return port_is_open(LOCAL_PORT) and not local_available()
431
432
433def connect_local_server() -> Server:
434 """Use a Postgres already running on 5432.
435
436 Credentials are the conventional `postgres`/`postgres`; if the local server
437 wants something else, that's a bring-your-own case and the user should set
438 `PLAIN_POSTGRES_URL` instead.
439 """
440 return Server(
441 host="127.0.0.1",
442 port=LOCAL_PORT,
443 user=DEV_USER,
444 password=DEV_PASSWORD,
445 backend="local",
446 )
447
448
449# --------------------------------------------------------------------------
450# Readiness
451# --------------------------------------------------------------------------
452
453
454def wait_until_ready(port: int, *, timeout: float = 60.0) -> None:
455 """Block until Postgres accepts a real connection, not just a TCP handshake.
456
457 A fresh container listens on the port well before `initdb` finishes, so the
458 TCP probe alone would hand back a server that refuses queries.
459 """
460 import psycopg
461
462 deadline = time.monotonic() + timeout
463 last_error: Exception | None = None
464 while time.monotonic() < deadline:
465 try:
466 psycopg.connect(
467 host="127.0.0.1",
468 port=port,
469 user=DEV_USER,
470 password=DEV_PASSWORD,
471 dbname="postgres",
472 connect_timeout=2,
473 ).close()
474 return
475 except psycopg.OperationalError as e:
476 last_error = e
477 time.sleep(0.5)
478
479 raise TimeoutError(
480 f"Postgres on port {port} did not become ready within {timeout:.0f}s: {last_error}"
481 )