You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

394 lines
12KB

  1. # postgresql/array.py
  2. # Copyright (C) 2005-2021 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: http://www.opensource.org/licenses/mit-license.php
  7. import re
  8. from ... import types as sqltypes
  9. from ... import util
  10. from ...sql import coercions
  11. from ...sql import expression
  12. from ...sql import operators
  13. from ...sql import roles
  14. def Any(other, arrexpr, operator=operators.eq):
  15. """A synonym for the :meth:`.ARRAY.Comparator.any` method.
  16. This method is legacy and is here for backwards-compatibility.
  17. .. seealso::
  18. :func:`_expression.any_`
  19. """
  20. return arrexpr.any(other, operator)
  21. def All(other, arrexpr, operator=operators.eq):
  22. """A synonym for the :meth:`.ARRAY.Comparator.all` method.
  23. This method is legacy and is here for backwards-compatibility.
  24. .. seealso::
  25. :func:`_expression.all_`
  26. """
  27. return arrexpr.all(other, operator)
  28. class array(expression.ClauseList, expression.ColumnElement):
  29. """A PostgreSQL ARRAY literal.
  30. This is used to produce ARRAY literals in SQL expressions, e.g.::
  31. from sqlalchemy.dialects.postgresql import array
  32. from sqlalchemy.dialects import postgresql
  33. from sqlalchemy import select, func
  34. stmt = select(array([1,2]) + array([3,4,5]))
  35. print(stmt.compile(dialect=postgresql.dialect()))
  36. Produces the SQL::
  37. SELECT ARRAY[%(param_1)s, %(param_2)s] ||
  38. ARRAY[%(param_3)s, %(param_4)s, %(param_5)s]) AS anon_1
  39. An instance of :class:`.array` will always have the datatype
  40. :class:`_types.ARRAY`. The "inner" type of the array is inferred from
  41. the values present, unless the ``type_`` keyword argument is passed::
  42. array(['foo', 'bar'], type_=CHAR)
  43. Multidimensional arrays are produced by nesting :class:`.array` constructs.
  44. The dimensionality of the final :class:`_types.ARRAY`
  45. type is calculated by
  46. recursively adding the dimensions of the inner :class:`_types.ARRAY`
  47. type::
  48. stmt = select(
  49. array([
  50. array([1, 2]), array([3, 4]), array([column('q'), column('x')])
  51. ])
  52. )
  53. print(stmt.compile(dialect=postgresql.dialect()))
  54. Produces::
  55. SELECT ARRAY[ARRAY[%(param_1)s, %(param_2)s],
  56. ARRAY[%(param_3)s, %(param_4)s], ARRAY[q, x]] AS anon_1
  57. .. versionadded:: 1.3.6 added support for multidimensional array literals
  58. .. seealso::
  59. :class:`_postgresql.ARRAY`
  60. """
  61. __visit_name__ = "array"
  62. stringify_dialect = "postgresql"
  63. def __init__(self, clauses, **kw):
  64. clauses = [
  65. coercions.expect(roles.ExpressionElementRole, c) for c in clauses
  66. ]
  67. super(array, self).__init__(*clauses, **kw)
  68. self._type_tuple = [arg.type for arg in clauses]
  69. main_type = kw.pop(
  70. "type_",
  71. self._type_tuple[0] if self._type_tuple else sqltypes.NULLTYPE,
  72. )
  73. if isinstance(main_type, ARRAY):
  74. self.type = ARRAY(
  75. main_type.item_type,
  76. dimensions=main_type.dimensions + 1
  77. if main_type.dimensions is not None
  78. else 2,
  79. )
  80. else:
  81. self.type = ARRAY(main_type)
  82. @property
  83. def _select_iterable(self):
  84. return (self,)
  85. def _bind_param(self, operator, obj, _assume_scalar=False, type_=None):
  86. if _assume_scalar or operator is operators.getitem:
  87. return expression.BindParameter(
  88. None,
  89. obj,
  90. _compared_to_operator=operator,
  91. type_=type_,
  92. _compared_to_type=self.type,
  93. unique=True,
  94. )
  95. else:
  96. return array(
  97. [
  98. self._bind_param(
  99. operator, o, _assume_scalar=True, type_=type_
  100. )
  101. for o in obj
  102. ]
  103. )
  104. def self_group(self, against=None):
  105. if against in (operators.any_op, operators.all_op, operators.getitem):
  106. return expression.Grouping(self)
  107. else:
  108. return self
  109. CONTAINS = operators.custom_op("@>", precedence=5)
  110. CONTAINED_BY = operators.custom_op("<@", precedence=5)
  111. OVERLAP = operators.custom_op("&&", precedence=5)
  112. class ARRAY(sqltypes.ARRAY):
  113. """PostgreSQL ARRAY type.
  114. .. versionchanged:: 1.1 The :class:`_postgresql.ARRAY` type is now
  115. a subclass of the core :class:`_types.ARRAY` type.
  116. The :class:`_postgresql.ARRAY` type is constructed in the same way
  117. as the core :class:`_types.ARRAY` type; a member type is required, and a
  118. number of dimensions is recommended if the type is to be used for more
  119. than one dimension::
  120. from sqlalchemy.dialects import postgresql
  121. mytable = Table("mytable", metadata,
  122. Column("data", postgresql.ARRAY(Integer, dimensions=2))
  123. )
  124. The :class:`_postgresql.ARRAY` type provides all operations defined on the
  125. core :class:`_types.ARRAY` type, including support for "dimensions",
  126. indexed access, and simple matching such as
  127. :meth:`.types.ARRAY.Comparator.any` and
  128. :meth:`.types.ARRAY.Comparator.all`. :class:`_postgresql.ARRAY`
  129. class also
  130. provides PostgreSQL-specific methods for containment operations, including
  131. :meth:`.postgresql.ARRAY.Comparator.contains`
  132. :meth:`.postgresql.ARRAY.Comparator.contained_by`, and
  133. :meth:`.postgresql.ARRAY.Comparator.overlap`, e.g.::
  134. mytable.c.data.contains([1, 2])
  135. The :class:`_postgresql.ARRAY` type may not be supported on all
  136. PostgreSQL DBAPIs; it is currently known to work on psycopg2 only.
  137. Additionally, the :class:`_postgresql.ARRAY`
  138. type does not work directly in
  139. conjunction with the :class:`.ENUM` type. For a workaround, see the
  140. special type at :ref:`postgresql_array_of_enum`.
  141. .. seealso::
  142. :class:`_types.ARRAY` - base array type
  143. :class:`_postgresql.array` - produces a literal array value.
  144. """
  145. class Comparator(sqltypes.ARRAY.Comparator):
  146. """Define comparison operations for :class:`_types.ARRAY`.
  147. Note that these operations are in addition to those provided
  148. by the base :class:`.types.ARRAY.Comparator` class, including
  149. :meth:`.types.ARRAY.Comparator.any` and
  150. :meth:`.types.ARRAY.Comparator.all`.
  151. """
  152. def contains(self, other, **kwargs):
  153. """Boolean expression. Test if elements are a superset of the
  154. elements of the argument array expression.
  155. """
  156. return self.operate(CONTAINS, other, result_type=sqltypes.Boolean)
  157. def contained_by(self, other):
  158. """Boolean expression. Test if elements are a proper subset of the
  159. elements of the argument array expression.
  160. """
  161. return self.operate(
  162. CONTAINED_BY, other, result_type=sqltypes.Boolean
  163. )
  164. def overlap(self, other):
  165. """Boolean expression. Test if array has elements in common with
  166. an argument array expression.
  167. """
  168. return self.operate(OVERLAP, other, result_type=sqltypes.Boolean)
  169. comparator_factory = Comparator
  170. def __init__(
  171. self, item_type, as_tuple=False, dimensions=None, zero_indexes=False
  172. ):
  173. """Construct an ARRAY.
  174. E.g.::
  175. Column('myarray', ARRAY(Integer))
  176. Arguments are:
  177. :param item_type: The data type of items of this array. Note that
  178. dimensionality is irrelevant here, so multi-dimensional arrays like
  179. ``INTEGER[][]``, are constructed as ``ARRAY(Integer)``, not as
  180. ``ARRAY(ARRAY(Integer))`` or such.
  181. :param as_tuple=False: Specify whether return results
  182. should be converted to tuples from lists. DBAPIs such
  183. as psycopg2 return lists by default. When tuples are
  184. returned, the results are hashable.
  185. :param dimensions: if non-None, the ARRAY will assume a fixed
  186. number of dimensions. This will cause the DDL emitted for this
  187. ARRAY to include the exact number of bracket clauses ``[]``,
  188. and will also optimize the performance of the type overall.
  189. Note that PG arrays are always implicitly "non-dimensioned",
  190. meaning they can store any number of dimensions no matter how
  191. they were declared.
  192. :param zero_indexes=False: when True, index values will be converted
  193. between Python zero-based and PostgreSQL one-based indexes, e.g.
  194. a value of one will be added to all index values before passing
  195. to the database.
  196. .. versionadded:: 0.9.5
  197. """
  198. if isinstance(item_type, ARRAY):
  199. raise ValueError(
  200. "Do not nest ARRAY types; ARRAY(basetype) "
  201. "handles multi-dimensional arrays of basetype"
  202. )
  203. if isinstance(item_type, type):
  204. item_type = item_type()
  205. self.item_type = item_type
  206. self.as_tuple = as_tuple
  207. self.dimensions = dimensions
  208. self.zero_indexes = zero_indexes
  209. @property
  210. def hashable(self):
  211. return self.as_tuple
  212. @property
  213. def python_type(self):
  214. return list
  215. def compare_values(self, x, y):
  216. return x == y
  217. def _proc_array(self, arr, itemproc, dim, collection):
  218. if dim is None:
  219. arr = list(arr)
  220. if (
  221. dim == 1
  222. or dim is None
  223. and (
  224. # this has to be (list, tuple), or at least
  225. # not hasattr('__iter__'), since Py3K strings
  226. # etc. have __iter__
  227. not arr
  228. or not isinstance(arr[0], (list, tuple))
  229. )
  230. ):
  231. if itemproc:
  232. return collection(itemproc(x) for x in arr)
  233. else:
  234. return collection(arr)
  235. else:
  236. return collection(
  237. self._proc_array(
  238. x,
  239. itemproc,
  240. dim - 1 if dim is not None else None,
  241. collection,
  242. )
  243. for x in arr
  244. )
  245. @util.memoized_property
  246. def _against_native_enum(self):
  247. return (
  248. isinstance(self.item_type, sqltypes.Enum)
  249. and self.item_type.native_enum
  250. )
  251. def bind_expression(self, bindvalue):
  252. return bindvalue
  253. def bind_processor(self, dialect):
  254. item_proc = self.item_type.dialect_impl(dialect).bind_processor(
  255. dialect
  256. )
  257. def process(value):
  258. if value is None:
  259. return value
  260. else:
  261. return self._proc_array(
  262. value, item_proc, self.dimensions, list
  263. )
  264. return process
  265. def result_processor(self, dialect, coltype):
  266. item_proc = self.item_type.dialect_impl(dialect).result_processor(
  267. dialect, coltype
  268. )
  269. def process(value):
  270. if value is None:
  271. return value
  272. else:
  273. return self._proc_array(
  274. value,
  275. item_proc,
  276. self.dimensions,
  277. tuple if self.as_tuple else list,
  278. )
  279. if self._against_native_enum:
  280. super_rp = process
  281. def handle_raw_string(value):
  282. inner = re.match(r"^{(.*)}$", value).group(1)
  283. return inner.split(",") if inner else []
  284. def process(value):
  285. if value is None:
  286. return value
  287. # isinstance(value, util.string_types) is required to handle
  288. # the # case where a TypeDecorator for and Array of Enum is
  289. # used like was required in sa < 1.3.17
  290. return super_rp(
  291. handle_raw_string(value)
  292. if isinstance(value, util.string_types)
  293. else value
  294. )
  295. return process