1from __future__ import annotations
2
3import _thread
4import warnings
5from collections import deque
6from collections.abc import Generator, Sequence
7from contextlib import contextmanager
8from typing import TYPE_CHECKING, Any, LiteralString, NamedTuple, cast
9
10import psycopg
11from psycopg import errors
12from psycopg import sql as psycopg_sql
13
14from plain.logs import get_framework_logger
15from plain.postgres import utils
16from plain.postgres.dialect import quote_name
17from plain.postgres.fields import GenericIPAddressField, TimeField, UUIDField
18from plain.postgres.schema import DatabaseSchemaEditor
19from plain.postgres.sources import ConnectionSource
20from plain.postgres.transaction import TransactionManagementError
21from plain.postgres.utils import CursorDebugWrapper as BaseCursorDebugWrapper
22from plain.postgres.utils import CursorWrapper, debug_transaction
23from plain.runtime import settings
24
25if TYPE_CHECKING:
26 from psycopg import Connection as PsycopgConnection
27
28 from plain.postgres.database_url import DatabaseConfig
29 from plain.postgres.fields import Field
30
31logger = get_framework_logger()
32
33
34def get_migratable_models() -> Generator[Any]:
35 """Return all models that should be included in migrations."""
36 from plain.packages import packages_registry
37 from plain.postgres import models_registry
38
39 return (
40 model
41 for package_config in packages_registry.get_package_configs()
42 for model in models_registry.get_models(
43 package_label=package_config.package_label
44 )
45 )
46
47
48class TableInfo(NamedTuple):
49 """Structure returned by DatabaseConnection.get_table_list()."""
50
51 name: str
52 type: str
53 comment: str | None
54
55
56class DatabaseConnection:
57 """
58 PostgreSQL database connection.
59
60 This is the only database backend supported by Plain.
61 """
62
63 queries_limit: int = 9000
64
65 ignored_tables: list[str] = []
66
67 def __init__(self, source: ConnectionSource):
68 # Lazy — acquired on first use via self._source.
69 self.connection: PsycopgConnection[Any] | None = None
70 self._source: ConnectionSource = source
71 # Query logging in debug mode or when explicitly enabled.
72 self.queries_log: deque[dict[str, Any]] = deque(maxlen=self.queries_limit)
73 self.force_debug_cursor: bool = False
74
75 # Transaction related attributes.
76 # Tracks if the connection is in autocommit mode. Per PEP 249, by
77 # default, it isn't.
78 self.autocommit: bool = False
79 # Tracks if the connection is in a transaction managed by 'atomic'.
80 self.in_atomic_block: bool = False
81 # Increment to generate unique savepoint ids.
82 self.savepoint_state: int = 0
83 # List of savepoints created by 'atomic'.
84 self.savepoint_ids: list[str | None] = []
85 # Stack of active 'atomic' blocks.
86 self.atomic_blocks: list[Any] = []
87 # Tracks if the transaction should be rolled back to the next
88 # available savepoint because of an exception in an inner block.
89 self.needs_rollback: bool = False
90 self.rollback_exc: Exception | None = None
91
92 # A list of no-argument functions to run when the transaction commits.
93 # Each entry is an (sids, func, robust) tuple, where sids is a set of
94 # the active savepoint IDs when this function was registered and robust
95 # specifies whether it's allowed for the function to fail.
96 self.run_on_commit: list[tuple[set[str | None], Any, bool]] = []
97
98 # Should we run the on-commit hooks the next time set_autocommit(True)
99 # is called?
100 self.run_commit_hooks_on_set_autocommit_on: bool = False
101
102 # A stack of wrappers to be invoked around execute()/executemany()
103 # calls. Each entry is a function taking five arguments: execute, sql,
104 # params, many, and context. It's the function's responsibility to
105 # call execute(sql, params, many, context).
106 self.execute_wrappers: list[Any] = []
107
108 def __repr__(self) -> str:
109 return f"<{self.__class__.__qualname__} vendor='postgresql'>"
110
111 def __del__(self) -> None:
112 # Safety net for wrappers GC'd without an explicit close() —
113 # e.g. inside a short-lived `asyncio.to_thread` context copy.
114 # Returns the pooled connection to the pool. Guards handle
115 # interpreter shutdown, when attrs may already be cleared.
116 conn = getattr(self, "connection", None)
117 if conn is None:
118 return
119 source = getattr(self, "_source", None)
120 if source is None:
121 return
122 try:
123 source.release(conn)
124 except Exception:
125 pass
126
127 @property
128 def settings_dict(self) -> DatabaseConfig:
129 """Config of the server this wrapper talks to. For pool-backed
130 wrappers this always reflects the live `POSTGRES_URL`."""
131 return self._source.config
132
133 @property
134 def queries_logged(self) -> bool:
135 return self.force_debug_cursor or settings.DEBUG
136
137 @property
138 def queries(self) -> list[dict[str, Any]]:
139 if len(self.queries_log) == self.queries_log.maxlen:
140 warnings.warn(
141 f"Limit for query logging exceeded, only the last {self.queries_log.maxlen} queries "
142 "will be returned."
143 )
144 return list(self.queries_log)
145
146 # ##### Connection and cursor methods #####
147
148 def _set_autocommit(self, autocommit: bool) -> None:
149 """Backend-specific implementation to enable or disable autocommit."""
150 assert self.connection is not None
151 self.connection.autocommit = autocommit
152
153 def check_constraints(self, table_names: list[str] | None = None) -> None:
154 """
155 Check constraints by setting them to immediate. Return them to deferred
156 afterward.
157 """
158 with self.cursor() as cursor:
159 cursor.execute("SET CONSTRAINTS ALL IMMEDIATE")
160 cursor.execute("SET CONSTRAINTS ALL DEFERRED")
161
162 def make_debug_cursor(self, cursor: psycopg.Cursor[Any]) -> CursorDebugWrapper:
163 return CursorDebugWrapper(cursor, self)
164
165 # ##### Connection lifecycle #####
166
167 def connect(self) -> None:
168 """Connect to the database. Assume that the connection is closed."""
169 self.connection = self._source.acquire()
170 self.set_autocommit(True)
171
172 def ensure_connection(self) -> None:
173 """Guarantee that a live connection to the database is established."""
174 if (
175 self.connection is not None
176 and self.connection.closed
177 and not self.in_atomic_block
178 ):
179 # The server closed this connection while we held it (restart,
180 # failover, idle timeout) — psycopg marks it closed once an
181 # operation fails on it. Mid-atomic, swapping it would silently
182 # run the rest of the block outside its transaction — leave the
183 # dead connection for Atomic.__exit__'s error recovery instead.
184 logger.warning("Discarding dead database connection")
185 self.close()
186 if self.connection is None:
187 self.connect()
188
189 # ##### PEP-249 connection method wrappers #####
190
191 def _prepare_cursor(self, cursor: psycopg.Cursor[Any]) -> utils.CursorWrapper:
192 """
193 Validate the connection is usable and perform database cursor wrapping.
194 """
195 if self.queries_logged:
196 wrapped_cursor = self.make_debug_cursor(cursor)
197 else:
198 wrapped_cursor = self.make_cursor(cursor)
199 return wrapped_cursor
200
201 def _cursor(self) -> utils.CursorWrapper:
202 self.ensure_connection()
203 assert self.connection is not None
204 return self._prepare_cursor(self.connection.cursor())
205
206 def _commit(self) -> None:
207 if self.connection is not None:
208 with debug_transaction(self, "COMMIT"):
209 return self.connection.commit()
210
211 def _rollback(self) -> None:
212 if self.connection is not None:
213 with debug_transaction(self, "ROLLBACK"):
214 return self.connection.rollback()
215
216 # ##### Generic wrappers for PEP-249 connection methods #####
217
218 def cursor(self) -> utils.CursorWrapper:
219 """Create a cursor, opening a connection if necessary."""
220 return self._cursor()
221
222 def commit(self) -> None:
223 """Commit a transaction and reset the dirty flag."""
224 self.validate_no_atomic_block()
225 self._commit()
226 self.run_commit_hooks_on_set_autocommit_on = True
227
228 def rollback(self) -> None:
229 """Roll back a transaction and reset the dirty flag."""
230 self.validate_no_atomic_block()
231 self._rollback()
232 self.needs_rollback = False
233 self.run_on_commit = []
234
235 def close(self) -> None:
236 """Close the connection to the database."""
237 # Closing mid-atomic would reopen a fresh autocommit connection on
238 # the next cursor() and silently run the rest of the block outside
239 # its transaction. Callers that drop a connection during error
240 # recovery (see Atomic.__exit__) unwind the atomic state first.
241 self.validate_no_atomic_block()
242
243 self.run_on_commit = []
244 if self.connection is None:
245 return
246 try:
247 self._source.release(self.connection)
248 finally:
249 # Null the reference so __del__ (and ensure_connection) can't
250 # touch an already-released psycopg connection.
251 self.connection = None
252
253 # ##### Savepoint management #####
254
255 def _savepoint(self, sid: str) -> None:
256 with self.cursor() as cursor:
257 cursor.execute(f"SAVEPOINT {quote_name(sid)}")
258
259 def _savepoint_rollback(self, sid: str) -> None:
260 with self.cursor() as cursor:
261 cursor.execute(f"ROLLBACK TO SAVEPOINT {quote_name(sid)}")
262
263 def _savepoint_commit(self, sid: str) -> None:
264 with self.cursor() as cursor:
265 cursor.execute(f"RELEASE SAVEPOINT {quote_name(sid)}")
266
267 # ##### Generic savepoint management methods #####
268
269 def savepoint(self) -> str | None:
270 """
271 Create a savepoint inside the current transaction. Return an
272 identifier for the savepoint that will be used for the subsequent
273 rollback or commit. Return None if in autocommit mode (no transaction).
274 """
275 if self.get_autocommit():
276 return None
277
278 thread_ident = _thread.get_ident()
279 tid = str(thread_ident).replace("-", "")
280
281 self.savepoint_state += 1
282 sid = "s%s_x%d" % (tid, self.savepoint_state) # noqa: UP031
283
284 self._savepoint(sid)
285
286 return sid
287
288 def savepoint_rollback(self, sid: str) -> None:
289 """
290 Roll back to a savepoint. Do nothing if in autocommit mode.
291 """
292 if self.get_autocommit():
293 return
294
295 self._savepoint_rollback(sid)
296
297 # Remove any callbacks registered while this savepoint was active.
298 self.run_on_commit = [
299 (sids, func, robust)
300 for (sids, func, robust) in self.run_on_commit
301 if sid not in sids
302 ]
303
304 def savepoint_commit(self, sid: str) -> None:
305 """
306 Release a savepoint. Do nothing if in autocommit mode.
307 """
308 if self.get_autocommit():
309 return
310
311 self._savepoint_commit(sid)
312
313 def clean_savepoints(self) -> None:
314 """
315 Reset the counter used to generate unique savepoint ids in this thread.
316 """
317 self.savepoint_state = 0
318
319 # ##### Generic transaction management methods #####
320
321 def get_autocommit(self) -> bool:
322 """Get the autocommit state."""
323 self.ensure_connection()
324 return self.autocommit
325
326 def set_autocommit(self, autocommit: bool) -> None:
327 """
328 Enable or disable autocommit.
329
330 Used internally by atomic() to manage transactions. Don't call this
331 directly — use atomic() instead.
332 """
333 self.validate_no_atomic_block()
334 self.ensure_connection()
335
336 if autocommit:
337 self._set_autocommit(autocommit)
338 else:
339 with debug_transaction(self, "BEGIN"):
340 self._set_autocommit(autocommit)
341 self.autocommit = autocommit
342
343 if autocommit and self.run_commit_hooks_on_set_autocommit_on:
344 self.run_and_clear_commit_hooks()
345 self.run_commit_hooks_on_set_autocommit_on = False
346
347 def get_rollback(self) -> bool:
348 """Get the "needs rollback" flag -- for *advanced use* only."""
349 if not self.in_atomic_block:
350 raise TransactionManagementError(
351 "The rollback flag doesn't work outside of an 'atomic' block."
352 )
353 return self.needs_rollback
354
355 def set_rollback(self, rollback: bool) -> None:
356 """
357 Set or unset the "needs rollback" flag -- for *advanced use* only.
358 """
359 if not self.in_atomic_block:
360 raise TransactionManagementError(
361 "The rollback flag doesn't work outside of an 'atomic' block."
362 )
363 self.needs_rollback = rollback
364
365 def validate_no_atomic_block(self) -> None:
366 """Raise an error if an atomic block is active."""
367 if self.in_atomic_block:
368 raise TransactionManagementError(
369 "This is forbidden when an 'atomic' block is active."
370 )
371
372 def validate_no_broken_transaction(self) -> None:
373 if self.needs_rollback:
374 raise TransactionManagementError(
375 "An error occurred in the current transaction. You can't "
376 "execute queries until the end of the 'atomic' block."
377 ) from self.rollback_exc
378
379 # ##### Miscellaneous #####
380
381 def make_cursor(self, cursor: psycopg.Cursor[Any]) -> utils.CursorWrapper:
382 """Create a cursor without debug logging."""
383 return utils.CursorWrapper(cursor, self)
384
385 def schema_editor(self, *args: Any, **kwargs: Any) -> DatabaseSchemaEditor:
386 """Return a new instance of the schema editor."""
387 return DatabaseSchemaEditor(self, *args, **kwargs)
388
389 def on_commit(self, func: Any, robust: bool = False) -> None:
390 if not callable(func):
391 raise TypeError("on_commit()'s callback must be a callable.")
392 if self.in_atomic_block:
393 # Transaction in progress; save for execution on commit.
394 self.run_on_commit.append((set(self.savepoint_ids), func, robust))
395 else:
396 # No transaction in progress; execute immediately.
397 if robust:
398 try:
399 func()
400 except Exception as e:
401 logger.error(
402 "Error calling on_commit() handler",
403 exc_info=True,
404 extra={"handler": func.__qualname__, "error": str(e)},
405 )
406 else:
407 func()
408
409 def run_and_clear_commit_hooks(self) -> None:
410 self.validate_no_atomic_block()
411 current_run_on_commit = self.run_on_commit
412 self.run_on_commit = []
413 while current_run_on_commit:
414 _, func, robust = current_run_on_commit.pop(0)
415 if robust:
416 try:
417 func()
418 except Exception as e:
419 logger.error(
420 "Error calling on_commit() handler during transaction",
421 exc_info=True,
422 extra={"handler": func.__qualname__, "error": str(e)},
423 )
424 else:
425 func()
426
427 @contextmanager
428 def execute_wrapper(self, wrapper: Any) -> Generator[None]:
429 """
430 Return a context manager under which the wrapper is applied to suitable
431 database query executions.
432 """
433 self.execute_wrappers.append(wrapper)
434 try:
435 yield
436 finally:
437 self.execute_wrappers.pop()
438
439 # ##### SQL generation methods that require connection state #####
440
441 def compose_sql(self, query: str, params: Any) -> str:
442 """
443 Compose a SQL query with parameters using psycopg's mogrify.
444
445 This requires an active connection because it uses the connection's
446 cursor to properly format parameters.
447 """
448 assert self.connection is not None
449 return psycopg.ClientCursor(self.connection).mogrify(
450 psycopg_sql.SQL(cast(LiteralString, query)), params
451 )
452
453 def last_executed_query(
454 self,
455 cursor: utils.CursorWrapper,
456 sql: str,
457 params: Any,
458 ) -> str | None:
459 """
460 Return a string of the query last executed by the given cursor, with
461 placeholders replaced with actual values.
462 """
463 try:
464 return self.compose_sql(sql, params)
465 except errors.DataError:
466 return None
467
468 def unification_cast_sql(self, output_field: Field) -> str:
469 """
470 Given a field instance, return the SQL that casts the result of a union
471 to that type. The resulting string should contain a '%s' placeholder
472 for the expression being cast.
473 """
474 if isinstance(output_field, GenericIPAddressField | TimeField | UUIDField):
475 # PostgreSQL will resolve a union as type 'text' if input types are
476 # 'unknown'.
477 # https://www.postgresql.org/docs/current/typeconv-union-case.html
478 # These fields cannot be implicitly cast back in the default
479 # PostgreSQL configuration so we need to explicitly cast them.
480 # We must also remove components of the type within brackets:
481 # varchar(255) -> varchar.
482 db_type = output_field.db_type()
483 if db_type:
484 return "CAST(%s AS {})".format(db_type.split("(")[0])
485 return "%s"
486
487 # ##### Introspection methods #####
488
489 def table_names(
490 self, cursor: CursorWrapper | None = None, include_views: bool = False
491 ) -> list[str]:
492 """
493 Return a list of names of all tables that exist in the database.
494 Sort the returned table list by Python's default sorting. Do NOT use
495 the database's ORDER BY here to avoid subtle differences in sorting
496 order between databases.
497 """
498
499 def get_names(cursor: CursorWrapper) -> list[str]:
500 return sorted(
501 ti.name
502 for ti in self.get_table_list(cursor)
503 if include_views or ti.type == "t"
504 )
505
506 if cursor is None:
507 with self.cursor() as cursor:
508 return get_names(cursor)
509 return get_names(cursor)
510
511 def get_table_list(self, cursor: CursorWrapper) -> Sequence[TableInfo]:
512 """
513 Return an unsorted list of TableInfo named tuples of all tables and
514 views that exist in the database.
515 """
516 cursor.execute(
517 """
518 SELECT
519 c.relname,
520 CASE
521 WHEN c.relispartition THEN 'p'
522 WHEN c.relkind IN ('m', 'v') THEN 'v'
523 ELSE 't'
524 END,
525 obj_description(c.oid, 'pg_class')
526 FROM pg_catalog.pg_class c
527 LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
528 WHERE c.relkind IN ('f', 'm', 'p', 'r', 'v')
529 AND n.nspname NOT IN ('pg_catalog', 'pg_toast')
530 AND pg_catalog.pg_table_is_visible(c.oid)
531 """
532 )
533 return [
534 TableInfo(*row)
535 for row in cursor.fetchall()
536 if row[0] not in self.ignored_tables
537 ]
538
539 def plain_table_names(
540 self, only_existing: bool = False, include_views: bool = True
541 ) -> list[str]:
542 """
543 Return a list of all table names that have associated Plain models and
544 are in INSTALLED_PACKAGES.
545
546 If only_existing is True, include only the tables in the database.
547 """
548 tables = set()
549 for model in get_migratable_models():
550 tables.add(model.model_options.db_table)
551 tables.update(
552 f.m2m_db_table() for f in model._model_meta.local_many_to_many
553 )
554 tables = list(tables)
555 if only_existing:
556 existing_tables = set(self.table_names(include_views=include_views))
557 tables = [t for t in tables if t in existing_tables]
558 return tables
559
560 def get_sequences(
561 self, cursor: CursorWrapper, table_name: str, table_fields: tuple[Any, ...] = ()
562 ) -> list[dict[str, Any]]:
563 """
564 Return a list of introspected sequences for table_name. Each sequence
565 is a dict: {'table': <table_name>, 'column': <column_name>, 'name': <sequence_name>}.
566 """
567 cursor.execute(
568 """
569 SELECT
570 s.relname AS sequence_name,
571 a.attname AS colname
572 FROM
573 pg_class s
574 JOIN pg_depend d ON d.objid = s.oid
575 AND d.classid = 'pg_class'::regclass
576 AND d.refclassid = 'pg_class'::regclass
577 JOIN pg_attribute a ON d.refobjid = a.attrelid
578 AND d.refobjsubid = a.attnum
579 JOIN pg_class tbl ON tbl.oid = d.refobjid
580 AND tbl.relname = %s
581 AND pg_catalog.pg_table_is_visible(tbl.oid)
582 WHERE
583 s.relkind = 'S';
584 """,
585 [table_name],
586 )
587 return [
588 {"name": row[0], "table": table_name, "column": row[1]}
589 for row in cursor.fetchall()
590 ]
591
592 def get_constraints(
593 self, cursor: CursorWrapper, table_name: str
594 ) -> dict[str, dict[str, Any]]:
595 """
596 Retrieve any constraints or keys (unique, pk, fk, check, index) across
597 one or more columns. Also retrieve the definition of expression-based
598 indexes.
599 """
600 constraints: dict[str, dict[str, Any]] = {}
601 # Loop over the key table, collecting things as constraints. The column
602 # array must return column names in the same order in which they were
603 # created.
604 cursor.execute(
605 """
606 SELECT
607 c.conname,
608 array(
609 SELECT attname
610 FROM unnest(c.conkey) WITH ORDINALITY cols(colid, arridx)
611 JOIN pg_attribute AS ca ON cols.colid = ca.attnum
612 WHERE ca.attrelid = c.conrelid
613 ORDER BY cols.arridx
614 ),
615 c.contype,
616 (SELECT fkc.relname || '.' || fka.attname
617 FROM pg_attribute AS fka
618 JOIN pg_class AS fkc ON fka.attrelid = fkc.oid
619 WHERE fka.attrelid = c.confrelid AND fka.attnum = c.confkey[1]),
620 c.convalidated,
621 pg_get_constraintdef(c.oid),
622 c.confdeltype
623 FROM pg_constraint AS c
624 JOIN pg_class AS cl ON c.conrelid = cl.oid
625 WHERE cl.relname = %s AND pg_catalog.pg_table_is_visible(cl.oid)
626 """,
627 [table_name],
628 )
629 for (
630 constraint,
631 columns,
632 kind,
633 used_cols,
634 validated,
635 constraintdef,
636 confdeltype,
637 ) in cursor.fetchall():
638 constraints[constraint] = {
639 "columns": columns,
640 "foreign_key": tuple(used_cols.split(".", 1)) if kind == "f" else None,
641 "contype": kind,
642 "index": False,
643 "definition": constraintdef,
644 "validated": validated,
645 "on_delete_action": confdeltype if kind == "f" else None,
646 }
647 # Now get indexes. Sort order, opclasses, INCLUDE, and predicates all
648 # ride along inside `pg_get_indexdef` and are compared via the
649 # normalized-tail round-trip in convergence — no need to introspect
650 # them here as separate columns.
651 cursor.execute(
652 """
653 SELECT
654 indexname,
655 array_agg(attname ORDER BY arridx),
656 indisunique,
657 amname,
658 exprdef,
659 indisvalid
660 FROM (
661 SELECT
662 c2.relname as indexname, idx.*, attr.attname, am.amname,
663 pg_get_indexdef(idx.indexrelid) AS exprdef
664 FROM (
665 SELECT *
666 FROM
667 pg_index i,
668 unnest(i.indkey)
669 WITH ORDINALITY koi(key, arridx)
670 ) idx
671 LEFT JOIN pg_class c ON idx.indrelid = c.oid
672 LEFT JOIN pg_class c2 ON idx.indexrelid = c2.oid
673 LEFT JOIN pg_am am ON c2.relam = am.oid
674 LEFT JOIN
675 pg_attribute attr ON attr.attrelid = c.oid AND attr.attnum = idx.key
676 WHERE c.relname = %s AND pg_catalog.pg_table_is_visible(c.oid)
677 ) s2
678 GROUP BY
679 indexname, indisunique, amname, exprdef, indisvalid;
680 """,
681 [table_name],
682 )
683 for (
684 index,
685 columns,
686 unique,
687 type_,
688 definition,
689 valid,
690 ) in cursor.fetchall():
691 if index not in constraints:
692 constraints[index] = {
693 "columns": columns if columns != [None] else [],
694 "unique": unique,
695 "index": True,
696 "type": type_,
697 "definition": definition,
698 "valid": valid,
699 }
700 return constraints
701
702
703class CursorDebugWrapper(BaseCursorDebugWrapper):
704 def copy(self, statement: Any) -> Any:
705 with self.debug_sql(statement):
706 return self.cursor.copy(statement)