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
8from .array import ARRAY
9from ... import util
10from ...sql import coercions
11from ...sql import elements
12from ...sql import expression
13from ...sql import functions
14from ...sql import roles
15from ...sql import schema
16from ...sql.schema import ColumnCollectionConstraint
17from ...sql.visitors import InternalTraversal
18
19
20class aggregate_order_by(expression.ColumnElement):
21 """Represent a PostgreSQL aggregate order by expression.
22
23 E.g.::
24
25 from sqlalchemy.dialects.postgresql import aggregate_order_by
26 expr = func.array_agg(aggregate_order_by(table.c.a, table.c.b.desc()))
27 stmt = select(expr)
28
29 would represent the expression::
30
31 SELECT array_agg(a ORDER BY b DESC) FROM table;
32
33 Similarly::
34
35 expr = func.string_agg(
36 table.c.a,
37 aggregate_order_by(literal_column("','"), table.c.a)
38 )
39 stmt = select(expr)
40
41 Would represent::
42
43 SELECT string_agg(a, ',' ORDER BY a) FROM table;
44
45 .. versionadded:: 1.1
46
47 .. versionchanged:: 1.2.13 - the ORDER BY argument may be multiple terms
48
49 .. seealso::
50
51 :class:`_functions.array_agg`
52
53 """
54
55 __visit_name__ = "aggregate_order_by"
56
57 stringify_dialect = "postgresql"
58 _traverse_internals = [
59 ("target", InternalTraversal.dp_clauseelement),
60 ("type", InternalTraversal.dp_type),
61 ("order_by", InternalTraversal.dp_clauseelement),
62 ]
63
64 def __init__(self, target, *order_by):
65 self.target = coercions.expect(roles.ExpressionElementRole, target)
66 self.type = self.target.type
67
68 _lob = len(order_by)
69 if _lob == 0:
70 raise TypeError("at least one ORDER BY element is required")
71 elif _lob == 1:
72 self.order_by = coercions.expect(
73 roles.ExpressionElementRole, order_by[0]
74 )
75 else:
76 self.order_by = elements.ClauseList(
77 *order_by, _literal_as_text_role=roles.ExpressionElementRole
78 )
79
80 def self_group(self, against=None):
81 return self
82
83 def get_children(self, **kwargs):
84 return self.target, self.order_by
85
86 def _copy_internals(self, clone=elements._clone, **kw):
87 self.target = clone(self.target, **kw)
88 self.order_by = clone(self.order_by, **kw)
89
90 @property
91 def _from_objects(self):
92 return self.target._from_objects + self.order_by._from_objects
93
94
95class ExcludeConstraint(ColumnCollectionConstraint):
96 """A table-level EXCLUDE constraint.
97
98 Defines an EXCLUDE constraint as described in the `PostgreSQL
99 documentation`__.
100
101 __ https://www.postgresql.org/docs/current/static/sql-createtable.html#SQL-CREATETABLE-EXCLUDE
102
103 """ # noqa
104
105 __visit_name__ = "exclude_constraint"
106
107 where = None
108 inherit_cache = False
109
110 create_drop_stringify_dialect = "postgresql"
111
112 @elements._document_text_coercion(
113 "where",
114 ":class:`.ExcludeConstraint`",
115 ":paramref:`.ExcludeConstraint.where`",
116 )
117 def __init__(self, *elements, **kw):
118 r"""
119 Create an :class:`.ExcludeConstraint` object.
120
121 E.g.::
122
123 const = ExcludeConstraint(
124 (Column('period'), '&&'),
125 (Column('group'), '='),
126 where=(Column('group') != 'some group'),
127 ops={'group': 'my_operator_class'}
128 )
129
130 The constraint is normally embedded into the :class:`_schema.Table`
131 construct
132 directly, or added later using :meth:`.append_constraint`::
133
134 some_table = Table(
135 'some_table', metadata,
136 Column('id', Integer, primary_key=True),
137 Column('period', TSRANGE()),
138 Column('group', String)
139 )
140
141 some_table.append_constraint(
142 ExcludeConstraint(
143 (some_table.c.period, '&&'),
144 (some_table.c.group, '='),
145 where=some_table.c.group != 'some group',
146 name='some_table_excl_const',
147 ops={'group': 'my_operator_class'}
148 )
149 )
150
151 :param \*elements:
152
153 A sequence of two tuples of the form ``(column, operator)`` where
154 "column" is a SQL expression element or a raw SQL string, most
155 typically a :class:`_schema.Column` object,
156 and "operator" is a string
157 containing the operator to use. In order to specify a column name
158 when a :class:`_schema.Column` object is not available,
159 while ensuring
160 that any necessary quoting rules take effect, an ad-hoc
161 :class:`_schema.Column` or :func:`_expression.column`
162 object should be
163 used.
164
165 :param name:
166 Optional, the in-database name of this constraint.
167
168 :param deferrable:
169 Optional bool. If set, emit DEFERRABLE or NOT DEFERRABLE when
170 issuing DDL for this constraint.
171
172 :param initially:
173 Optional string. If set, emit INITIALLY <value> when issuing DDL
174 for this constraint.
175
176 :param using:
177 Optional string. If set, emit USING <index_method> when issuing DDL
178 for this constraint. Defaults to 'gist'.
179
180 :param where:
181 Optional SQL expression construct or literal SQL string.
182 If set, emit WHERE <predicate> when issuing DDL
183 for this constraint.
184
185 :param ops:
186 Optional dictionary. Used to define operator classes for the
187 elements; works the same way as that of the
188 :ref:`postgresql_ops <postgresql_operator_classes>`
189 parameter specified to the :class:`_schema.Index` construct.
190
191 .. versionadded:: 1.3.21
192
193 .. seealso::
194
195 :ref:`postgresql_operator_classes` - general description of how
196 PostgreSQL operator classes are specified.
197
198 """
199 columns = []
200 render_exprs = []
201 self.operators = {}
202
203 expressions, operators = zip(*elements)
204
205 for (expr, column, strname, add_element), operator in zip(
206 coercions.expect_col_expression_collection(
207 roles.DDLConstraintColumnRole, expressions
208 ),
209 operators,
210 ):
211 if add_element is not None:
212 columns.append(add_element)
213
214 name = column.name if column is not None else strname
215
216 if name is not None:
217 # backwards compat
218 self.operators[name] = operator
219
220 render_exprs.append((expr, name, operator))
221
222 self._render_exprs = render_exprs
223
224 ColumnCollectionConstraint.__init__(
225 self,
226 *columns,
227 name=kw.get("name"),
228 deferrable=kw.get("deferrable"),
229 initially=kw.get("initially")
230 )
231 self.using = kw.get("using", "gist")
232 where = kw.get("where")
233 if where is not None:
234 self.where = coercions.expect(roles.StatementOptionRole, where)
235
236 self.ops = kw.get("ops", {})
237
238 def _set_parent(self, table, **kw):
239 super(ExcludeConstraint, self)._set_parent(table)
240
241 self._render_exprs = [
242 (
243 expr if isinstance(expr, elements.ClauseElement) else colexpr,
244 name,
245 operator,
246 )
247 for (expr, name, operator), colexpr in util.zip_longest(
248 self._render_exprs, self.columns
249 )
250 ]
251
252 def _copy(self, target_table=None, **kw):
253 elements = [
254 (
255 schema._copy_expression(expr, self.parent, target_table),
256 self.operators[expr.name],
257 )
258 for expr in self.columns
259 ]
260 c = self.__class__(
261 *elements,
262 name=self.name,
263 deferrable=self.deferrable,
264 initially=self.initially,
265 where=self.where,
266 using=self.using
267 )
268 c.dispatch._update(self.dispatch)
269 return c
270
271
272def array_agg(*arg, **kw):
273 """PostgreSQL-specific form of :class:`_functions.array_agg`, ensures
274 return type is :class:`_postgresql.ARRAY` and not
275 the plain :class:`_types.ARRAY`, unless an explicit ``type_``
276 is passed.
277
278 .. versionadded:: 1.1
279
280 """
281 kw["_default_array_type"] = ARRAY
282 return functions.func.array_agg(*arg, **kw)