dml.py 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274
  1. # postgresql/on_conflict.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 . import ext
  8. from ... import util
  9. from ...sql import coercions
  10. from ...sql import roles
  11. from ...sql import schema
  12. from ...sql.base import _exclusive_against
  13. from ...sql.base import _generative
  14. from ...sql.base import ColumnCollection
  15. from ...sql.dml import Insert as StandardInsert
  16. from ...sql.elements import ClauseElement
  17. from ...sql.expression import alias
  18. from ...util.langhelpers import public_factory
  19. __all__ = ("Insert", "insert")
  20. class Insert(StandardInsert):
  21. """PostgreSQL-specific implementation of INSERT.
  22. Adds methods for PG-specific syntaxes such as ON CONFLICT.
  23. The :class:`_postgresql.Insert` object is created using the
  24. :func:`sqlalchemy.dialects.postgresql.insert` function.
  25. .. versionadded:: 1.1
  26. """
  27. stringify_dialect = "postgresql"
  28. inherit_cache = False
  29. @util.memoized_property
  30. def excluded(self):
  31. """Provide the ``excluded`` namespace for an ON CONFLICT statement
  32. PG's ON CONFLICT clause allows reference to the row that would
  33. be inserted, known as ``excluded``. This attribute provides
  34. all columns in this row to be referenceable.
  35. .. tip:: The :attr:`_postgresql.Insert.excluded` attribute is an
  36. instance of :class:`_expression.ColumnCollection`, which provides
  37. an interface the same as that of the :attr:`_schema.Table.c`
  38. collection described at :ref:`metadata_tables_and_columns`.
  39. With this collection, ordinary names are accessible like attributes
  40. (e.g. ``stmt.excluded.some_column``), but special names and
  41. dictionary method names should be accessed using indexed access,
  42. such as ``stmt.excluded["column name"]`` or
  43. ``stmt.excluded["values"]``. See the docstring for
  44. :class:`_expression.ColumnCollection` for further examples.
  45. .. seealso::
  46. :ref:`postgresql_insert_on_conflict` - example of how
  47. to use :attr:`_expression.Insert.excluded`
  48. """
  49. return alias(self.table, name="excluded").columns
  50. _on_conflict_exclusive = _exclusive_against(
  51. "_post_values_clause",
  52. msgs={
  53. "_post_values_clause": "This Insert construct already has "
  54. "an ON CONFLICT clause established"
  55. },
  56. )
  57. @_generative
  58. @_on_conflict_exclusive
  59. def on_conflict_do_update(
  60. self,
  61. constraint=None,
  62. index_elements=None,
  63. index_where=None,
  64. set_=None,
  65. where=None,
  66. ):
  67. r"""
  68. Specifies a DO UPDATE SET action for ON CONFLICT clause.
  69. Either the ``constraint`` or ``index_elements`` argument is
  70. required, but only one of these can be specified.
  71. :param constraint:
  72. The name of a unique or exclusion constraint on the table,
  73. or the constraint object itself if it has a .name attribute.
  74. :param index_elements:
  75. A sequence consisting of string column names, :class:`_schema.Column`
  76. objects, or other column expression objects that will be used
  77. to infer a target index.
  78. :param index_where:
  79. Additional WHERE criterion that can be used to infer a
  80. conditional target index.
  81. :param set\_:
  82. A dictionary or other mapping object
  83. where the keys are either names of columns in the target table,
  84. or :class:`_schema.Column` objects or other ORM-mapped columns
  85. matching that of the target table, and expressions or literals
  86. as values, specifying the ``SET`` actions to take.
  87. .. versionadded:: 1.4 The
  88. :paramref:`_postgresql.Insert.on_conflict_do_update.set_`
  89. parameter supports :class:`_schema.Column` objects from the target
  90. :class:`_schema.Table` as keys.
  91. .. warning:: This dictionary does **not** take into account
  92. Python-specified default UPDATE values or generation functions,
  93. e.g. those specified using :paramref:`_schema.Column.onupdate`.
  94. These values will not be exercised for an ON CONFLICT style of
  95. UPDATE, unless they are manually specified in the
  96. :paramref:`.Insert.on_conflict_do_update.set_` dictionary.
  97. :param where:
  98. Optional argument. If present, can be a literal SQL
  99. string or an acceptable expression for a ``WHERE`` clause
  100. that restricts the rows affected by ``DO UPDATE SET``. Rows
  101. not meeting the ``WHERE`` condition will not be updated
  102. (effectively a ``DO NOTHING`` for those rows).
  103. .. versionadded:: 1.1
  104. .. seealso::
  105. :ref:`postgresql_insert_on_conflict`
  106. """
  107. self._post_values_clause = OnConflictDoUpdate(
  108. constraint, index_elements, index_where, set_, where
  109. )
  110. @_generative
  111. @_on_conflict_exclusive
  112. def on_conflict_do_nothing(
  113. self, constraint=None, index_elements=None, index_where=None
  114. ):
  115. """
  116. Specifies a DO NOTHING action for ON CONFLICT clause.
  117. The ``constraint`` and ``index_elements`` arguments
  118. are optional, but only one of these can be specified.
  119. :param constraint:
  120. The name of a unique or exclusion constraint on the table,
  121. or the constraint object itself if it has a .name attribute.
  122. :param index_elements:
  123. A sequence consisting of string column names, :class:`_schema.Column`
  124. objects, or other column expression objects that will be used
  125. to infer a target index.
  126. :param index_where:
  127. Additional WHERE criterion that can be used to infer a
  128. conditional target index.
  129. .. versionadded:: 1.1
  130. .. seealso::
  131. :ref:`postgresql_insert_on_conflict`
  132. """
  133. self._post_values_clause = OnConflictDoNothing(
  134. constraint, index_elements, index_where
  135. )
  136. insert = public_factory(
  137. Insert, ".dialects.postgresql.insert", ".dialects.postgresql.Insert"
  138. )
  139. class OnConflictClause(ClauseElement):
  140. stringify_dialect = "postgresql"
  141. def __init__(self, constraint=None, index_elements=None, index_where=None):
  142. if constraint is not None:
  143. if not isinstance(constraint, util.string_types) and isinstance(
  144. constraint,
  145. (schema.Index, schema.Constraint, ext.ExcludeConstraint),
  146. ):
  147. constraint = getattr(constraint, "name") or constraint
  148. if constraint is not None:
  149. if index_elements is not None:
  150. raise ValueError(
  151. "'constraint' and 'index_elements' are mutually exclusive"
  152. )
  153. if isinstance(constraint, util.string_types):
  154. self.constraint_target = constraint
  155. self.inferred_target_elements = None
  156. self.inferred_target_whereclause = None
  157. elif isinstance(constraint, schema.Index):
  158. index_elements = constraint.expressions
  159. index_where = constraint.dialect_options["postgresql"].get(
  160. "where"
  161. )
  162. elif isinstance(constraint, ext.ExcludeConstraint):
  163. index_elements = constraint.columns
  164. index_where = constraint.where
  165. else:
  166. index_elements = constraint.columns
  167. index_where = constraint.dialect_options["postgresql"].get(
  168. "where"
  169. )
  170. if index_elements is not None:
  171. self.constraint_target = None
  172. self.inferred_target_elements = index_elements
  173. self.inferred_target_whereclause = index_where
  174. elif constraint is None:
  175. self.constraint_target = (
  176. self.inferred_target_elements
  177. ) = self.inferred_target_whereclause = None
  178. class OnConflictDoNothing(OnConflictClause):
  179. __visit_name__ = "on_conflict_do_nothing"
  180. class OnConflictDoUpdate(OnConflictClause):
  181. __visit_name__ = "on_conflict_do_update"
  182. def __init__(
  183. self,
  184. constraint=None,
  185. index_elements=None,
  186. index_where=None,
  187. set_=None,
  188. where=None,
  189. ):
  190. super(OnConflictDoUpdate, self).__init__(
  191. constraint=constraint,
  192. index_elements=index_elements,
  193. index_where=index_where,
  194. )
  195. if (
  196. self.inferred_target_elements is None
  197. and self.constraint_target is None
  198. ):
  199. raise ValueError(
  200. "Either constraint or index_elements, "
  201. "but not both, must be specified unless DO NOTHING"
  202. )
  203. if isinstance(set_, dict):
  204. if not set_:
  205. raise ValueError("set parameter dictionary must not be empty")
  206. elif isinstance(set_, ColumnCollection):
  207. set_ = dict(set_)
  208. else:
  209. raise ValueError(
  210. "set parameter must be a non-empty dictionary "
  211. "or a ColumnCollection such as the `.c.` collection "
  212. "of a Table object"
  213. )
  214. self.update_values_to_set = [
  215. (coercions.expect(roles.DMLColumnRole, key), value)
  216. for key, value in set_.items()
  217. ]
  218. self.update_whereclause = where