codekingpro/portable-devtools
115k
1# dialects/postgresql/ext.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: ignore-errors
8from __future__ import annotations
9
10from typing import Any
11from typing import Iterable
12from typing import List
13from typing import Optional
14from typing import overload
15from typing import Tuple
16from typing import TYPE_CHECKING
17from typing import TypeVar
18
19from . import types
20from .array import ARRAY
21from ...sql import coercions
22from ...sql import elements
23from ...sql import expression
24from ...sql import functions
25from ...sql import roles
26from ...sql import schema
27from ...sql.schema import ColumnCollectionConstraint
28from ...sql.sqltypes import TEXT
29from ...sql.visitors import InternalTraversal
30
31if TYPE_CHECKING:
32 from ...sql._typing import _ColumnExpressionArgument
33 from ...sql._typing import _DDLColumnArgument
34 from ...sql.elements import ClauseElement
35 from ...sql.elements import ColumnElement
36 from ...sql.operators import OperatorType
37 from ...sql.selectable import FromClause
38 from ...sql.visitors import _CloneCallableType
39 from ...sql.visitors import _TraverseInternalsType
40
41_T = TypeVar("_T", bound=Any)
42
43
44class aggregate_order_by(expression.ColumnElement[_T]):
45 """Represent a PostgreSQL aggregate order by expression.
46
47 E.g.::
48
49 from sqlalchemy.dialects.postgresql import aggregate_order_by
50
51 expr = func.array_agg(aggregate_order_by(table.c.a, table.c.b.desc()))
52 stmt = select(expr)
53
54 would represent the expression:
55
56 .. sourcecode:: sql
57
58 SELECT array_agg(a ORDER BY b DESC) FROM table;
59
60 Similarly::
61
62 expr = func.string_agg(
63 table.c.a, aggregate_order_by(literal_column("','"), table.c.a)
64 )
65 stmt = select(expr)
66
67 Would represent:
68
69 .. sourcecode:: sql
70
71 SELECT string_agg(a, ',' ORDER BY a) FROM table;
72
73 .. versionchanged:: 1.2.13 - the ORDER BY argument may be multiple terms
74
75 .. seealso::
76
77 :class:`_functions.array_agg`
78
79 """
80
81 __visit_name__ = "aggregate_order_by"
82
83 stringify_dialect = "postgresql"
84 _traverse_internals: _TraverseInternalsType = [
85 ("target", InternalTraversal.dp_clauseelement),
86 ("type", InternalTraversal.dp_type),
87 ("order_by", InternalTraversal.dp_clauseelement),
88 ]
89
90 @overload
91 def __init__(
92 self,
93 target: ColumnElement[_T],
94 *order_by: _ColumnExpressionArgument[Any],
95 ): ...
96
97 @overload
98 def __init__(
99 self,
100 target: _ColumnExpressionArgument[_T],
101 *order_by: _ColumnExpressionArgument[Any],
102 ): ...
103
104 def __init__(
105 self,
106 target: _ColumnExpressionArgument[_T],
107 *order_by: _ColumnExpressionArgument[Any],
108 ):
109 self.target: ClauseElement = coercions.expect(
110 roles.ExpressionElementRole, target
111 )
112 self.type = self.target.type
113
114 _lob = len(order_by)
115 self.order_by: ClauseElement
116 if _lob == 0:
117 raise TypeError("at least one ORDER BY element is required")
118 elif _lob == 1:
119 self.order_by = coercions.expect(
120 roles.ExpressionElementRole, order_by[0]
121 )
122 else:
123 self.order_by = elements.ClauseList(
124 *order_by, _literal_as_text_role=roles.ExpressionElementRole
125 )
126
127 def self_group(
128 self, against: Optional[OperatorType] = None
129 ) -> ClauseElement:
130 return self
131
132 def get_children(self, **kwargs: Any) -> Iterable[ClauseElement]:
133 return self.target, self.order_by
134
135 def _copy_internals(
136 self, clone: _CloneCallableType = elements._clone, **kw: Any
137 ) -> None:
138 self.target = clone(self.target, **kw)
139 self.order_by = clone(self.order_by, **kw)
140
141 @property
142 def _from_objects(self) -> List[FromClause]:
143 return self.target._from_objects + self.order_by._from_objects
144
145
146class ExcludeConstraint(ColumnCollectionConstraint):
147 """A table-level EXCLUDE constraint.
148
149 Defines an EXCLUDE constraint as described in the `PostgreSQL
150 documentation`__.
151
152 __ https://www.postgresql.org/docs/current/static/sql-createtable.html#SQL-CREATETABLE-EXCLUDE
153
154 """ # noqa
155
156 __visit_name__ = "exclude_constraint"
157
158 where = None
159 inherit_cache = False
160
161 create_drop_stringify_dialect = "postgresql"
162
163 @elements._document_text_coercion(
164 "where",
165 ":class:`.ExcludeConstraint`",
166 ":paramref:`.ExcludeConstraint.where`",
167 )
168 def __init__(
169 self, *elements: Tuple[_DDLColumnArgument, str], **kw: Any
170 ) -> None:
171 r"""
172 Create an :class:`.ExcludeConstraint` object.
173
174 E.g.::
175
176 const = ExcludeConstraint(
177 (Column("period"), "&&"),
178 (Column("group"), "="),
179 where=(Column("group") != "some group"),
180 ops={"group": "my_operator_class"},
181 )
182
183 The constraint is normally embedded into the :class:`_schema.Table`
184 construct
185 directly, or added later using :meth:`.append_constraint`::
186
187 some_table = Table(
188 "some_table",
189 metadata,
190 Column("id", Integer, primary_key=True),
191 Column("period", TSRANGE()),
192 Column("group", String),
193 )
194
195 some_table.append_constraint(
196 ExcludeConstraint(
197 (some_table.c.period, "&&"),
198 (some_table.c.group, "="),
199 where=some_table.c.group != "some group",
200 name="some_table_excl_const",
201 ops={"group": "my_operator_class"},
202 )
203 )
204
205 The exclude constraint defined in this example requires the
206 ``btree_gist`` extension, that can be created using the
207 command ``CREATE EXTENSION btree_gist;``.
208
209 :param \*elements:
210
211 A sequence of two tuples of the form ``(column, operator)`` where
212 "column" is either a :class:`_schema.Column` object, or a SQL
213 expression element (e.g. ``func.int8range(table.from, table.to)``)
214 or the name of a column as string, and "operator" is a string
215 containing the operator to use (e.g. `"&&"` or `"="`).
216
217 In order to specify a column name when a :class:`_schema.Column`
218 object is not available, while ensuring
219 that any necessary quoting rules take effect, an ad-hoc
220 :class:`_schema.Column` or :func:`_expression.column`
221 object should be used.
222 The ``column`` may also be a string SQL expression when
223 passed as :func:`_expression.literal_column` or
224 :func:`_expression.text`
225
226 :param name:
227 Optional, the in-database name of this constraint.
228
229 :param deferrable:
230 Optional bool. If set, emit DEFERRABLE or NOT DEFERRABLE when
231 issuing DDL for this constraint.
232
233 :param initially:
234 Optional string. If set, emit INITIALLY <value> when issuing DDL
235 for this constraint.
236
237 :param using:
238 Optional string. If set, emit USING <index_method> when issuing DDL
239 for this constraint. Defaults to 'gist'.
240
241 :param where:
242 Optional SQL expression construct or literal SQL string.
243 If set, emit WHERE <predicate> when issuing DDL
244 for this constraint.
245
246 :param ops:
247 Optional dictionary. Used to define operator classes for the
248 elements; works the same way as that of the
249 :ref:`postgresql_ops <postgresql_operator_classes>`
250 parameter specified to the :class:`_schema.Index` construct.
251
252 .. versionadded:: 1.3.21
253
254 .. seealso::
255
256 :ref:`postgresql_operator_classes` - general description of how
257 PostgreSQL operator classes are specified.
258
259 """
260 columns = []
261 render_exprs = []
262 self.operators = {}
263
264 expressions, operators = zip(*elements)
265
266 for (expr, column, strname, add_element), operator in zip(
267 coercions.expect_col_expression_collection(
268 roles.DDLConstraintColumnRole, expressions
269 ),
270 operators,
271 ):
272 if add_element is not None:
273 columns.append(add_element)
274
275 name = column.name if column is not None else strname
276
277 if name is not None:
278 # backwards compat
279 self.operators[name] = operator
280
281 render_exprs.append((expr, name, operator))
282
283 self._render_exprs = render_exprs
284
285 ColumnCollectionConstraint.__init__(
286 self,
287 *columns,
288 name=kw.get("name"),
289 deferrable=kw.get("deferrable"),
290 initially=kw.get("initially"),
291 )
292 self.using = kw.get("using", "gist")
293 where = kw.get("where")
294 if where is not None:
295 self.where = coercions.expect(roles.StatementOptionRole, where)
296
297 self.ops = kw.get("ops", {})
298
299 def _set_parent(self, table, **kw):
300 super()._set_parent(table)
301
302 self._render_exprs = [
303 (
304 expr if not isinstance(expr, str) else table.c[expr],
305 name,
306 operator,
307 )
308 for expr, name, operator in (self._render_exprs)
309 ]
310
311 def _copy(self, target_table=None, **kw):
312 elements = [
313 (
314 schema._copy_expression(expr, self.parent, target_table),
315 operator,
316 )
317 for expr, _, operator in self._render_exprs
318 ]
319 c = self.__class__(
320 *elements,
321 name=self.name,
322 deferrable=self.deferrable,
323 initially=self.initially,
324 where=self.where,
325 using=self.using,
326 )
327 c.dispatch._update(self.dispatch)
328 return c
329
330
331def array_agg(*arg, **kw):
332 """PostgreSQL-specific form of :class:`_functions.array_agg`, ensures
333 return type is :class:`_postgresql.ARRAY` and not
334 the plain :class:`_types.ARRAY`, unless an explicit ``type_``
335 is passed.
336
337 """
338 kw["_default_array_type"] = ARRAY
339 return functions.func.array_agg(*arg, **kw)
340
341
342class _regconfig_fn(functions.GenericFunction[_T]):
343 inherit_cache = True
344
345 def __init__(self, *args, **kwargs):
346 args = list(args)
347 if len(args) > 1:
348 initial_arg = coercions.expect(
349 roles.ExpressionElementRole,
350 args.pop(0),
351 name=getattr(self, "name", None),
352 apply_propagate_attrs=self,
353 type_=types.REGCONFIG,
354 )
355 initial_arg = [initial_arg]
356 else:
357 initial_arg = []
358
359 addtl_args = [
360 coercions.expect(
361 roles.ExpressionElementRole,
362 c,
363 name=getattr(self, "name", None),
364 apply_propagate_attrs=self,
365 )
366 for c in args
367 ]
368 super().__init__(*(initial_arg + addtl_args), **kwargs)
369
370
371class to_tsvector(_regconfig_fn):
372 """The PostgreSQL ``to_tsvector`` SQL function.
373
374 This function applies automatic casting of the REGCONFIG argument
375 to use the :class:`_postgresql.REGCONFIG` datatype automatically,
376 and applies a return type of :class:`_postgresql.TSVECTOR`.
377
378 Assuming the PostgreSQL dialect has been imported, either by invoking
379 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL
380 engine using ``create_engine("postgresql...")``,
381 :class:`_postgresql.to_tsvector` will be used automatically when invoking
382 ``sqlalchemy.func.to_tsvector()``, ensuring the correct argument and return
383 type handlers are used at compile and execution time.
384
385 .. versionadded:: 2.0.0rc1
386
387 """
388
389 inherit_cache = True
390 type = types.TSVECTOR
391
392
393class to_tsquery(_regconfig_fn):
394 """The PostgreSQL ``to_tsquery`` SQL function.
395
396 This function applies automatic casting of the REGCONFIG argument
397 to use the :class:`_postgresql.REGCONFIG` datatype automatically,
398 and applies a return type of :class:`_postgresql.TSQUERY`.
399
400 Assuming the PostgreSQL dialect has been imported, either by invoking
401 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL
402 engine using ``create_engine("postgresql...")``,
403 :class:`_postgresql.to_tsquery` will be used automatically when invoking
404 ``sqlalchemy.func.to_tsquery()``, ensuring the correct argument and return
405 type handlers are used at compile and execution time.
406
407 .. versionadded:: 2.0.0rc1
408
409 """
410
411 inherit_cache = True
412 type = types.TSQUERY
413
414
415class plainto_tsquery(_regconfig_fn):
416 """The PostgreSQL ``plainto_tsquery`` SQL function.
417
418 This function applies automatic casting of the REGCONFIG argument
419 to use the :class:`_postgresql.REGCONFIG` datatype automatically,
420 and applies a return type of :class:`_postgresql.TSQUERY`.
421
422 Assuming the PostgreSQL dialect has been imported, either by invoking
423 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL
424 engine using ``create_engine("postgresql...")``,
425 :class:`_postgresql.plainto_tsquery` will be used automatically when
426 invoking ``sqlalchemy.func.plainto_tsquery()``, ensuring the correct
427 argument and return type handlers are used at compile and execution time.
428
429 .. versionadded:: 2.0.0rc1
430
431 """
432
433 inherit_cache = True
434 type = types.TSQUERY
435
436
437class phraseto_tsquery(_regconfig_fn):
438 """The PostgreSQL ``phraseto_tsquery`` SQL function.
439
440 This function applies automatic casting of the REGCONFIG argument
441 to use the :class:`_postgresql.REGCONFIG` datatype automatically,
442 and applies a return type of :class:`_postgresql.TSQUERY`.
443
444 Assuming the PostgreSQL dialect has been imported, either by invoking
445 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL
446 engine using ``create_engine("postgresql...")``,
447 :class:`_postgresql.phraseto_tsquery` will be used automatically when
448 invoking ``sqlalchemy.func.phraseto_tsquery()``, ensuring the correct
449 argument and return type handlers are used at compile and execution time.
450
451 .. versionadded:: 2.0.0rc1
452
453 """
454
455 inherit_cache = True
456 type = types.TSQUERY
457
458
459class websearch_to_tsquery(_regconfig_fn):
460 """The PostgreSQL ``websearch_to_tsquery`` SQL function.
461
462 This function applies automatic casting of the REGCONFIG argument
463 to use the :class:`_postgresql.REGCONFIG` datatype automatically,
464 and applies a return type of :class:`_postgresql.TSQUERY`.
465
466 Assuming the PostgreSQL dialect has been imported, either by invoking
467 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL
468 engine using ``create_engine("postgresql...")``,
469 :class:`_postgresql.websearch_to_tsquery` will be used automatically when
470 invoking ``sqlalchemy.func.websearch_to_tsquery()``, ensuring the correct
471 argument and return type handlers are used at compile and execution time.
472
473 .. versionadded:: 2.0.0rc1
474
475 """
476
477 inherit_cache = True
478 type = types.TSQUERY
479
480
481class ts_headline(_regconfig_fn):
482 """The PostgreSQL ``ts_headline`` SQL function.
483
484 This function applies automatic casting of the REGCONFIG argument
485 to use the :class:`_postgresql.REGCONFIG` datatype automatically,
486 and applies a return type of :class:`_types.TEXT`.
487
488 Assuming the PostgreSQL dialect has been imported, either by invoking
489 ``from sqlalchemy.dialects import postgresql``, or by creating a PostgreSQL
490 engine using ``create_engine("postgresql...")``,
491 :class:`_postgresql.ts_headline` will be used automatically when invoking
492 ``sqlalchemy.func.ts_headline()``, ensuring the correct argument and return
493 type handlers are used at compile and execution time.
494
495 .. versionadded:: 2.0.0rc1
496
497 """
498
499 inherit_cache = True
500 type = TEXT
501
502 def __init__(self, *args, **kwargs):
503 args = list(args)
504
505 # parse types according to
506 # https://www.postgresql.org/docs/current/textsearch-controls.html#TEXTSEARCH-HEADLINE
507 if len(args) < 2:
508 # invalid args; don't do anything
509 has_regconfig = False
510 elif (
511 isinstance(args[1], elements.ColumnElement)
512 and args[1].type._type_affinity is types.TSQUERY
513 ):
514 # tsquery is second argument, no regconfig argument
515 has_regconfig = False
516 else:
517 has_regconfig = True
518
519 if has_regconfig:
520 initial_arg = coercions.expect(
521 roles.ExpressionElementRole,
522 args.pop(0),
523 apply_propagate_attrs=self,
524 name=getattr(self, "name", None),
525 type_=types.REGCONFIG,
526 )
527 initial_arg = [initial_arg]
528 else:
529 initial_arg = []
530
531 addtl_args = [
532 coercions.expect(
533 roles.ExpressionElementRole,
534 c,
535 name=getattr(self, "name", None),
536 apply_propagate_attrs=self,
537 )
538 for c in args
539 ]
540 super().__init__(*(initial_arg + addtl_args), **kwargs)
541