1from __future__ import annotations
2
3from abc import ABC, abstractmethod
4from dataclasses import dataclass
5from typing import TYPE_CHECKING, ClassVar
6
7import psycopg
8import psycopg.sql
9
10from plain.logs import get_framework_logger
11from plain.runtime import settings as plain_settings
12
13from ..constraints import BaseConstraint, CheckConstraint, UniqueConstraint
14from ..db import get_connection
15from ..dialect import build_timeout_set_clauses, quote_name
16from ..indexes import Index
17
18if TYPE_CHECKING:
19 from ..base import Model
20
21
22logger = get_framework_logger()
23
24
25def _convergence_prelude(*, blocking: bool, local: bool) -> str:
26 """Return the SET prelude for convergence DDL.
27
28 `blocking=True` (ACCESS EXCLUSIVE) sets both lock_timeout and
29 statement_timeout. `blocking=False` (SHARE UPDATE EXCLUSIVE — VALIDATE,
30 CONCURRENTLY) omits statement_timeout so the non-blocking statement can
31 run to completion on any table size.
32 """
33 return build_timeout_set_clauses(
34 lock_timeout=plain_settings.POSTGRES_CONVERGENCE_LOCK_TIMEOUT,
35 statement_timeout=(
36 plain_settings.POSTGRES_CONVERGENCE_STATEMENT_TIMEOUT if blocking else None
37 ),
38 local=local,
39 )
40
41
42def _execute_and_commit(sql: str | list[str], *, blocking: bool = True) -> None:
43 """Execute DDL in a committed transaction with convergence timeouts.
44
45 Accepts a single SQL string or a list of statements that must share one
46 transaction (e.g. SetNotNullCorrection step 3: SET NOT NULL + DROP temp check
47 together so no orphan constraint can remain after a partial failure).
48 """
49 prelude = _convergence_prelude(blocking=blocking, local=True)
50 if isinstance(sql, str):
51 script = prelude + sql
52 else:
53 script = prelude + "; ".join(sql)
54
55 # psycopg3 simple-query protocol: a single execute() with a multi-statement
56 # script runs every statement in one transaction. params must be None —
57 # these are literal DDL strings assembled above.
58 conn = get_connection()
59 try:
60 with conn.cursor() as cursor:
61 cursor.execute(script)
62 conn.commit()
63 except Exception:
64 conn.rollback()
65 raise
66
67
68def _execute_autocommit(sql: str) -> None:
69 """Execute DDL in autocommit mode (required for CONCURRENTLY operations).
70
71 CONCURRENTLY holds SHARE UPDATE EXCLUSIVE (non-blocking). Only
72 lock_timeout is set — statement_timeout is omitted so the operation can
73 run to completion on any table size.
74
75 Autocommit forbids SET LOCAL, and bundling SET with CONCURRENTLY in a
76 single query forms an implicit transaction block that CONCURRENTLY
77 rejects. So SET, DDL, and RESET are three separate cursor.execute()
78 calls. RESET runs in a finally so the session-level timeout doesn't leak
79 to the next caller even if the DDL fails.
80 """
81 conn = get_connection()
82 if conn.in_atomic_block:
83 raise RuntimeError("Cannot use CONCURRENTLY inside an atomic block")
84 old_autocommit = conn.get_autocommit()
85 if not old_autocommit:
86 conn.commit()
87 conn.set_autocommit(True)
88 try:
89 # Session-level SET (no LOCAL — autocommit has no transaction to
90 # scope to). Routed through the dialect helper so the setting value
91 # is validated the same way as every other timeout SET.
92 set_prelude = build_timeout_set_clauses(
93 lock_timeout=plain_settings.POSTGRES_CONVERGENCE_LOCK_TIMEOUT,
94 statement_timeout=None,
95 local=False,
96 )
97 try:
98 with conn.cursor() as cursor:
99 cursor.execute(set_prelude)
100 cursor.execute(sql)
101 finally:
102 # Clear the session-level timeouts. Both are RESET even though
103 # only lock_timeout is set today, so a future change that
104 # introduces statement_timeout on this path can't silently leak
105 # it to the next caller. RESET of an unset value is a no-op. If
106 # RESET itself fails, close the connection so the pool can't
107 # hand it back with leaked session state.
108 try:
109 with conn.cursor() as cursor:
110 cursor.execute("RESET lock_timeout; RESET statement_timeout")
111 except Exception:
112 logger.warning(
113 "Failed to RESET session timeouts after autocommit DDL; "
114 "closing connection to avoid leaking session state",
115 exc_info=True,
116 )
117 conn.close()
118 finally:
119 # Restore autocommit only if the connection is still open — if we
120 # closed it above due to RESET failure, calling set_autocommit would
121 # force a reconnect just to flip a flag, and any failure there would
122 # mask the original DDL error.
123 if not old_autocommit and conn.connection is not None:
124 conn.set_autocommit(False)
125
126
127class Correction(ABC):
128 """Concrete executable SQL operation for convergence."""
129
130 pass_order: ClassVar[int]
131
132 @abstractmethod
133 def describe(self) -> str: ...
134
135 @abstractmethod
136 def apply(self) -> str: ...
137
138
139@dataclass
140class RebuildIndexCorrection(Correction):
141 """Drop an INVALID index and recreate it CONCURRENTLY."""
142
143 pass_order = 0
144
145 table: str
146 index: Index
147 model: type[Model]
148
149 def describe(self) -> str:
150 return f"{self.table}: rebuild index {self.index.name}"
151
152 def apply(self) -> str:
153 drop_sql = f"DROP INDEX CONCURRENTLY IF EXISTS {quote_name(self.index.name)}"
154 _execute_autocommit(drop_sql)
155 create_sql = self.index.to_sql(self.model)
156 _execute_autocommit(create_sql)
157 return f"{drop_sql}; {create_sql}"
158
159
160@dataclass
161class RenameIndexCorrection(Correction):
162 """Rename an index (catalog-only, instant)."""
163
164 pass_order = 1
165
166 table: str
167 old_name: str
168 new_name: str
169
170 def describe(self) -> str:
171 return f"{self.table}: rename index {self.old_name} -> {self.new_name}"
172
173 def apply(self) -> str:
174 sql = f"ALTER INDEX {quote_name(self.old_name)} RENAME TO {quote_name(self.new_name)}"
175 _execute_and_commit(sql)
176 return sql
177
178
179@dataclass
180class CreateIndexCorrection(Correction):
181 """Create a missing index using CONCURRENTLY (doesn't block writes)."""
182
183 pass_order = 1
184
185 table: str
186 index: Index
187 model: type[Model]
188
189 def describe(self) -> str:
190 return f"{self.table}: create index {self.index.name}"
191
192 def apply(self) -> str:
193 sql = self.index.to_sql(self.model)
194 _execute_autocommit(sql)
195 return sql
196
197
198@dataclass
199class AddConstraintCorrection(Correction):
200 """Add a missing constraint.
201
202 Check constraints use ADD CONSTRAINT ... NOT VALID + VALIDATE CONSTRAINT
203 in a single apply() — the add is catalog-only (brief lock) and the
204 validate uses SHARE UPDATE EXCLUSIVE which doesn't block writes, so
205 there's no benefit to deferring validation to a later run.
206 Unique constraints use CREATE UNIQUE INDEX CONCURRENTLY + USING INDEX
207 to avoid blocking writes.
208 """
209
210 pass_order = 2
211
212 table: str
213 constraint: BaseConstraint
214 model: type[Model]
215
216 def describe(self) -> str:
217 return f"{self.table}: add constraint {self.constraint.name}"
218
219 def apply(self) -> str:
220 if isinstance(self.constraint, UniqueConstraint):
221 return self._apply_unique(self.constraint)
222 return self._apply_other()
223
224 def _apply_unique(self, constraint: UniqueConstraint) -> str:
225 # Step 1: Create unique index concurrently (non-blocking)
226 create_idx = constraint.to_sql(self.model, concurrently=True)
227 _execute_autocommit(create_idx)
228
229 # Step 2: Attach as constraint — but only for variants PostgreSQL
230 # accepts. Partial indexes, expression indexes, and non-default
231 # operator class indexes cannot be attached as constraints; they
232 # remain as unique indexes (same enforcement, no pg_constraint row).
233 if constraint.index_only:
234 return create_idx
235
236 add_constraint = constraint.to_attach_sql(self.model)
237 try:
238 _execute_and_commit(add_constraint)
239 except Exception:
240 # Clean up the orphaned index if the constraint attachment fails
241 name = quote_name(constraint.name)
242 _execute_autocommit(f"DROP INDEX CONCURRENTLY IF EXISTS {name}")
243 raise
244
245 return f"{create_idx}; {add_constraint}"
246
247 def _apply_other(self) -> str:
248 if isinstance(self.constraint, CheckConstraint):
249 add_sql = self.constraint.to_sql(self.model, not_valid=True)
250 _execute_and_commit(add_sql)
251
252 validate_sql = (
253 f"ALTER TABLE {quote_name(self.table)}"
254 f" VALIDATE CONSTRAINT {quote_name(self.constraint.name)}"
255 )
256 _execute_and_commit(validate_sql, blocking=False)
257
258 return f"{add_sql}; {validate_sql}"
259
260 sql = self.constraint.to_sql(self.model)
261 _execute_and_commit(sql)
262 return sql
263
264
265@dataclass
266class AddForeignKeyCorrection(Correction):
267 """Add a missing FK constraint using NOT VALID, then validate immediately.
268
269 Step 1: ADD CONSTRAINT ... NOT VALID (SHARE ROW EXCLUSIVE, no scan)
270 Step 2: VALIDATE CONSTRAINT (SHARE UPDATE EXCLUSIVE, scans data)
271
272 Both steps run in a single apply() because the validate lock is weaker
273 than the add lock — there's no benefit to deferring validation.
274 """
275
276 pass_order = 2
277
278 table: str
279 constraint_name: str
280 column: str
281 target_table: str
282 target_column: str
283 on_delete_clause: str = "" # e.g. " ON DELETE CASCADE" or "" for NO ACTION
284
285 def describe(self) -> str:
286 return f"{self.table}: add FK {self.constraint_name} ({self.column} → {self.target_table}.{self.target_column})"
287
288 def apply(self) -> str:
289 add_sql = (
290 f"ALTER TABLE {quote_name(self.table)}"
291 f" ADD CONSTRAINT {quote_name(self.constraint_name)}"
292 f" FOREIGN KEY ({quote_name(self.column)})"
293 f" REFERENCES {quote_name(self.target_table)} ({quote_name(self.target_column)})"
294 f"{self.on_delete_clause}"
295 f" DEFERRABLE INITIALLY DEFERRED"
296 f" NOT VALID"
297 )
298 _execute_and_commit(add_sql)
299
300 validate_sql = (
301 f"ALTER TABLE {quote_name(self.table)}"
302 f" VALIDATE CONSTRAINT {quote_name(self.constraint_name)}"
303 )
304 _execute_and_commit(validate_sql, blocking=False)
305
306 return f"{add_sql}; {validate_sql}"
307
308
309@dataclass
310class ReplaceForeignKeyCorrection(Correction):
311 """Swap a FK's ON DELETE action, then validate.
312
313 Step 1: ALTER TABLE DROP CONSTRAINT + ADD CONSTRAINT ... NOT VALID
314 (ACCESS EXCLUSIVE on the referenced table's DROP, but the
315 statement is catalog-only — no table scan)
316 Step 2: VALIDATE CONSTRAINT (SHARE UPDATE EXCLUSIVE, scans data)
317
318 Between steps the constraint exists as NOT VALID, which still enforces
319 on new inserts/updates — so there is no window of unsafe writes. The
320 DROP and ADD share a single ALTER TABLE statement so the old
321 constraint is never absent in the catalog.
322
323 VALIDATE still scans the table, but under a weaker lock than a naive
324 DROP + ADD would take for the validation scan. Existing rows were
325 already valid under the previous constraint (on_delete action doesn't
326 affect what is validated — only what happens at delete time), so
327 the scan will pass unless the data was corrupted out-of-band.
328 """
329
330 pass_order = 2
331
332 table: str
333 constraint_name: str
334 column: str
335 target_table: str
336 target_column: str
337 on_delete_clause: str
338
339 def describe(self) -> str:
340 return f"{self.table}: update FK {self.constraint_name} on_delete"
341
342 def apply(self) -> str:
343 replace_sql = (
344 f"ALTER TABLE {quote_name(self.table)}"
345 f" DROP CONSTRAINT {quote_name(self.constraint_name)},"
346 f" ADD CONSTRAINT {quote_name(self.constraint_name)}"
347 f" FOREIGN KEY ({quote_name(self.column)})"
348 f" REFERENCES {quote_name(self.target_table)} ({quote_name(self.target_column)})"
349 f"{self.on_delete_clause}"
350 f" DEFERRABLE INITIALLY DEFERRED"
351 f" NOT VALID"
352 )
353 _execute_and_commit(replace_sql)
354
355 validate_sql = (
356 f"ALTER TABLE {quote_name(self.table)}"
357 f" VALIDATE CONSTRAINT {quote_name(self.constraint_name)}"
358 )
359 _execute_and_commit(validate_sql, blocking=False)
360
361 return f"{replace_sql}; {validate_sql}"
362
363
364@dataclass
365class SetNotNullCorrection(Correction):
366 """Enforce NOT NULL via CHECK NOT VALID → VALIDATE → SET NOT NULL.
367
368 A bare SET NOT NULL acquires ACCESS EXCLUSIVE and scans the whole table
369 while holding it. With a validated IS NOT NULL check constraint already
370 in place, Postgres (12+) skips the scan and the lock is brief.
371
372 Transaction boundaries are chosen to keep lock windows narrow:
373
374 1. ADD CHECK (col IS NOT NULL) NOT VALID — catalog-only, brief lock
375 2. VALIDATE CONSTRAINT — SHARE UPDATE EXCLUSIVE scan
376 3. SET NOT NULL + DROP temp check — atomic, instant catalog ops
377
378 Steps 3+4 share a single commit so no orphaned temp check can remain
379 after the column becomes NOT NULL. If earlier steps fail, the leftover
380 temp check is cleaned up at the start of the next run (analysis also
381 ignores framework-owned temp checks so they never block sync).
382 """
383
384 pass_order = 2
385
386 table: str
387 column: str
388
389 def describe(self) -> str:
390 return f"{self.table}: set NOT NULL on {self.column}"
391
392 def apply(self) -> str:
393 from .analysis import generate_notnull_check_name
394
395 t = quote_name(self.table)
396 c = quote_name(self.column)
397 check = quote_name(generate_notnull_check_name(self.table, self.column))
398
399 # Clean up any leftover temp constraint from a previous failed run
400 _execute_and_commit(f"ALTER TABLE {t} DROP CONSTRAINT IF EXISTS {check}")
401
402 # Step 1: Add NOT VALID check (no scan, brief lock)
403 add_sql = (
404 f"ALTER TABLE {t} ADD CONSTRAINT {check} CHECK ({c} IS NOT NULL) NOT VALID"
405 )
406 _execute_and_commit(add_sql)
407
408 # Step 2: Validate (SHARE UPDATE EXCLUSIVE — non-blocking scan)
409 validate_sql = f"ALTER TABLE {t} VALIDATE CONSTRAINT {check}"
410 _execute_and_commit(validate_sql, blocking=False)
411
412 # Step 3: SET NOT NULL + drop temp check in one commit.
413 # Both are instant catalog operations (SET NOT NULL skips the scan
414 # thanks to the validated check). Combining them ensures no orphaned
415 # temp check if SET NOT NULL succeeds.
416 set_sql = f"ALTER TABLE {t} ALTER COLUMN {c} SET NOT NULL"
417 drop_sql = f"ALTER TABLE {t} DROP CONSTRAINT {check}"
418 _execute_and_commit([set_sql, drop_sql])
419
420 return f"{add_sql}; {validate_sql}; {set_sql}; {drop_sql}"
421
422
423@dataclass
424class DropNotNullCorrection(Correction):
425 """Remove NOT NULL from a column (model now allows NULL).
426
427 DROP NOT NULL is a catalog-only change — no data scan, instant.
428 """
429
430 pass_order = 2
431
432 table: str
433 column: str
434
435 def describe(self) -> str:
436 return f"{self.table}: drop NOT NULL on {self.column}"
437
438 def apply(self) -> str:
439 sql = f"ALTER TABLE {quote_name(self.table)} ALTER COLUMN {quote_name(self.column)} DROP NOT NULL"
440 _execute_and_commit(sql)
441 return sql
442
443
444@dataclass
445class SetColumnDefaultCorrection(Correction):
446 """Set (or replace) a column's DEFAULT (catalog-only, instant)."""
447
448 pass_order = 2
449
450 table: str
451 column: str
452 default_sql: str
453
454 def describe(self) -> str:
455 return f"{self.table}: set DEFAULT {self.default_sql} on {self.column}"
456
457 def apply(self) -> str:
458 sql = (
459 f"ALTER TABLE {quote_name(self.table)}"
460 f" ALTER COLUMN {quote_name(self.column)}"
461 f" SET DEFAULT {self.default_sql}"
462 )
463 _execute_and_commit(sql)
464 return sql
465
466
467@dataclass
468class DropColumnDefaultCorrection(Correction):
469 """Drop a column's DEFAULT (catalog-only, instant)."""
470
471 pass_order = 2
472
473 table: str
474 column: str
475
476 def describe(self) -> str:
477 return f"{self.table}: drop DEFAULT on {self.column}"
478
479 def apply(self) -> str:
480 sql = (
481 f"ALTER TABLE {quote_name(self.table)}"
482 f" ALTER COLUMN {quote_name(self.column)}"
483 f" DROP DEFAULT"
484 )
485 _execute_and_commit(sql)
486 return sql
487
488
489@dataclass
490class RenameConstraintCorrection(Correction):
491 """Rename a constraint (catalog-only, instant).
492
493 For unique constraints, Postgres automatically renames the backing index.
494 """
495
496 pass_order = 2
497
498 table: str
499 old_name: str
500 new_name: str
501
502 def describe(self) -> str:
503 return f"{self.table}: rename constraint {self.old_name} -> {self.new_name}"
504
505 def apply(self) -> str:
506 sql = f"ALTER TABLE {quote_name(self.table)} RENAME CONSTRAINT {quote_name(self.old_name)} TO {quote_name(self.new_name)}"
507 _execute_and_commit(sql)
508 return sql
509
510
511@dataclass
512class ValidateConstraintCorrection(Correction):
513 """Validate a NOT VALID constraint (SHARE UPDATE EXCLUSIVE — doesn't block writes)."""
514
515 pass_order = 3
516
517 table: str
518 name: str
519
520 def describe(self) -> str:
521 return f"{self.table}: validate constraint {self.name}"
522
523 def apply(self) -> str:
524 sql = f"ALTER TABLE {quote_name(self.table)} VALIDATE CONSTRAINT {quote_name(self.name)}"
525 _execute_and_commit(sql, blocking=False)
526 return sql
527
528
529@dataclass
530class DropConstraintCorrection(Correction):
531 pass_order = 4
532
533 table: str
534 name: str
535
536 def describe(self) -> str:
537 return f"{self.table}: drop constraint {self.name}"
538
539 def apply(self) -> str:
540 sql = f"ALTER TABLE {quote_name(self.table)} DROP CONSTRAINT {quote_name(self.name)}"
541 _execute_and_commit(sql)
542 return sql
543
544
545@dataclass
546class DropIndexCorrection(Correction):
547 pass_order = 5
548
549 table: str
550 name: str
551
552 def describe(self) -> str:
553 return f"{self.table}: drop index {self.name}"
554
555 def apply(self) -> str:
556 sql = f"DROP INDEX CONCURRENTLY IF EXISTS {quote_name(self.name)}"
557 _execute_autocommit(sql)
558 return sql
559
560
561@dataclass
562class SetStorageParameterCorrection(Correction):
563 """Set a single `pg_class.reloptions` parameter (catalog-only, instant)."""
564
565 pass_order = 2
566
567 table: str
568 key: str
569 value: str
570
571 def describe(self) -> str:
572 return f"{self.table}: set storage parameter {self.key} = {self.value}"
573
574 def apply(self) -> str:
575 conn = get_connection()
576 quoted = psycopg.sql.quote(self.value, conn.connection)
577 sql = f"ALTER TABLE {quote_name(self.table)} SET ({self.key} = {quoted})"
578 _execute_and_commit(sql)
579 return sql
580
581
582@dataclass
583class ResetStorageParameterCorrection(Correction):
584 """Reset a single `pg_class.reloptions` parameter (catalog-only, instant)."""
585
586 pass_order = 2
587
588 table: str
589 key: str
590
591 def describe(self) -> str:
592 return f"{self.table}: reset storage parameter {self.key}"
593
594 def apply(self) -> str:
595 sql = f"ALTER TABLE {quote_name(self.table)} RESET ({self.key})"
596 _execute_and_commit(sql)
597 return sql