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