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.
- pgdevkit/__init__.py +0 -0
- pgdevkit/cli.py +217 -0
- pgdevkit/connection.py +56 -0
- pgdevkit/db/__init__.py +37 -0
- pgdevkit/db/connection.py +74 -0
- pgdevkit/db/crud.py +184 -0
- pgdevkit/db/loader.py +19 -0
- pgdevkit/db/model.py +23 -0
- pgdevkit/diff.py +253 -0
- pgdevkit/docs/database-layout.md +145 -0
- pgdevkit/fetch_missing.py +229 -0
- pgdevkit/introspect.py +190 -0
- pgdevkit/lakebase.py +90 -0
- pgdevkit/models.py +102 -0
- pgdevkit/parser.py +336 -0
- pgdevkit/skills/pgdevkit/SKILL.md +283 -0
- pgdevkit/skills/pgdevkit/references/complex_helper.py +234 -0
- pgdevkit/skills/pgdevkit/references/dynamic-sql.md +92 -0
- pgdevkit/skills/pgdevkit/references/temporal-tables.md +33 -0
- pgdevkit/testdb/__init__.py +3 -0
- pgdevkit/testdb/api.py +137 -0
- pgdevkit/testdb/config.py +51 -0
- pgdevkit/testdb/constants.py +11 -0
- pgdevkit/testdb/container.py +95 -0
- pgdevkit/testdb/naming.py +41 -0
- pgdevkit/testdb/query.py +65 -0
- pgdevkit/testdb/schema.py +188 -0
- pgdevkit-0.1.0.dist-info/METADATA +140 -0
- pgdevkit-0.1.0.dist-info/RECORD +31 -0
- pgdevkit-0.1.0.dist-info/WHEEL +4 -0
- pgdevkit-0.1.0.dist-info/entry_points.txt +2 -0
|
@@ -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).
|
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()
|