ext.py 8.5 KB

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