From 8860117c9655a4bdeafebab1c6ef12c6a6198e66 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sat, 16 Jan 2021 12:39:51 -0500
Subject: introduce generalized decorator to prevent invalid method calls

This introduces the ``_exclusive_against()`` utility decorator
that can be used to prevent repeated invocations of methods that
typically should only be called once.

An informative error message is now raised for a selected set of DML
methods (currently all part of :class:`_dml.Insert` constructs) if they are
called a second time, which would implicitly cancel out the previous
setting.  The methods altered include:
:class:`_sqlite.Insert.on_conflict_do_update`,
:class:`_sqlite.Insert.on_conflict_do_nothing` (SQLite),
:class:`_postgresql.Insert.on_conflict_do_update`,
:class:`_postgresql.Insert.on_conflict_do_nothing` (PostgreSQL),
:class:`_mysql.Insert.on_duplicate_key_update` (MySQL)

Fixes: #5169
Change-Id: I9278fa87cd3470dcf296ff96bb0fb17a3236d49d
---
 lib/sqlalchemy/sql/base.py | 25 +++++++++++++++++++++++++
 lib/sqlalchemy/sql/dml.py  | 30 +++++++++++++++++-------------
 2 files changed, 42 insertions(+), 13 deletions(-)

(limited to 'lib/sqlalchemy/sql')

diff --git a/lib/sqlalchemy/sql/base.py b/lib/sqlalchemy/sql/base.py
index 550111020..220bbb115 100644
--- a/lib/sqlalchemy/sql/base.py
+++ b/lib/sqlalchemy/sql/base.py
@@ -102,6 +102,31 @@ def _generative(fn):
     return decorated
 
 
+def _exclusive_against(*names, **kw):
+    msgs = kw.pop("msgs", {})
+
+    defaults = kw.pop("defaults", {})
+
+    getters = [
+        (name, operator.attrgetter(name), defaults.get(name, None))
+        for name in names
+    ]
+
+    @util.decorator
+    def check(fn, self, *args, **kw):
+        for name, getter, default_ in getters:
+            if getter(self) is not default_:
+                msg = msgs.get(
+                    name,
+                    "Method %s() has already been invoked on this %s construct"
+                    % (fn.__name__, self.__class__),
+                )
+                raise exc.InvalidRequestError(msg)
+        return fn(self, *args, **kw)
+
+    return check
+
+
 def _clone(element, **kw):
     return element._clone()
 
diff --git a/lib/sqlalchemy/sql/dml.py b/lib/sqlalchemy/sql/dml.py
index c402de121..3f492a490 100644
--- a/lib/sqlalchemy/sql/dml.py
+++ b/lib/sqlalchemy/sql/dml.py
@@ -14,6 +14,7 @@ from . import coercions
 from . import roles
 from . import util as sql_util
 from .base import _entity_namespace_key
+from .base import _exclusive_against
 from .base import _from_objects
 from .base import _generative
 from .base import ColumnCollection
@@ -495,6 +496,15 @@ class ValuesBase(UpdateBase):
             self._setup_prefixes(prefixes)
 
     @_generative
+    @_exclusive_against(
+        "_select_names",
+        "_ordered_values",
+        msgs={
+            "_select_names": "This construct already inserts from a SELECT",
+            "_ordered_values": "This statement already has ordered "
+            "values present",
+        },
+    )
     def values(self, *args, **kwargs):
         r"""Specify a fixed VALUES clause for an INSERT statement, or the SET
         clause for an UPDATE.
@@ -607,15 +617,6 @@ class ValuesBase(UpdateBase):
 
 
         """
-        if self._select_names:
-            raise exc.InvalidRequestError(
-                "This construct already inserts from a SELECT"
-            )
-        elif self._ordered_values:
-            raise exc.ArgumentError(
-                "This statement already has ordered values present"
-            )
-
         if args:
             # positional case.  this is currently expensive.   we don't
             # yet have positional-only args so we have to check the length.
@@ -699,6 +700,13 @@ class ValuesBase(UpdateBase):
                 self._values = util.immutabledict(arg)
 
     @_generative
+    @_exclusive_against(
+        "_returning",
+        msgs={
+            "_returning": "RETURNING is already configured on this statement"
+        },
+        defaults={"_returning": _returning},
+    )
     def return_defaults(self, *cols):
         """Make use of a :term:`RETURNING` clause for the purpose
         of fetching server-side expressions and defaults.
@@ -783,10 +791,6 @@ class ValuesBase(UpdateBase):
             :attr:`_engine.CursorResult.inserted_primary_key_rows`
 
         """
-        if self._returning:
-            raise exc.InvalidRequestError(
-                "RETURNING is already configured on this statement"
-            )
         self._return_defaults = cols or True
 
 
-- 
cgit v1.2.1