From de11c5217b4c62f86dfd05a28689159095ab1024 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 15 May 2017 09:39:19 -0400
Subject: Add with_for_update() support in session.refresh()

Session.refresh() is still hardcoded to legacy lockmode,
come up with a new API so that the newer argument style
works with it.

Added new argument :paramref:`.with_for_update` to the
:meth:`.Session.refresh` method.  When the :meth:`.Query.with_lockmode`
method were deprecated in favor of :meth:`.Query.with_for_update`,
the :meth:`.Session.refresh` method was never updated to reflect
the new option.

Change-Id: Ia02a653746b7024699b515451525a88d7a17d63a
Fixes: #3991
---
 lib/sqlalchemy/orm/loading.py    |  7 ++++---
 lib/sqlalchemy/orm/session.py    | 31 +++++++++++++++++++++++++++++--
 lib/sqlalchemy/sql/selectable.py | 13 +++++++++++++
 3 files changed, 46 insertions(+), 5 deletions(-)

(limited to 'lib/sqlalchemy')

diff --git a/lib/sqlalchemy/orm/loading.py b/lib/sqlalchemy/orm/loading.py
index 3733d50e1..7feec660d 100644
--- a/lib/sqlalchemy/orm/loading.py
+++ b/lib/sqlalchemy/orm/loading.py
@@ -169,7 +169,7 @@ def get_from_identity(session, key, passive):
 
 
 def load_on_ident(query, key,
-                  refresh_state=None, lockmode=None,
+                  refresh_state=None, with_for_update=None,
                   only_load_props=None):
     """Load the given identity key from the database."""
 
@@ -209,9 +209,10 @@ def load_on_ident(query, key,
 
         q._params = params
 
-    if lockmode is not None:
+    # with_for_update needs to be query.LockmodeArg()
+    if with_for_update is not None:
         version_check = True
-        q = q.with_lockmode(lockmode)
+        q._for_update_arg = with_for_update
     elif query._for_update_arg is not None:
         version_check = True
         q._for_update_arg = query._for_update_arg
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index 6186ac4f7..7c313e635 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -1421,7 +1421,9 @@ class Session(_SessionClassMethods):
                     "flush is occurring prematurely")
                 util.raise_from_cause(e)
 
-    def refresh(self, instance, attribute_names=None, lockmode=None):
+    def refresh(
+            self, instance, attribute_names=None, with_for_update=None,
+            lockmode=None):
         """Expire and refresh the attributes on the given instance.
 
         A query will be issued to the database and all attributes will be
@@ -1445,8 +1447,17 @@ class Session(_SessionClassMethods):
           string attribute names indicating a subset of attributes to
           be refreshed.
 
+        :param with_for_update: optional boolean ``True`` indicating FOR UPDATE
+          should be used, or may be a dictionary containing flags to
+          indicate a more specific set of FOR UPDATE flags for the SELECT;
+          flags should match the parameters of :meth:`.Query.with_for_update`.
+          Supersedes the :paramref:`.Session.refresh.lockmode` parameter.
+
+          .. versionadded:: 1.2
+
         :param lockmode: Passed to the :class:`~sqlalchemy.orm.query.Query`
           as used by :meth:`~sqlalchemy.orm.query.Query.with_lockmode`.
+          Superseded by :paramref:`.Session.refresh.with_for_update`.
 
         .. seealso::
 
@@ -1464,10 +1475,26 @@ class Session(_SessionClassMethods):
 
         self._expire_state(state, attribute_names)
 
+        if with_for_update == {}:
+            raise sa_exc.ArgumentError(
+                "with_for_update should be the boolean value "
+                "True, or a dictionary with options.  "
+                "A blank dictionary is ambiguous.")
+
+        if lockmode:
+            with_for_update = query.LockmodeArg.parse_legacy_query(lockmode)
+        elif with_for_update is not None:
+            if with_for_update is True:
+                with_for_update = query.LockmodeArg()
+            elif with_for_update:
+                with_for_update = query.LockmodeArg(**with_for_update)
+            else:
+                with_for_update = None
+
         if loading.load_on_ident(
                 self.query(object_mapper(instance)),
                 state.key, refresh_state=state,
-                lockmode=lockmode,
+                with_for_update=with_for_update,
                 only_load_props=attribute_names) is None:
             raise sa_exc.InvalidRequestError(
                 "Could not refresh instance '%s'" %
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index a6ed6b0ce..a1e3abcb4 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -1841,6 +1841,19 @@ class ForUpdateArg(ClauseElement):
         else:
             return True
 
+    def __eq__(self, other):
+        return (
+            isinstance(other, ForUpdateArg) and
+            other.nowait == self.nowait and
+            other.read == self.read and
+            other.skip_locked == self.skip_locked and
+            other.key_share == self.key_share and
+            other.of is self.of
+        )
+
+    def __hash__(self):
+        return id(self)
+
     def _copy_internals(self, clone=_clone, **kw):
         if self.of is not None:
             self.of = [clone(col, **kw) for col in self.of]
-- 
cgit v1.2.1