summaryrefslogtreecommitdiff
path: root/lib/sqlalchemy/dialects/postgresql/dml.py
blob: bfdfbfa36498ef37146b54e59c212840ae28e889 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
# postgresql/on_conflict.py
# Copyright (C) 2005-2017 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: http://www.opensource.org/licenses/mit-license.php

from ...sql.elements import ClauseElement, _literal_as_binds
from ...sql.dml import Insert as StandardInsert
from ...sql.expression import alias
from ...sql import schema
from ...util.langhelpers import public_factory
from ...sql.base import _generative
from ... import util
from . import ext

__all__ = ('Insert', 'insert')


class Insert(StandardInsert):
    """PostgreSQL-specific implementation of INSERT.

    Adds methods for PG-specific syntaxes such as ON CONFLICT.

    .. versionadded:: 1.1

    """

    @util.memoized_property
    def excluded(self):
        """Provide the ``excluded`` namespace for an ON CONFLICT statement

        PG's ON CONFLICT clause allows reference to the row that would
        be inserted, known as ``excluded``.  This attribute provides
        all columns in this row to be referenaceable.

        .. seealso::

            :ref:`postgresql_insert_on_conflict` - example of how
            to use :attr:`.Insert.excluded`

        """
        return alias(self.table, name='excluded').columns

    @_generative
    def on_conflict_do_update(
            self,
            constraint=None, index_elements=None,
            index_where=None, set_=None, where=None):
        """
        Specifies a DO UPDATE SET action for ON CONFLICT clause.

        Either the ``constraint`` or ``index_elements`` argument is
        required, but only one of these can be specified.

        :param constraint:
        The name of a unique or exclusion constraint on the table,
        or the constraint object itself if it has a .name attribute.

        :param index_elements:
        A sequence consisting of string column names, :class:`.Column`
        objects, or other column expression objects that will be used
        to infer a target index.

        :param index_where:
        Additional WHERE criterion that can be used to infer a
        conditional target index.

        :param set_:
        Required argument. A dictionary or other mapping object
        with column names as keys and expressions or literals as values,
        specifying the ``SET`` actions to take.
        If the target :class:`.Column` specifies a ".key" attribute distinct
        from the column name, that key should be used.

        .. warning:: This dictionary does **not** take into account
           Python-specified default UPDATE values or generation functions,
           e.g. those specified using :paramref:`.Column.onupdate`.
           These values will not be exercised for an ON CONFLICT style of
           UPDATE, unless they are manually specified in the
           :paramref:`.Insert.on_conflict_do_update.set_` dictionary.

        :param where:
        Optional argument. If present, can be a literal SQL
        string or an acceptable expression for a ``WHERE`` clause
        that restricts the rows affected by ``DO UPDATE SET``. Rows
        not meeting the ``WHERE`` condition will not be updated
        (effectively a ``DO NOTHING`` for those rows).

        .. versionadded:: 1.1

        .. seealso::

            :ref:`postgresql_insert_on_conflict`

        """
        self._post_values_clause = OnConflictDoUpdate(
            constraint, index_elements, index_where, set_, where)
        return self

    @_generative
    def on_conflict_do_nothing(
            self,
            constraint=None, index_elements=None, index_where=None):
        """
        Specifies a DO NOTHING action for ON CONFLICT clause.

        The ``constraint`` and ``index_elements`` arguments
        are optional, but only one of these can be specified.

        :param constraint:
        The name of a unique or exclusion constraint on the table,
        or the constraint object itself if it has a .name attribute.

        :param index_elements:
        A sequence consisting of string column names, :class:`.Column`
        objects, or other column expression objects that will be used
        to infer a target index.

        :param index_where:
        Additional WHERE criterion that can be used to infer a
        conditional target index.

        .. versionadded:: 1.1

        .. seealso::

            :ref:`postgresql_insert_on_conflict`

        """
        self._post_values_clause = OnConflictDoNothing(
            constraint, index_elements, index_where)
        return self

insert = public_factory(Insert, '.dialects.postgresql.insert')


class OnConflictClause(ClauseElement):
    def __init__(
            self,
            constraint=None,
            index_elements=None,
            index_where=None):

        if constraint is not None:
            if not isinstance(constraint, util.string_types) and \
                    isinstance(constraint, (
                        schema.Index, schema.Constraint,
                        ext.ExcludeConstraint)):
                constraint = getattr(constraint, 'name') or constraint

        if constraint is not None:
            if index_elements is not None:
                raise ValueError(
                    "'constraint' and 'index_elements' are mutually exclusive")

            if isinstance(constraint, util.string_types):
                self.constraint_target = constraint
                self.inferred_target_elements = None
                self.inferred_target_whereclause = None
            elif isinstance(constraint, schema.Index):
                index_elements = constraint.expressions
                index_where = \
                    constraint.dialect_options['postgresql'].get("where")
            elif isinstance(constraint, ext.ExcludeConstraint):
                index_elements = constraint.columns
                index_where = constraint.where
            else:
                index_elements = constraint.columns
                index_where = \
                    constraint.dialect_options['postgresql'].get("where")

        if index_elements is not None:
            self.constraint_target = None
            self.inferred_target_elements = index_elements
            self.inferred_target_whereclause = index_where
        elif constraint is None:
            self.constraint_target = self.inferred_target_elements = \
                self.inferred_target_whereclause = None


class OnConflictDoNothing(OnConflictClause):
    __visit_name__ = 'on_conflict_do_nothing'


class OnConflictDoUpdate(OnConflictClause):
    __visit_name__ = 'on_conflict_do_update'

    def __init__(
            self,
            constraint=None,
            index_elements=None,
            index_where=None,
            set_=None,
            where=None):
        super(OnConflictDoUpdate, self).__init__(
            constraint=constraint,
            index_elements=index_elements,
            index_where=index_where)

        if self.inferred_target_elements is None and \
                self.constraint_target is None:
            raise ValueError(
                "Either constraint or index_elements, "
                "but not both, must be specified unless DO NOTHING")

        if (not isinstance(set_, dict) or not set_):
            raise ValueError("set parameter must be a non-empty dictionary")
        self.update_values_to_set = [
            (key, value)
            for key, value in set_.items()
        ]
        self.update_whereclause = where