codekingpro/portable-devtools
114k
1# dialects/postgresql/array.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
8
9
10from __future__ import annotations
11
12import re
13from typing import Any
14from typing import Optional
15from typing import TypeVar
16
17from .operators import CONTAINED_BY
18from .operators import CONTAINS
19from .operators import OVERLAP
20from ... import types as sqltypes
21from ... import util
22from ...sql import expression
23from ...sql import operators
24from ...sql._typing import _TypeEngineArgument
25
26
27_T = TypeVar("_T", bound=Any)
28
29
30def Any(other, arrexpr, operator=operators.eq):
31 """A synonym for the ARRAY-level :meth:`.ARRAY.Comparator.any` method.
32 See that method for details.
33
34 """
35
36 return arrexpr.any(other, operator)
37
38
39def All(other, arrexpr, operator=operators.eq):
40 """A synonym for the ARRAY-level :meth:`.ARRAY.Comparator.all` method.
41 See that method for details.
42
43 """
44
45 return arrexpr.all(other, operator)
46
47
48class array(expression.ExpressionClauseList[_T]):
49 """A PostgreSQL ARRAY literal.
50
51 This is used to produce ARRAY literals in SQL expressions, e.g.::
52
53 from sqlalchemy.dialects.postgresql import array
54 from sqlalchemy.dialects import postgresql
55 from sqlalchemy import select, func
56
57 stmt = select(array([1,2]) + array([3,4,5]))
58
59 print(stmt.compile(dialect=postgresql.dialect()))
60
61 Produces the SQL::
62
63 SELECT ARRAY[%(param_1)s, %(param_2)s] ||
64 ARRAY[%(param_3)s, %(param_4)s, %(param_5)s]) AS anon_1
65
66 An instance of :class:`.array` will always have the datatype
67 :class:`_types.ARRAY`. The "inner" type of the array is inferred from
68 the values present, unless the ``type_`` keyword argument is passed::
69
70 array(['foo', 'bar'], type_=CHAR)
71
72 Multidimensional arrays are produced by nesting :class:`.array` constructs.
73 The dimensionality of the final :class:`_types.ARRAY`
74 type is calculated by
75 recursively adding the dimensions of the inner :class:`_types.ARRAY`
76 type::
77
78 stmt = select(
79 array([
80 array([1, 2]), array([3, 4]), array([column('q'), column('x')])
81 ])
82 )
83 print(stmt.compile(dialect=postgresql.dialect()))
84
85 Produces::
86
87 SELECT ARRAY[ARRAY[%(param_1)s, %(param_2)s],
88 ARRAY[%(param_3)s, %(param_4)s], ARRAY[q, x]] AS anon_1
89
90 .. versionadded:: 1.3.6 added support for multidimensional array literals
91
92 .. seealso::
93
94 :class:`_postgresql.ARRAY`
95
96 """
97
98 __visit_name__ = "array"
99
100 stringify_dialect = "postgresql"
101 inherit_cache = True
102
103 def __init__(self, clauses, **kw):
104 type_arg = kw.pop("type_", None)
105 super().__init__(operators.comma_op, *clauses, **kw)
106
107 self._type_tuple = [arg.type for arg in self.clauses]
108
109 main_type = (
110 type_arg
111 if type_arg is not None
112 else self._type_tuple[0] if self._type_tuple else sqltypes.NULLTYPE
113 )
114
115 if isinstance(main_type, ARRAY):
116 self.type = ARRAY(
117 main_type.item_type,
118 dimensions=(
119 main_type.dimensions + 1
120 if main_type.dimensions is not None
121 else 2
122 ),
123 )
124 else:
125 self.type = ARRAY(main_type)
126
127 @property
128 def _select_iterable(self):
129 return (self,)
130
131 def _bind_param(self, operator, obj, _assume_scalar=False, type_=None):
132 if _assume_scalar or operator is operators.getitem:
133 return expression.BindParameter(
134 None,
135 obj,
136 _compared_to_operator=operator,
137 type_=type_,
138 _compared_to_type=self.type,
139 unique=True,
140 )
141
142 else:
143 return array(
144 [
145 self._bind_param(
146 operator, o, _assume_scalar=True, type_=type_
147 )
148 for o in obj
149 ]
150 )
151
152 def self_group(self, against=None):
153 if against in (operators.any_op, operators.all_op, operators.getitem):
154 return expression.Grouping(self)
155 else:
156 return self
157
158
159class ARRAY(sqltypes.ARRAY):
160 """PostgreSQL ARRAY type.
161
162 The :class:`_postgresql.ARRAY` type is constructed in the same way
163 as the core :class:`_types.ARRAY` type; a member type is required, and a
164 number of dimensions is recommended if the type is to be used for more
165 than one dimension::
166
167 from sqlalchemy.dialects import postgresql
168
169 mytable = Table("mytable", metadata,
170 Column("data", postgresql.ARRAY(Integer, dimensions=2))
171 )
172
173 The :class:`_postgresql.ARRAY` type provides all operations defined on the
174 core :class:`_types.ARRAY` type, including support for "dimensions",
175 indexed access, and simple matching such as
176 :meth:`.types.ARRAY.Comparator.any` and
177 :meth:`.types.ARRAY.Comparator.all`. :class:`_postgresql.ARRAY`
178 class also
179 provides PostgreSQL-specific methods for containment operations, including
180 :meth:`.postgresql.ARRAY.Comparator.contains`
181 :meth:`.postgresql.ARRAY.Comparator.contained_by`, and
182 :meth:`.postgresql.ARRAY.Comparator.overlap`, e.g.::
183
184 mytable.c.data.contains([1, 2])
185
186 Indexed access is one-based by default, to match that of PostgreSQL;
187 for zero-based indexed access, set
188 :paramref:`_postgresql.ARRAY.zero_indexes`.
189
190 Additionally, the :class:`_postgresql.ARRAY`
191 type does not work directly in
192 conjunction with the :class:`.ENUM` type. For a workaround, see the
193 special type at :ref:`postgresql_array_of_enum`.
194
195 .. container:: topic
196
197 **Detecting Changes in ARRAY columns when using the ORM**
198
199 The :class:`_postgresql.ARRAY` type, when used with the SQLAlchemy ORM,
200 does not detect in-place mutations to the array. In order to detect
201 these, the :mod:`sqlalchemy.ext.mutable` extension must be used, using
202 the :class:`.MutableList` class::
203
204 from sqlalchemy.dialects.postgresql import ARRAY
205 from sqlalchemy.ext.mutable import MutableList
206
207 class SomeOrmClass(Base):
208 # ...
209
210 data = Column(MutableList.as_mutable(ARRAY(Integer)))
211
212 This extension will allow "in-place" changes such to the array
213 such as ``.append()`` to produce events which will be detected by the
214 unit of work. Note that changes to elements **inside** the array,
215 including subarrays that are mutated in place, are **not** detected.
216
217 Alternatively, assigning a new array value to an ORM element that
218 replaces the old one will always trigger a change event.
219
220 .. seealso::
221
222 :class:`_types.ARRAY` - base array type
223
224 :class:`_postgresql.array` - produces a literal array value.
225
226 """
227
228 def __init__(
229 self,
230 item_type: _TypeEngineArgument[Any],
231 as_tuple: bool = False,
232 dimensions: Optional[int] = None,
233 zero_indexes: bool = False,
234 ):
235 """Construct an ARRAY.
236
237 E.g.::
238
239 Column('myarray', ARRAY(Integer))
240
241 Arguments are:
242
243 :param item_type: The data type of items of this array. Note that
244 dimensionality is irrelevant here, so multi-dimensional arrays like
245 ``INTEGER[][]``, are constructed as ``ARRAY(Integer)``, not as
246 ``ARRAY(ARRAY(Integer))`` or such.
247
248 :param as_tuple=False: Specify whether return results
249 should be converted to tuples from lists. DBAPIs such
250 as psycopg2 return lists by default. When tuples are
251 returned, the results are hashable.
252
253 :param dimensions: if non-None, the ARRAY will assume a fixed
254 number of dimensions. This will cause the DDL emitted for this
255 ARRAY to include the exact number of bracket clauses ``[]``,
256 and will also optimize the performance of the type overall.
257 Note that PG arrays are always implicitly "non-dimensioned",
258 meaning they can store any number of dimensions no matter how
259 they were declared.
260
261 :param zero_indexes=False: when True, index values will be converted
262 between Python zero-based and PostgreSQL one-based indexes, e.g.
263 a value of one will be added to all index values before passing
264 to the database.
265
266 """
267 if isinstance(item_type, ARRAY):
268 raise ValueError(
269 "Do not nest ARRAY types; ARRAY(basetype) "
270 "handles multi-dimensional arrays of basetype"
271 )
272 if isinstance(item_type, type):
273 item_type = item_type()
274 self.item_type = item_type
275 self.as_tuple = as_tuple
276 self.dimensions = dimensions
277 self.zero_indexes = zero_indexes
278
279 class Comparator(sqltypes.ARRAY.Comparator):
280 """Define comparison operations for :class:`_types.ARRAY`.
281
282 Note that these operations are in addition to those provided
283 by the base :class:`.types.ARRAY.Comparator` class, including
284 :meth:`.types.ARRAY.Comparator.any` and
285 :meth:`.types.ARRAY.Comparator.all`.
286
287 """
288
289 def contains(self, other, **kwargs):
290 """Boolean expression. Test if elements are a superset of the
291 elements of the argument array expression.
292
293 kwargs may be ignored by this operator but are required for API
294 conformance.
295 """
296 return self.operate(CONTAINS, other, result_type=sqltypes.Boolean)
297
298 def contained_by(self, other):
299 """Boolean expression. Test if elements are a proper subset of the
300 elements of the argument array expression.
301 """
302 return self.operate(
303 CONTAINED_BY, other, result_type=sqltypes.Boolean
304 )
305
306 def overlap(self, other):
307 """Boolean expression. Test if array has elements in common with
308 an argument array expression.
309 """
310 return self.operate(OVERLAP, other, result_type=sqltypes.Boolean)
311
312 comparator_factory = Comparator
313
314 @property
315 def hashable(self):
316 return self.as_tuple
317
318 @property
319 def python_type(self):
320 return list
321
322 def compare_values(self, x, y):
323 return x == y
324
325 @util.memoized_property
326 def _against_native_enum(self):
327 return (
328 isinstance(self.item_type, sqltypes.Enum)
329 and self.item_type.native_enum
330 )
331
332 def literal_processor(self, dialect):
333 item_proc = self.item_type.dialect_impl(dialect).literal_processor(
334 dialect
335 )
336 if item_proc is None:
337 return None
338
339 def to_str(elements):
340 return f"ARRAY[{', '.join(elements)}]"
341
342 def process(value):
343 inner = self._apply_item_processor(
344 value, item_proc, self.dimensions, to_str
345 )
346 return inner
347
348 return process
349
350 def bind_processor(self, dialect):
351 item_proc = self.item_type.dialect_impl(dialect).bind_processor(
352 dialect
353 )
354
355 def process(value):
356 if value is None:
357 return value
358 else:
359 return self._apply_item_processor(
360 value, item_proc, self.dimensions, list
361 )
362
363 return process
364
365 def result_processor(self, dialect, coltype):
366 item_proc = self.item_type.dialect_impl(dialect).result_processor(
367 dialect, coltype
368 )
369
370 def process(value):
371 if value is None:
372 return value
373 else:
374 return self._apply_item_processor(
375 value,
376 item_proc,
377 self.dimensions,
378 tuple if self.as_tuple else list,
379 )
380
381 if self._against_native_enum:
382 super_rp = process
383 pattern = re.compile(r"^{(.*)}$")
384
385 def handle_raw_string(value):
386 inner = pattern.match(value).group(1)
387 return _split_enum_values(inner)
388
389 def process(value):
390 if value is None:
391 return value
392 # isinstance(value, str) is required to handle
393 # the case where a TypeDecorator for and Array of Enum is
394 # used like was required in sa < 1.3.17
395 return super_rp(
396 handle_raw_string(value)
397 if isinstance(value, str)
398 else value
399 )
400
401 return process
402
403
404def _split_enum_values(array_string):
405 if '"' not in array_string:
406 # no escape char is present so it can just split on the comma
407 return array_string.split(",") if array_string else []
408
409 # handles quoted strings from:
410 # r'abc,"quoted","also\\\\quoted", "quoted, comma", "esc \" quot", qpr'
411 # returns
412 # ['abc', 'quoted', 'also\\quoted', 'quoted, comma', 'esc " quot', 'qpr']
413 text = array_string.replace(r"\"", "_$ESC_QUOTE$_")
414 text = text.replace(r"\\", "\\")
415 result = []
416 on_quotes = re.split(r'(")', text)
417 in_quotes = False
418 for tok in on_quotes:
419 if tok == '"':
420 in_quotes = not in_quotes
421 elif in_quotes:
422 result.append(tok.replace("_$ESC_QUOTE$_", '"'))
423 else:
424 result.extend(re.findall(r"([^\s,]+),?", tok))
425 return result
426