codekingpro/portable-devtools

1# dialects/postgresql/dml.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
7from __future__ import annotations
8
9from typing import Any
10from typing import Optional
11
12from . import ext
13from .._typing import _OnConflictConstraintT
14from .._typing import _OnConflictIndexElementsT
15from .._typing import _OnConflictIndexWhereT
16from .._typing import _OnConflictSetT
17from .._typing import _OnConflictWhereT
18from ... import util
19from ...sql import coercions
20from ...sql import roles
21from ...sql import schema
22from ...sql._typing import _DMLTableArgument
23from ...sql.base import _exclusive_against
24from ...sql.base import _generative
25from ...sql.base import ColumnCollection
26from ...sql.base import ReadOnlyColumnCollection
27from ...sql.dml import Insert as StandardInsert
28from ...sql.elements import ClauseElement
29from ...sql.elements import KeyedColumnElement
30from ...sql.expression import alias
31from ...util.typing import Self
32
33
34__all__ = ("Insert", "insert")
35
36
37def insert(table: _DMLTableArgument) -> Insert:
38    """Construct a PostgreSQL-specific variant :class:`_postgresql.Insert`
39    construct.
40
41    .. container:: inherited_member
42
43        The :func:`sqlalchemy.dialects.postgresql.insert` function creates
44        a :class:`sqlalchemy.dialects.postgresql.Insert`.  This class is based
45        on the dialect-agnostic :class:`_sql.Insert` construct which may
46        be constructed using the :func:`_sql.insert` function in
47        SQLAlchemy Core.
48
49    The :class:`_postgresql.Insert` construct includes additional methods
50    :meth:`_postgresql.Insert.on_conflict_do_update`,
51    :meth:`_postgresql.Insert.on_conflict_do_nothing`.
52
53    """
54    return Insert(table)
55
56
57class Insert(StandardInsert):
58    """PostgreSQL-specific implementation of INSERT.
59
60    Adds methods for PG-specific syntaxes such as ON CONFLICT.
61
62    The :class:`_postgresql.Insert` object is created using the
63    :func:`sqlalchemy.dialects.postgresql.insert` function.
64
65    """
66
67    stringify_dialect = "postgresql"
68    inherit_cache = False
69
70    @util.memoized_property
71    def excluded(
72        self,
73    ) -> ReadOnlyColumnCollection[str, KeyedColumnElement[Any]]:
74        """Provide the ``excluded`` namespace for an ON CONFLICT statement
75
76        PG's ON CONFLICT clause allows reference to the row that would
77        be inserted, known as ``excluded``.  This attribute provides
78        all columns in this row to be referenceable.
79
80        .. tip::  The :attr:`_postgresql.Insert.excluded` attribute is an
81            instance of :class:`_expression.ColumnCollection`, which provides
82            an interface the same as that of the :attr:`_schema.Table.c`
83            collection described at :ref:`metadata_tables_and_columns`.
84            With this collection, ordinary names are accessible like attributes
85            (e.g. ``stmt.excluded.some_column``), but special names and
86            dictionary method names should be accessed using indexed access,
87            such as ``stmt.excluded["column name"]`` or
88            ``stmt.excluded["values"]``.   See the docstring for
89            :class:`_expression.ColumnCollection` for further examples.
90
91        .. seealso::
92
93            :ref:`postgresql_insert_on_conflict` - example of how
94            to use :attr:`_expression.Insert.excluded`
95
96        """
97        return alias(self.table, name="excluded").columns
98
99    _on_conflict_exclusive = _exclusive_against(
100        "_post_values_clause",
101        msgs={
102            "_post_values_clause": "This Insert construct already has "
103            "an ON CONFLICT clause established"
104        },
105    )
106
107    @_generative
108    @_on_conflict_exclusive
109    def on_conflict_do_update(
110        self,
111        constraint: _OnConflictConstraintT = None,
112        index_elements: _OnConflictIndexElementsT = None,
113        index_where: _OnConflictIndexWhereT = None,
114        set_: _OnConflictSetT = None,
115        where: _OnConflictWhereT = None,
116    ) -> Self:
117        r"""
118        Specifies a DO UPDATE SET action for ON CONFLICT clause.
119
120        Either the ``constraint`` or ``index_elements`` argument is
121        required, but only one of these can be specified.
122
123        :param constraint:
124         The name of a unique or exclusion constraint on the table,
125         or the constraint object itself if it has a .name attribute.
126
127        :param index_elements:
128         A sequence consisting of string column names, :class:`_schema.Column`
129         objects, or other column expression objects that will be used
130         to infer a target index.
131
132        :param index_where:
133         Additional WHERE criterion that can be used to infer a
134         conditional target index.
135
136        :param set\_:
137         A dictionary or other mapping object
138         where the keys are either names of columns in the target table,
139         or :class:`_schema.Column` objects or other ORM-mapped columns
140         matching that of the target table, and expressions or literals
141         as values, specifying the ``SET`` actions to take.
142
143         .. versionadded:: 1.4 The
144            :paramref:`_postgresql.Insert.on_conflict_do_update.set_`
145            parameter supports :class:`_schema.Column` objects from the target
146            :class:`_schema.Table` as keys.
147
148         .. warning:: This dictionary does **not** take into account
149            Python-specified default UPDATE values or generation functions,
150            e.g. those specified using :paramref:`_schema.Column.onupdate`.
151            These values will not be exercised for an ON CONFLICT style of
152            UPDATE, unless they are manually specified in the
153            :paramref:`.Insert.on_conflict_do_update.set_` dictionary.
154
155        :param where:
156         Optional argument. If present, can be a literal SQL
157         string or an acceptable expression for a ``WHERE`` clause
158         that restricts the rows affected by ``DO UPDATE SET``. Rows
159         not meeting the ``WHERE`` condition will not be updated
160         (effectively a ``DO NOTHING`` for those rows).
161
162
163        .. seealso::
164
165            :ref:`postgresql_insert_on_conflict`
166
167        """
168        self._post_values_clause = OnConflictDoUpdate(
169            constraint, index_elements, index_where, set_, where
170        )
171        return self
172
173    @_generative
174    @_on_conflict_exclusive
175    def on_conflict_do_nothing(
176        self,
177        constraint: _OnConflictConstraintT = None,
178        index_elements: _OnConflictIndexElementsT = None,
179        index_where: _OnConflictIndexWhereT = None,
180    ) -> Self:
181        """
182        Specifies a DO NOTHING action for ON CONFLICT clause.
183
184        The ``constraint`` and ``index_elements`` arguments
185        are optional, but only one of these can be specified.
186
187        :param constraint:
188         The name of a unique or exclusion constraint on the table,
189         or the constraint object itself if it has a .name attribute.
190
191        :param index_elements:
192         A sequence consisting of string column names, :class:`_schema.Column`
193         objects, or other column expression objects that will be used
194         to infer a target index.
195
196        :param index_where:
197         Additional WHERE criterion that can be used to infer a
198         conditional target index.
199
200        .. seealso::
201
202            :ref:`postgresql_insert_on_conflict`
203
204        """
205        self._post_values_clause = OnConflictDoNothing(
206            constraint, index_elements, index_where
207        )
208        return self
209
210
211class OnConflictClause(ClauseElement):
212    stringify_dialect = "postgresql"
213
214    constraint_target: Optional[str]
215    inferred_target_elements: _OnConflictIndexElementsT
216    inferred_target_whereclause: _OnConflictIndexWhereT
217
218    def __init__(
219        self,
220        constraint: _OnConflictConstraintT = None,
221        index_elements: _OnConflictIndexElementsT = None,
222        index_where: _OnConflictIndexWhereT = None,
223    ):
224        if constraint is not None:
225            if not isinstance(constraint, str) and isinstance(
226                constraint,
227                (schema.Constraint, ext.ExcludeConstraint),
228            ):
229                constraint = getattr(constraint, "name") or constraint
230
231        if constraint is not None:
232            if index_elements is not None:
233                raise ValueError(
234                    "'constraint' and 'index_elements' are mutually exclusive"
235                )
236
237            if isinstance(constraint, str):
238                self.constraint_target = constraint
239                self.inferred_target_elements = None
240                self.inferred_target_whereclause = None
241            elif isinstance(constraint, schema.Index):
242                index_elements = constraint.expressions
243                index_where = constraint.dialect_options["postgresql"].get(
244                    "where"
245                )
246            elif isinstance(constraint, ext.ExcludeConstraint):
247                index_elements = constraint.columns
248                index_where = constraint.where
249            else:
250                index_elements = constraint.columns
251                index_where = constraint.dialect_options["postgresql"].get(
252                    "where"
253                )
254
255        if index_elements is not None:
256            self.constraint_target = None
257            self.inferred_target_elements = index_elements
258            self.inferred_target_whereclause = index_where
259        elif constraint is None:
260            self.constraint_target = self.inferred_target_elements = (
261                self.inferred_target_whereclause
262            ) = None
263
264
265class OnConflictDoNothing(OnConflictClause):
266    __visit_name__ = "on_conflict_do_nothing"
267
268
269class OnConflictDoUpdate(OnConflictClause):
270    __visit_name__ = "on_conflict_do_update"
271
272    def __init__(
273        self,
274        constraint: _OnConflictConstraintT = None,
275        index_elements: _OnConflictIndexElementsT = None,
276        index_where: _OnConflictIndexWhereT = None,
277        set_: _OnConflictSetT = None,
278        where: _OnConflictWhereT = None,
279    ):
280        super().__init__(
281            constraint=constraint,
282            index_elements=index_elements,
283            index_where=index_where,
284        )
285
286        if (
287            self.inferred_target_elements is None
288            and self.constraint_target is None
289        ):
290            raise ValueError(
291                "Either constraint or index_elements, "
292                "but not both, must be specified unless DO NOTHING"
293            )
294
295        if isinstance(set_, dict):
296            if not set_:
297                raise ValueError("set parameter dictionary must not be empty")
298        elif isinstance(set_, ColumnCollection):
299            set_ = dict(set_)
300        else:
301            raise ValueError(
302                "set parameter must be a non-empty dictionary "
303                "or a ColumnCollection such as the `.c.` collection "
304                "of a Table object"
305            )
306        self.update_values_to_set = [
307            (coercions.expect(roles.DMLColumnRole, key), value)
308            for key, value in set_.items()
309        ]
310        self.update_whereclause = where
311