1from __future__ import annotations
2
3import functools
4import time
5from collections.abc import Generator, Iterator, Mapping, Sequence
6from contextlib import contextmanager
7from hashlib import md5
8from types import TracebackType
9from typing import TYPE_CHECKING, Any, Self
10
11import psycopg
12from plain.logs import get_framework_logger
13from plain.postgres.otel import db_span
14from plain.utils.dateparse import parse_time
15
16if TYPE_CHECKING:
17 from plain.postgres.connection import DatabaseConnection
18
19logger = get_framework_logger()
20
21
22def make_model_tuple(model: Any) -> tuple[str, str]:
23 """
24 Take a model or a string of the form "package_label.ModelName" and return a
25 corresponding ("package_label", "modelname") tuple. If a tuple is passed in,
26 assume it's a valid model tuple already and return it unchanged.
27 """
28 try:
29 if isinstance(model, tuple):
30 model_tuple = model
31 elif isinstance(model, str):
32 package_label, model_name = model.split(".")
33 model_tuple = package_label, model_name.lower()
34 else:
35 model_tuple = (
36 model.model_options.package_label,
37 model.model_options.model_name,
38 )
39 assert len(model_tuple) == 2
40 return model_tuple
41 except (ValueError, AssertionError):
42 raise ValueError(
43 f"Invalid model reference '{model}'. String model references "
44 "must be of the form 'package_label.ModelName'."
45 )
46
47
48def resolve_callables(
49 mapping: dict[str, Any],
50) -> Generator[tuple[str, Any]]:
51 """
52 Generate key/value pairs for the given mapping where the values are
53 evaluated if they're callable.
54 """
55 for k, v in mapping.items():
56 yield k, v() if callable(v) else v
57
58
59class CursorWrapper:
60 def __init__(self, cursor: Any, db: DatabaseConnection) -> None:
61 self.cursor = cursor
62 self.db = db
63
64 def __getattr__(self, attr: str) -> Any:
65 return getattr(self.cursor, attr)
66
67 def __iter__(self) -> Iterator[tuple[Any, ...]]:
68 yield from self.cursor
69
70 def fetchone(self) -> tuple[Any, ...] | None:
71 return self.cursor.fetchone()
72
73 def fetchmany(self, size: int | None = None) -> list[tuple[Any, ...]]:
74 if size is None:
75 return self.cursor.fetchmany()
76 return self.cursor.fetchmany(size)
77
78 def fetchall(self) -> list[tuple[Any, ...]]:
79 return self.cursor.fetchall()
80
81 def __enter__(self) -> Self:
82 return self
83
84 def __exit__(
85 self,
86 type: type[BaseException] | None,
87 value: BaseException | None,
88 traceback: TracebackType | None,
89 ) -> None:
90 # Close instead of passing through to avoid backend-specific behavior
91 # (#17671). Catch errors liberally because errors in cleanup code
92 # aren't useful.
93 try:
94 self.close()
95 except psycopg.Error:
96 pass
97
98 def stream(
99 self, sql: str, params: Sequence[Any] | None = None
100 ) -> Generator[tuple[Any, ...]]:
101 self.db.validate_no_broken_transaction()
102 # psycopg's server-side cursor leaves rowcount at -1, so count rows as
103 # they're yielded and feed db_span via the closure.
104 count = 0
105 with db_span(self.db, sql, params=params, row_count_provider=lambda: count):
106 try:
107 iterator = (
108 self.cursor.stream(sql)
109 if params is None
110 else self.cursor.stream(sql, params)
111 )
112 for row in iterator:
113 count += 1
114 yield row
115 finally:
116 try:
117 self.close()
118 except psycopg.Error:
119 pass
120
121 # execute() and executemany() cannot be implemented in __getattr__ because
122 # the code must run when the method is invoked, not just when it is accessed.
123
124 def execute(
125 self, sql: str, params: Sequence[Any] | Mapping[str, Any] | None = None
126 ) -> Self:
127 return self._execute_with_wrappers(
128 sql, params, many=False, executor=self._execute
129 )
130
131 def executemany(self, sql: str, param_list: Sequence[Sequence[Any]]) -> Self:
132 return self._execute_with_wrappers(
133 sql, param_list, many=True, executor=self._executemany
134 )
135
136 def _execute_with_wrappers(
137 self, sql: str, params: Any, many: bool, executor: Any
138 ) -> Self:
139 context: dict[str, Any] = {"connection": self.db, "cursor": self}
140 for wrapper in reversed(self.db.execute_wrappers):
141 executor = functools.partial(wrapper, executor)
142 executor(sql, params, many, context)
143 return self
144
145 def _execute(self, sql: str, params: Any, *ignored_wrapper_args: Any) -> None:
146 with db_span(
147 self.db, sql, params=params, row_count_provider=lambda: self.cursor.rowcount
148 ):
149 self.db.validate_no_broken_transaction()
150 if params is None:
151 self.cursor.execute(sql)
152 else:
153 self.cursor.execute(sql, params)
154
155 def _executemany(
156 self, sql: str, param_list: Any, *ignored_wrapper_args: Any
157 ) -> None:
158 with db_span(
159 self.db,
160 sql,
161 many=True,
162 params=param_list,
163 row_count_provider=lambda: self.cursor.rowcount,
164 ):
165 self.db.validate_no_broken_transaction()
166 self.cursor.executemany(sql, param_list)
167
168
169class CursorDebugWrapper(CursorWrapper):
170 def stream(
171 self, sql: str, params: Sequence[Any] | None = None
172 ) -> Generator[tuple[Any, ...]]:
173 with self.debug_sql(sql, params, use_last_executed_query=True):
174 yield from super().stream(sql, params)
175
176 def execute(
177 self, sql: str, params: Sequence[Any] | Mapping[str, Any] | None = None
178 ) -> Self:
179 with self.debug_sql(sql, params, use_last_executed_query=True):
180 super().execute(sql, params)
181 return self
182
183 def executemany(self, sql: str, param_list: Sequence[Sequence[Any]]) -> Self:
184 with self.debug_sql(sql, param_list, many=True):
185 super().executemany(sql, param_list)
186 return self
187
188 @contextmanager
189 def debug_sql(
190 self,
191 sql: str | None = None,
192 params: Any = None,
193 use_last_executed_query: bool = False,
194 many: bool = False,
195 ) -> Generator[None]:
196 start = time.monotonic()
197 try:
198 yield
199 finally:
200 stop = time.monotonic()
201 duration = stop - start
202 if use_last_executed_query:
203 sql = self.db.last_executed_query(self.cursor, sql, params) # ty: ignore[invalid-argument-type]
204 try:
205 times = len(params) if many else ""
206 except TypeError:
207 # params could be an iterator.
208 times = "?"
209 self.db.queries_log.append(
210 {
211 "sql": f"{times} times: {sql}" if many else sql,
212 "time": f"{duration:.3f}",
213 }
214 )
215 logger.debug(
216 "Query executed",
217 extra={
218 "duration": round(duration, 3),
219 "sql": sql,
220 "params": params,
221 },
222 )
223
224
225@contextmanager
226def debug_transaction(connection: DatabaseConnection, sql: str) -> Generator[None]:
227 start = time.monotonic()
228 try:
229 yield
230 finally:
231 if connection.queries_logged:
232 stop = time.monotonic()
233 duration = stop - start
234 connection.queries_log.append(
235 {
236 "sql": f"{sql}",
237 "time": f"{duration:.3f}",
238 }
239 )
240 logger.debug(
241 "Transaction command",
242 extra={
243 "duration": round(duration, 3),
244 "sql": sql,
245 },
246 )
247
248
249def split_tzname_delta(tzname: str) -> tuple[str, str | None, str | None]:
250 """
251 Split a time zone name into a 3-tuple of (name, sign, offset).
252 """
253 for sign in ["+", "-"]:
254 if sign in tzname:
255 name, offset = tzname.rsplit(sign, 1)
256 if offset and parse_time(offset):
257 return name, sign, offset
258 return tzname, None, None
259
260
261###############################################
262# Converters from Python to database (string) #
263###############################################
264
265
266def split_identifier(identifier: str) -> tuple[str, str]:
267 """
268 Split an SQL identifier into a two element tuple of (namespace, name).
269
270 The identifier could be a table, column, or sequence name might be prefixed
271 by a namespace.
272 """
273 try:
274 namespace, name = identifier.split('"."')
275 except ValueError:
276 namespace, name = "", identifier
277 return namespace.strip('"'), name.strip('"')
278
279
280def truncate_name(identifier: str, length: int | None = None, hash_len: int = 4) -> str:
281 """
282 Shorten an SQL identifier to a repeatable mangled version with the given
283 length.
284
285 If a quote stripped name contains a namespace, e.g. USERNAME"."TABLE,
286 truncate the table portion only.
287 """
288 namespace, name = split_identifier(identifier)
289
290 if length is None or len(name) <= length:
291 return identifier
292
293 digest = names_digest(name, length=hash_len)
294 return "{}{}{}".format(
295 f'{namespace}"."' if namespace else "",
296 name[: length - hash_len],
297 digest,
298 )
299
300
301def names_digest(*args: str, length: int) -> str:
302 """
303 Generate a 32-bit digest of a set of arguments that can be used to shorten
304 identifying names.
305 """
306 h = md5(usedforsecurity=False)
307 for arg in args:
308 h.update(arg.encode())
309 return h.hexdigest()[:length]
310
311
312def generate_fk_constraint_name(
313 table: str, column: str, target_table: str, target_column: str
314) -> str:
315 """The deterministic name of a foreign key constraint."""
316 _, target_table_name = split_identifier(target_table)
317 suffix = f"_fk_{target_table_name}_{target_column}"
318 return generate_identifier_name(table, [column], suffix)
319
320
321def generate_identifier_name(
322 table_name: str, column_names: list[str], suffix: str = ""
323) -> str:
324 """Generate a deterministic name for an index or constraint.
325
326 The name is composed of the table name, column names, a hash digest,
327 and an optional suffix. Long names are truncated proportionally.
328 """
329 from .dialect import MAX_NAME_LENGTH
330
331 _, table_name = split_identifier(table_name)
332 hash_suffix_part = f"{names_digest(table_name, *column_names, length=8)}{suffix}"
333 max_length = MAX_NAME_LENGTH
334 name = f"{table_name}_{'_'.join(column_names)}_{hash_suffix_part}"
335 if len(name) <= max_length:
336 return name
337 if len(hash_suffix_part) > max_length / 3:
338 hash_suffix_part = hash_suffix_part[: max_length // 3]
339 other_length = (max_length - len(hash_suffix_part)) // 2 - 1
340 name = f"{table_name[:other_length]}_{'_'.join(column_names)[:other_length]}_{hash_suffix_part}"
341 if name[0] == "_" or name[0].isdigit():
342 name = f"D{name[:-1]}"
343 return name
344
345
346def strip_quotes(table_name: str) -> str:
347 """
348 Strip quotes off of quoted table names to make them safe for use in index
349 names, sequence names, etc.
350 """
351 has_quotes = table_name.startswith('"') and table_name.endswith('"')
352 return table_name[1:-1] if has_quotes else table_name