pgdevkit 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,234 @@
1
+ """
2
+ ComplexHelper — psycopg adapter for PostgreSQL composite types, enums, and JSONB.
3
+
4
+ Use this when your project has custom PostgreSQL types (composite types, enums)
5
+ that need to be registered with psycopg before inserting test data. Plain columns
6
+ and JSONB are handled automatically by pgdevkit.testdb.schema's _insert_test_data;
7
+ you only need this class for USER-DEFINED types, which that function does not
8
+ yet support natively — wrap the connection it's given before the INSERT:
9
+
10
+ from pgdevkit.testdb import schema # for reference, not a public extension point yet
11
+ from your_project.pg_complex_helper import ComplexHelper
12
+
13
+ async def insert_test_data_with_complex_types(json_file, table, force_reset, con):
14
+ helper = ComplexHelper(con)
15
+ schema_name, table_name = table.split(".")
16
+ complex_types = await helper.load_all_complex_types((schema_name, table_name))
17
+ rows = json.loads(json_file.read_text(encoding="utf-8"))
18
+ for row in rows:
19
+ for col, info in complex_types.items():
20
+ if col in row:
21
+ row[col] = await helper.recursive_convert(row[col], info, con)
22
+ # then insert `rows` the same way schema.py's _insert_test_data does
23
+ """
24
+
25
+ from psycopg import AsyncConnection
26
+ from psycopg.rows import dict_row
27
+ from typing import Any
28
+ from psycopg.sql import Identifier
29
+ from psycopg.types.composite import CompositeInfo, register_composite
30
+ from psycopg.types.json import Jsonb
31
+ from psycopg.types.enum import EnumInfo, register_enum
32
+
33
+
34
+ class ComplexHelper:
35
+ complex_types: dict[tuple[str, str], CompositeInfo | EnumInfo] = {}
36
+
37
+ def __init__(self, con: AsyncConnection):
38
+ self.con = con
39
+ self.system_complex_type_dict = None
40
+
41
+ self.registered: set[CompositeInfo | EnumInfo] = set()
42
+
43
+ async def load_complex_type_dict(self):
44
+ async with self.con.cursor(row_factory=dict_row) as cur:
45
+ await cur.execute("""
46
+ SELECT t.oid,
47
+ pg_catalog.format_type ( t.oid, NULL ) AS obj_name,
48
+ t.typtype
49
+ FROM pg_catalog.pg_type t
50
+ JOIN pg_catalog.pg_namespace n
51
+ ON n.oid = t.typnamespace
52
+ WHERE ( t.typrelid = 0
53
+ OR ( SELECT c.relkind = 'c'
54
+ FROM pg_catalog.pg_class c
55
+ WHERE c.oid = t.typrelid ) )
56
+ AND n.nspname <> 'pg_catalog'
57
+ AND n.nspname <> 'information_schema'
58
+ AND n.nspname !~ '^pg_toast'""")
59
+ system_complex_types = await cur.fetchall()
60
+ self.system_complex_type_dict = {
61
+ r["oid"]: (r["obj_name"], r["typtype"]) for r in system_complex_types
62
+ }
63
+
64
+ async def _load_complex_type_from_colinfos(
65
+ self, res: dict[str, Any] | None
66
+ ) -> CompositeInfo | EnumInfo | type[Jsonb] | None:
67
+ if not res:
68
+ return None
69
+ if res["data_type"].lower() == "jsonb":
70
+ return Jsonb
71
+ if res["data_type"].upper() == "ARRAY" and res["udt_name"] == "_jsonb":
72
+ return Jsonb
73
+ if (
74
+ not res["is_enum"]
75
+ and not res["is_user_defined"]
76
+ and not (res["data_type"] == "ARRAY" and res["udt_schema"] != "pg_catalog")
77
+ ):
78
+ return None
79
+ udt_schema: str = res["udt_schema"]
80
+ udt_name: str = res["udt_name"]
81
+ c = await self._get_complex_type(
82
+ f"{udt_schema}.{udt_name}", res["is_enum"], self.con
83
+ )
84
+ await self._recurse_register(c, self.con)
85
+ return c
86
+
87
+ async def load_all_complex_types(
88
+ self, table_name: tuple[str, str], include_generated: bool = False
89
+ ) -> dict[str, CompositeInfo | type[Jsonb] | EnumInfo | None]:
90
+ if self.system_complex_type_dict is None:
91
+ await self.load_complex_type_dict()
92
+ colquery = """
93
+ with enum_types as (
94
+ select n.nspname as enum_schema, t.typname as enum_name from pg_type t
95
+ inner join pg_namespace n on n.oid=t.typnamespace
96
+ where typtype='e'
97
+ )
98
+ select column_name, data_type,
99
+ data_type='USER-DEFINED' as is_user_defined,
100
+ udt_schema, udt_name,
101
+ e.enum_name is not null as is_enum
102
+ from information_schema.columns c
103
+ left join enum_types e on e.enum_schema=c.udt_schema and e.enum_name=c.udt_name
104
+ where table_schema=%(schema)s and table_name = %(tbl)s and (is_generated <> 'ALWAYS' or %(include_generated)s)"""
105
+ async with self.con.cursor(row_factory=dict_row) as cur:
106
+ await cur.execute(
107
+ colquery,
108
+ {
109
+ "schema": table_name[0],
110
+ "tbl": table_name[1],
111
+ "include_generated": include_generated,
112
+ },
113
+ )
114
+ res = await cur.fetchall()
115
+ return {
116
+ r["column_name"]: await self._load_complex_type_from_colinfos(r)
117
+ for r in res
118
+ }
119
+
120
+ async def load_complex_type(
121
+ self, table_name: tuple[str, str], col_name: str
122
+ ) -> CompositeInfo | type[Jsonb] | EnumInfo | None:
123
+ if self.system_complex_type_dict is None:
124
+ await self.load_complex_type_dict()
125
+ colquery = """
126
+ with enum_types as (
127
+ select n.nspname as enum_schema, t.typname as enum_name from pg_type t
128
+ inner join pg_namespace n on n.oid=t.typnamespace
129
+ where typtype='e'
130
+ )
131
+ select column_name, data_type,
132
+ data_type='USER-DEFINED' as is_user_defined,
133
+ udt_schema, udt_name,
134
+ e.enum_name is not null as is_enum
135
+ from information_schema.columns c
136
+ left join enum_types e on e.enum_schema=c.udt_schema and e.enum_name=c.udt_name
137
+ where table_schema=%(schema)s and table_name = %(tbl)s
138
+ and column_name = %(col)s"""
139
+ async with self.con.cursor(row_factory=dict_row) as cur:
140
+ await cur.execute(
141
+ colquery,
142
+ {"schema": table_name[0], "tbl": table_name[1], "col": col_name},
143
+ )
144
+ res = await cur.fetchone()
145
+
146
+ return await self._load_complex_type_from_colinfos(res)
147
+
148
+ async def _get_complex_type(
149
+ self, name: str, is_enum: bool, con: AsyncConnection
150
+ ) -> CompositeInfo | EnumInfo:
151
+ if name.endswith("[]"):
152
+ name = name[:-2]
153
+ schema, type_name = name.split(".")
154
+ if type_name.startswith(
155
+ "_"
156
+ ): # the array type in PostgreSQL starts with an underscore
157
+ type_name = type_name[1:]
158
+ if is_enum:
159
+ ci = await EnumInfo.fetch(con, Identifier(schema, type_name))
160
+ assert ci is not None, f"Enum type {name} not found in database"
161
+ self.complex_types[(schema, type_name)] = ci
162
+ if (schema, type_name) not in self.complex_types:
163
+ ci = await CompositeInfo.fetch(con, Identifier(schema, type_name))
164
+ assert ci is not None, f"Complex type {name} not found in database"
165
+ self.complex_types[(schema, type_name)] = ci
166
+ return self.complex_types[(schema, type_name)]
167
+
168
+ async def _recurse_register(
169
+ self, info: CompositeInfo | EnumInfo, con: AsyncConnection
170
+ ):
171
+ assert self.system_complex_type_dict is not None, (
172
+ "System complex type dictionary not loaded"
173
+ )
174
+ if info not in self.registered:
175
+ if isinstance(info, EnumInfo):
176
+ register_enum(info, con)
177
+ else:
178
+ register_composite(info, con)
179
+ self.registered.add(info)
180
+ if isinstance(info, EnumInfo):
181
+ return
182
+ for t in info.field_types:
183
+ if t in self.system_complex_type_dict:
184
+ name, typtype = self.system_complex_type_dict[t]
185
+ ci = await self._get_complex_type(name, typtype == "e", con)
186
+ await self._recurse_register(ci, con)
187
+
188
+ async def recursive_convert(
189
+ self,
190
+ value: Any,
191
+ info: CompositeInfo | EnumInfo | type[Jsonb] | None,
192
+ con: AsyncConnection,
193
+ ) -> Any:
194
+ if info is None:
195
+ return value
196
+ if value is None:
197
+ return None
198
+ if self.system_complex_type_dict is None:
199
+ await self.load_complex_type_dict()
200
+ if isinstance(value, list):
201
+ return [await self.recursive_convert(item, info, con) for item in value]
202
+ prms = {}
203
+ if info == Jsonb:
204
+ return Jsonb(value)
205
+ if isinstance(value, str):
206
+ assert isinstance(info, EnumInfo), f"Expected EnumInfo, got {type(info)}"
207
+ return getattr(info.enum, value) # Enum
208
+ assert isinstance(info, CompositeInfo), (
209
+ f"Expected CompositeInfo, got {type(info)}"
210
+ )
211
+ assert self.system_complex_type_dict is not None, (
212
+ "System complex type dictionary not loaded"
213
+ )
214
+ for k, v in value.items():
215
+ if v is None:
216
+ prms[k] = None
217
+ continue
218
+ fi = info.field_names.index(k)
219
+ type_oid = info.field_types[fi]
220
+ if type_oid in self.system_complex_type_dict:
221
+ name, typtype = self.system_complex_type_dict[type_oid]
222
+ ci = await self._get_complex_type(name, typtype == "e", con)
223
+ if name.endswith("[]"):
224
+ prms[k] = [
225
+ await self.recursive_convert(item, ci, con) for item in v
226
+ ]
227
+ else:
228
+ prms[k] = await self.recursive_convert(v, ci, con)
229
+ else:
230
+ prms[k] = v
231
+ assert info.python_type is not None, (
232
+ f"Python type for {info.name} is null, maybe an array?"
233
+ )
234
+ return info.python_type(**prms) if prms else None
@@ -0,0 +1,92 @@
1
+ # Dynamic SQL
2
+
3
+ Avoid dynamic SQL whenever possible — a static `.sql` file is always clearer. When column or table names genuinely vary at runtime, check `pyproject.toml` (or `[project] requires-python`) to pick an approach at authoring time, not with a runtime `sys.version_info` check.
4
+
5
+ ---
6
+
7
+ ## Python 3.14+ — t-string templates
8
+
9
+ T-strings look like f-strings but are evaluated by psycopg — values are always sent as bound parameters, never interpolated.
10
+
11
+ **Format specifiers:**
12
+
13
+ | Specifier | Meaning |
14
+ |-----------|---------|
15
+ | `{val}` / `{val:s}` | Bound parameter, automatic format (default) |
16
+ | `{val:b}` | Bound parameter, binary format |
17
+ | `{val:t}` | Bound parameter, text format |
18
+ | `{name:i}` | SQL identifier (table/column name) — double-quoted |
19
+ | `{val:l}` | Literal value merged client-side (use sparingly) |
20
+ | `{snippet:q}` | SQL snippet — another t-string or `sql.SQL`/`Composed` instance |
21
+
22
+ **Basic parameter:**
23
+
24
+ ```python
25
+ await cur.execute(t"SELECT * FROM users WHERE id = {user_id}")
26
+ ```
27
+
28
+ **Dynamic identifier (`:i`):**
29
+
30
+ ```python
31
+ column = "email"
32
+ await cur.execute(t"SELECT {column:i} FROM users WHERE active = {active}")
33
+ ```
34
+
35
+ **NOTIFY — requires client-side composition (`:i` and `:l`):**
36
+
37
+ ```python
38
+ def send_notify(conn: Connection, channel: str, payload: str) -> None:
39
+ conn.execute(t"NOTIFY {channel:i}, {payload:l}")
40
+ ```
41
+
42
+ **Nested templates with `:q` (dynamic WHERE clause):**
43
+
44
+ ```python
45
+ from psycopg import sql
46
+
47
+ def search_users(
48
+ conn: Connection,
49
+ ids: Sequence[int] | None = None,
50
+ name_pattern: str | None = None,
51
+ ) -> list[UserRow]:
52
+ filters = []
53
+ if ids is not None:
54
+ filters.append(t"u.id = ANY({list(ids)})")
55
+ if name_pattern is not None:
56
+ filters.append(t"u.name ~* {name_pattern}")
57
+ if not filters:
58
+ raise TypeError("at least one filter required")
59
+ joined = sql.SQL(" AND ").join(filters)
60
+ cur = conn.cursor(row_factory=class_row(UserRow))
61
+ cur.execute(t"SELECT * FROM users AS u WHERE {joined:q}")
62
+ return cur.fetchall()
63
+ ```
64
+
65
+ **Inspect composed SQL without executing:**
66
+
67
+ ```python
68
+ from psycopg import sql
69
+
70
+ name = "O'Reilly"
71
+ dob = datetime.date(1970, 1, 1)
72
+ print(sql.as_string(t"INSERT INTO tbl VALUES ({name}, {dob})"))
73
+ # INSERT INTO tbl VALUES ('O''Reilly', '1970-01-01'::date)
74
+ ```
75
+
76
+ ---
77
+
78
+ ## Python < 3.14 — `psycopg.sql`
79
+
80
+ ```python
81
+ from psycopg import sql
82
+
83
+ column = "email"
84
+ query = sql.SQL("SELECT {col} FROM users WHERE active = %(active)s").format(
85
+ col=sql.Identifier(column),
86
+ )
87
+ await cur.execute(query, {"active": True})
88
+ ```
89
+
90
+ `sql.Identifier` quotes the identifier at the driver level — SQL injection via column/table names is impossible. Values always stay as bound parameters.
91
+
92
+ **Never** use f-strings or string concatenation for identifiers or values.
@@ -0,0 +1,33 @@
1
+ # Temporal Tables (row-level history)
2
+
3
+ Use [nearform/temporal_tables](https://github.com/nearform/temporal_tables) — a pure PL/pgSQL trigger that archives every changed row to a parallel history table. Copy `versioning_function.sql` and `system_time_function.sql` from the repo into `db/migrations/` and apply them as a baseline migration.
4
+
5
+ For each versioned table, add `sys_period`, create the history table with `LIKE`, and attach the trigger:
6
+
7
+ ```sql
8
+ alter table users
9
+ add column sys_period tstzrange not null default tstzrange(current_timestamp, null);
10
+
11
+ create table users_history (like users);
12
+
13
+ create index on users_history using gist (id, sys_period);
14
+
15
+ create trigger users_versioning
16
+ before insert or update or delete on users
17
+ for each row execute procedure versioning(
18
+ 'sys_period', -- system-period column
19
+ 'users_history', -- history table name
20
+ true, -- conflict mitigation
21
+ false -- skip entry when no values changed
22
+ );
23
+ ```
24
+
25
+ `CREATE TABLE … LIKE` copies all columns and their types. Do not add a PK to the history table — the same `id` appears once per version.
26
+
27
+ Query point-in-time with `@>`:
28
+
29
+ ```sql
30
+ select * from users_history where id = %(id)s and sys_period @> %(as_of)s::timestamptz
31
+ ```
32
+
33
+ `set_system_time('2023-01-01+00')` / `set_system_time(null)` backdates operations (useful for migrations).
@@ -0,0 +1,3 @@
1
+ from .api import clean_testdb, dsn_for, ensure_testdb, reset_testdb, run_sql, status
2
+
3
+ __all__ = ["clean_testdb", "dsn_for", "ensure_testdb", "reset_testdb", "run_sql", "status"]
pgdevkit/testdb/api.py ADDED
@@ -0,0 +1,137 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ from pathlib import Path
5
+
6
+ import psycopg
7
+ from psycopg.sql import SQL, Identifier
8
+
9
+ from . import constants, query
10
+ from .config import ProjectConfig, load_config
11
+ from .container import ensure_container
12
+ from .naming import current_branch, slugify, workspace_db_name
13
+ from .schema import apply_schema
14
+
15
+
16
+ def _admin_dsn() -> str:
17
+ return f"postgresql://{constants.USER}:{constants.PASSWORD}@{constants.HOST}:{constants.PORT}/postgres?connect_timeout=10"
18
+
19
+
20
+ def _db_dsn(db_name: str) -> str:
21
+ return f"postgresql://{constants.USER}:{constants.PASSWORD}@{constants.HOST}:{constants.PORT}/{db_name}?connect_timeout=10"
22
+
23
+
24
+ def _resolve(project_root: Path | None) -> tuple[ProjectConfig, str]:
25
+ config = load_config(project_root)
26
+ branch = current_branch(config.root)
27
+ db_name = workspace_db_name(config.name, branch)
28
+ return config, db_name
29
+
30
+
31
+ def _env_for(config: ProjectConfig, db_name: str) -> dict[str, str]:
32
+ prefix = config.env_prefix
33
+ return {
34
+ f"{prefix}POSTGRES_HOST": constants.HOST,
35
+ f"{prefix}POSTGRES_PORT": str(constants.PORT),
36
+ f"{prefix}POSTGRES_DB": db_name,
37
+ f"{prefix}POSTGRES_USER": constants.USER,
38
+ f"{prefix}POSTGRES_PASSWORD": constants.PASSWORD,
39
+ }
40
+
41
+
42
+ async def _ensure_database(db_name: str) -> None:
43
+ async with await psycopg.AsyncConnection.connect(_admin_dsn(), autocommit=True) as con:
44
+ result = await con.execute("SELECT 1 FROM pg_database WHERE datname = %(db)s", {"db": db_name})
45
+ if await result.fetchone():
46
+ return
47
+ try:
48
+ await con.execute(SQL("CREATE DATABASE {}").format(Identifier(db_name)))
49
+ except psycopg.errors.DuplicateDatabase:
50
+ pass
51
+
52
+
53
+ async def _drop_database(db_name: str) -> None:
54
+ async with await psycopg.AsyncConnection.connect(_admin_dsn(), autocommit=True) as con:
55
+ await con.execute(
56
+ "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = %(db)s",
57
+ {"db": db_name},
58
+ )
59
+ await con.execute(SQL("DROP DATABASE IF EXISTS {}").format(Identifier(db_name)))
60
+
61
+
62
+ async def _apply(config: ProjectConfig, db_name: str, force_reset: bool) -> None:
63
+ async with await psycopg.AsyncConnection.connect(_db_dsn(db_name), autocommit=True) as con:
64
+ await apply_schema(
65
+ con,
66
+ config.root / config.database_dir,
67
+ extensions=config.extensions,
68
+ force_reset=force_reset,
69
+ )
70
+
71
+
72
+ def ensure_testdb(project_root: Path | None = None, force_reset: bool = False) -> dict[str, str]:
73
+ """Ensure the shared container is running, this workspace's database
74
+ exists, and its schema is applied. Returns the {PREFIX}POSTGRES_* env
75
+ vars for this workspace."""
76
+ ensure_container()
77
+ config, db_name = _resolve(project_root)
78
+
79
+ async def _run() -> None:
80
+ if force_reset:
81
+ await _drop_database(db_name)
82
+ await _ensure_database(db_name)
83
+ await _apply(config, db_name, force_reset)
84
+
85
+ asyncio.run(_run())
86
+ return _env_for(config, db_name)
87
+
88
+
89
+ def reset_testdb(project_root: Path | None = None) -> dict[str, str]:
90
+ """Drop and recreate only this workspace's database, then reapply
91
+ schema and seed data."""
92
+ return ensure_testdb(project_root, force_reset=True)
93
+
94
+
95
+ def clean_testdb(project_root: Path | None = None, all: bool = False) -> None:
96
+ """Drop this workspace's database. With all=True, drop every database
97
+ belonging to this project (matched by its name-slug prefix), across
98
+ every worktree/branch."""
99
+ config, db_name = _resolve(project_root)
100
+
101
+ async def _run() -> None:
102
+ if not all:
103
+ await _drop_database(db_name)
104
+ return
105
+ prefix = f"{slugify(config.name)}_"
106
+ escaped_prefix = prefix.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
107
+ async with await psycopg.AsyncConnection.connect(_admin_dsn(), autocommit=True) as con:
108
+ result = await con.execute(
109
+ "SELECT datname FROM pg_database WHERE datname LIKE %(pattern)s ESCAPE '\\'",
110
+ {"pattern": f"{escaped_prefix}%"},
111
+ )
112
+ names = [row[0] for row in await result.fetchall()]
113
+ for name in names:
114
+ await _drop_database(name)
115
+
116
+ asyncio.run(_run())
117
+
118
+
119
+ def status(project_root: Path | None = None) -> dict[str, str]:
120
+ config, db_name = _resolve(project_root)
121
+ return {
122
+ "container": constants.CONTAINER_NAME,
123
+ "host": constants.HOST,
124
+ "port": str(constants.PORT),
125
+ "database": db_name,
126
+ "dsn": _db_dsn(db_name),
127
+ }
128
+
129
+
130
+ def run_sql(sql: str, project_root: Path | None = None) -> list[dict] | None:
131
+ _, db_name = _resolve(project_root)
132
+ return asyncio.run(query.execute(_db_dsn(db_name), sql))
133
+
134
+
135
+ def dsn_for(project_root: Path | None = None) -> str:
136
+ _, db_name = _resolve(project_root)
137
+ return _db_dsn(db_name)
@@ -0,0 +1,51 @@
1
+ from __future__ import annotations
2
+
3
+ import tomllib
4
+ from dataclasses import dataclass, field
5
+ from pathlib import Path
6
+
7
+
8
+ @dataclass(frozen=True)
9
+ class ProjectConfig:
10
+ name: str
11
+ database_dir: str = "database"
12
+ env_prefix: str = ""
13
+ extensions: tuple[str, ...] = ()
14
+ root: Path = field(default_factory=Path)
15
+
16
+ def __post_init__(self) -> None:
17
+ if not self.env_prefix:
18
+ object.__setattr__(self, "env_prefix", f"{self.name.upper()}_")
19
+
20
+
21
+ def _find_pyproject(start: Path) -> Path | None:
22
+ for directory in [start, *start.parents]:
23
+ candidate = directory / "pyproject.toml"
24
+ if candidate.exists():
25
+ return candidate
26
+ return None
27
+
28
+
29
+ def load_config(start: Path | None = None) -> ProjectConfig:
30
+ start = (start or Path.cwd()).resolve()
31
+ pyproject = _find_pyproject(start)
32
+ root = pyproject.parent if pyproject else start
33
+
34
+ section: dict = {}
35
+ if pyproject is not None:
36
+ data = tomllib.loads(pyproject.read_text(encoding="utf-8"))
37
+ section = data.get("tool", {}).get("pgdevkit", {})
38
+
39
+ extensions = section.get("extensions", [])
40
+ if not isinstance(extensions, list):
41
+ raise TypeError(
42
+ f"[tool.pgdevkit].extensions in {pyproject} must be a list, got {type(extensions).__name__}"
43
+ )
44
+
45
+ return ProjectConfig(
46
+ name=section.get("name") or root.name,
47
+ database_dir=section.get("database_dir", "database"),
48
+ env_prefix=section.get("env_prefix", ""),
49
+ extensions=tuple(extensions),
50
+ root=root,
51
+ )
@@ -0,0 +1,11 @@
1
+ from __future__ import annotations
2
+
3
+ import os
4
+
5
+ CONTAINER_NAME = "pgdevkit-postgres"
6
+ IMAGE = "pgvector/pgvector:pg18-trixie"
7
+ HOST = os.environ.get("PGDEVKIT_TESTDB_HOST", "localhost")
8
+ PORT = int(os.environ.get("PGDEVKIT_TESTDB_PORT", "54322"))
9
+ USER = os.environ.get("PGDEVKIT_TESTDB_USER", "postgres")
10
+ PASSWORD = os.environ.get("PGDEVKIT_TESTDB_PASSWORD", "testpwd")
11
+ PG_SPEED_FLAGS = ["-c", "fsync=off", "-c", "synchronous_commit=off", "-c", "full_page_writes=off"]
@@ -0,0 +1,95 @@
1
+ from __future__ import annotations
2
+
3
+ import os
4
+ import subprocess
5
+ import time
6
+
7
+ import psycopg
8
+
9
+ from . import constants
10
+
11
+
12
+ def _podman(*args: str, check: bool = True) -> subprocess.CompletedProcess:
13
+ return subprocess.run(["podman", *args], capture_output=True, text=True, check=check)
14
+
15
+
16
+ def _available(timeout: float = 3.0) -> bool:
17
+ """Quick check (short timeout) for whether Postgres is already reachable
18
+ at HOST:PORT, so a database started outside pgdevkit's control (or the
19
+ container from a previous run) doesn't trigger another podman/docker
20
+ lifecycle call."""
21
+ # libpq's connect_timeout is whole seconds; anything below 1 means "wait
22
+ # indefinitely" instead of a short timeout, so it's clamped up to 1.
23
+ connect_timeout = max(1, round(timeout))
24
+ dsn = (
25
+ f"postgresql://{constants.USER}:{constants.PASSWORD}"
26
+ f"@{constants.HOST}:{constants.PORT}/postgres?connect_timeout={connect_timeout}"
27
+ )
28
+ try:
29
+ with psycopg.connect(dsn):
30
+ return True
31
+ except Exception: # noqa: BLE001
32
+ return False
33
+
34
+
35
+ def _container_status() -> str | None:
36
+ """Return 'running', 'exited', etc., or None if the container doesn't exist."""
37
+ result = _podman(
38
+ "inspect", constants.CONTAINER_NAME, "--format", "{{.State.Status}}", check=False
39
+ )
40
+ if result.returncode != 0:
41
+ return None
42
+ return result.stdout.strip()
43
+
44
+
45
+ def _create_container() -> None:
46
+ result = _podman(
47
+ "run", "-d",
48
+ "--name", constants.CONTAINER_NAME,
49
+ "-p", f"{constants.PORT}:5432",
50
+ "-e", f"POSTGRES_USER={constants.USER}",
51
+ "-e", f"POSTGRES_PASSWORD={constants.PASSWORD}",
52
+ constants.IMAGE,
53
+ "postgres", *constants.PG_SPEED_FLAGS,
54
+ check=False,
55
+ )
56
+ if result.returncode != 0 and "already in use" in result.stderr:
57
+ _podman("start", constants.CONTAINER_NAME)
58
+ return
59
+ if result.returncode != 0:
60
+ raise RuntimeError(f"podman run failed: {result.stderr}")
61
+
62
+
63
+ def _wait_ready(timeout: float = 30.0) -> None:
64
+ deadline = time.monotonic() + timeout
65
+ dsn = (
66
+ f"postgresql://{constants.USER}:{constants.PASSWORD}"
67
+ f"@{constants.HOST}:{constants.PORT}/postgres?connect_timeout=2"
68
+ )
69
+ last_error: Exception | None = None
70
+ while time.monotonic() < deadline:
71
+ try:
72
+ with psycopg.connect(dsn):
73
+ return
74
+ except Exception as e: # noqa: BLE001
75
+ last_error = e
76
+ time.sleep(0.5)
77
+ raise RuntimeError(f"Postgres did not become ready within {timeout}s: {last_error}")
78
+
79
+
80
+ def ensure_container() -> None:
81
+ """Idempotently ensure the shared pgdevkit-postgres container is running
82
+ and accepting connections. Never touches podman/docker if Postgres is
83
+ already reachable, or if PGDEVKIT_SKIP_CONTAINER says to assume it is."""
84
+ if os.environ.get("PGDEVKIT_SKIP_CONTAINER"):
85
+ return
86
+ if _available():
87
+ return
88
+ status = _container_status()
89
+ if status == "running":
90
+ return
91
+ if status is not None:
92
+ _podman("start", constants.CONTAINER_NAME)
93
+ else:
94
+ _create_container()
95
+ _wait_ready()