array.py 17 KB

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