- add support for specifying tables or entities for "of"

- implement Query with_for_update()
- rework docs and tests

4 files changed, 114 insertions, 67 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 69b0fb040..7ad266b58 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1021,8 +1021,9 @@ class PGCompiler(compiler.SQLCompiler):
             tmp = " FOR UPDATE"
 
         if select._for_update_arg.of:
-            # TODO: assuming simplistic c.table here
-            tables = set(c.table for c in select._for_update_arg.of)
+            tables = util.OrderedSet(
+                            c.table if isinstance(c, expression.ColumnClause)
+                            else c for c in select._for_update_arg.of)
             tmp += " OF " + ", ".join(
                                 self.process(table, ashint=True)
                                 for table in tables
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 173ad038e..14e8c31be 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -37,7 +37,6 @@ from ..sql import (
         expression, visitors
     )
 from ..sql.base import ColumnCollection
-from ..sql import operators
 from . import properties
 
 __all__ = ['Query', 'QueryContext', 'aliased']
@@ -69,7 +68,6 @@ class Query(object):
     _with_labels = False
     _criterion = None
     _yield_per = None
-    _lockmode = None
     _order_by = False
     _group_by = False
     _having = None
@@ -77,6 +75,7 @@ class Query(object):
     _prefixes = None
     _offset = None
     _limit = None
+    _for_update_arg = None
     _statement = None
     _correlate = frozenset()
     _populate_existing = False
@@ -797,7 +796,7 @@ class Query(object):
 
         if not self._populate_existing and \
                 not mapper.always_refresh and \
-                self._lockmode is None:
+                self._for_update_arg is None:
 
             instance = loading.get_from_identity(
                 self.session, key, attributes.PASSIVE_OFF)
@@ -1125,43 +1124,63 @@ class Query(object):
 
     @_generative()
     def with_lockmode(self, mode):
-        """Return a new Query object with the specified locking mode.
+        """Return a new :class:`.Query` object with the specified "locking mode",
+        which essentially refers to the ``FOR UPDATE`` clause.
 
         .. deprecated:: 0.9.0b2 superseded by :meth:`.Query.with_for_update`.
 
-        :param mode: a string representing the desired locking mode. A
-            corresponding :meth:`~sqlalchemy.orm.query.LockmodeArgs` object
-            is passed to the ``for_update`` parameter of
-            :meth:`~sqlalchemy.sql.expression.select` when the
-            query is executed. Valid values are:
+        :param mode: a string representing the desired locking mode.
+         Valid values are:
 
-            ``None`` - translates to no lockmode
+         * ``None`` - translates to no lockmode
 
-            ``'update'`` - translates to ``FOR UPDATE``
-            (standard SQL, supported by most dialects)
+         * ``'update'`` - translates to ``FOR UPDATE``
+           (standard SQL, supported by most dialects)
 
-            ``'update_nowait'`` - translates to ``FOR UPDATE NOWAIT``
-            (supported by Oracle, PostgreSQL 8.1 upwards)
+         * ``'update_nowait'`` - translates to ``FOR UPDATE NOWAIT``
+           (supported by Oracle, PostgreSQL 8.1 upwards)
 
-            ``'read'`` - translates to ``LOCK IN SHARE MODE`` (for MySQL),
-            and ``FOR SHARE`` (for PostgreSQL)
+         * ``'read'`` - translates to ``LOCK IN SHARE MODE`` (for MySQL),
+           and ``FOR SHARE`` (for PostgreSQL)
 
-            .. versionadded:: 0.7.7
-                ``FOR SHARE`` and ``FOR SHARE NOWAIT`` (PostgreSQL).
+        .. seealso::
 
-         :param of: either a column descriptor, or list of column
-            descriptors, representing the optional OF part of the
-            clause. This passes the descriptor to the
-            corresponding :meth:`~sqlalchemy.orm.query.LockmodeArgs` object,
-            and translates to ``FOR UPDATE OF table [NOWAIT]`` respectively
-            ``FOR UPDATE OF table, table [NOWAIT]`` (PostgreSQL), or
-            ``FOR UPDATE OF table.column [NOWAIT]`` respectively
-            ``FOR UPDATE OF table.column, table.column [NOWAIT]`` (Oracle).
+            :meth:`.Query.with_for_update` - improved API for
+            specifying the ``FOR UPDATE`` clause.
 
-            .. versionadded:: 0.9.0b2
         """
+        self._for_update_arg = LockmodeArg.parse_legacy_query(mode)
+
+    @_generative()
+    def with_for_update(self, read=False, nowait=False, of=None):
+        """return a new :class:`.Query` with the specified options for the
+        ``FOR UPDATE`` clause.
+
+        The behavior of this method is identical to that of
+        :meth:`.SelectBase.with_for_update`.  When called with no arguments,
+        the resulting ``SELECT`` statement will have a ``FOR UPDATE`` clause
+        appended.  When additional arguments are specified, backend-specific
+        options such as ``FOR UPDATE NOWAIT`` or ``LOCK IN SHARE MODE``
+        can take effect.
+
+        E.g.::
+
+            q = sess.query(User).with_for_update(nowait=True, of=User)
+
+        The above query on a Postgresql backend will render like::
+
+            SELECT users.id AS users_id FROM users FOR UPDATE OF users NOWAIT
+
+        .. versionadded:: 0.9.0b2 :meth:`.Query.with_for_update` supersedes
+           the :meth:`.Query.with_lockmode` method.
+
+        .. seealso::
+
+            :meth:`.SelectBase.with_for_update` - Core level method with
+            full argument and behavioral description.
 
-        self._lockmode = LockmodeArgs(mode=mode, of=of)
+        """
+        self._for_update_arg = LockmodeArg(read=read, nowait=nowait, of=of)
 
     @_generative()
     def params(self, *args, **kwargs):
@@ -2703,12 +2722,7 @@ class Query(object):
 
         context.labels = labels
 
-        if isinstance(self._lockmode, bool) and self._lockmode:
-            context.for_update = LockmodeArgs(mode='update')
-        elif isinstance(self._lockmode, LockmodeArgs):
-            if self._lockmode.mode not in LockmodeArgs.lockmodes:
-                raise sa_exc.ArgumentError('Unknown lockmode %r' % self._lockmode.mode)
-            context.for_update = self._lockmode
+        context._for_update_arg = self._for_update_arg
 
         for entity in self._entities:
             entity.setup_context(self, context)
@@ -2793,9 +2807,10 @@ class Query(object):
 
         statement = sql.select(
                             [inner] + context.secondary_columns,
-                            for_update=context.for_update,
                             use_labels=context.labels)
 
+        statement._for_update_arg = context._for_update_arg
+
         from_clause = inner
         for eager_join in context.eager_joins.values():
             # EagerLoader places a 'stop_on' attribute on the join,
@@ -2838,11 +2853,12 @@ class Query(object):
                         context.whereclause,
                         from_obj=context.froms,
                         use_labels=context.labels,
-                        for_update=context.for_update,
                         order_by=context.order_by,
                         **self._select_args
                     )
 
+        statement._for_update_arg = context._for_update_arg
+
         for hint in self._with_hints:
             statement = statement.with_hint(*hint)
 
@@ -2877,6 +2893,27 @@ class Query(object):
     def __str__(self):
         return str(self._compile_context().statement)
 
+from ..sql.selectable import ForUpdateArg
+
+class LockmodeArg(ForUpdateArg):
+    @classmethod
+    def parse_legacy_query(self, mode):
+        if mode in (None, False):
+            return None
+
+        if mode == "read":
+            read = True
+            nowait = False
+        elif mode == "update":
+            read = nowait = False
+        elif mode == "update_nowait":
+            nowait = True
+            read = False
+        else:
+            raise sa_exc.ArgumentError(
+                        "Unknown with_lockmode argument: %r" % mode)
+
+        return LockmodeArg(read=read, nowait=nowait)
 
 class _QueryEntity(object):
     """represent an entity column returned within a Query result."""
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index 01c803f3b..28c757a66 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -1175,7 +1175,7 @@ class ForUpdateArg(ClauseElement):
             ``FOR SHARE NOWAIT`` (PostgreSQL).
 
         """
-        if arg is None:
+        if arg in (None, False):
             return None
 
         nowait = read = False
@@ -1214,7 +1214,7 @@ class ForUpdateArg(ClauseElement):
         self.nowait = nowait
         self.read = read
         if of is not None:
-            self.of = [_only_column_elements(elem, "of")
+            self.of = [_interpret_as_column_or_from(elem)
                         for elem in util.to_list(of)]
         else:
             self.of = None
@@ -1262,24 +1262,38 @@ class SelectBase(Executable, FromClause):
 
     @property
     def for_update(self):
-        """Provide legacy dialect support for the ``for_update`` attribute
-           as a getter.
-
+        """Provide legacy dialect support for the ``for_update`` attribute.
         """
         if self._for_update_arg is not None:
             return self._for_update_arg.legacy_for_update_value
         else:
             return None
 
+    @for_update.setter
+    def for_update(self, value):
+        self._for_update_arg = ForUpdateArg.parse_legacy_select(value)
+
     @_generative
     def with_for_update(self, nowait=False, read=False, of=None):
-        """apply FOR UPDATE to this :class:`.SelectBase`.
+        """Specify a ``FOR UPDATE`` clause for this :class:`.SelectBase`.
 
         E.g.::
 
             stmt = select([table]).with_for_update(nowait=True)
 
-        Additional keyword arguments are provided for common database-specific
+        On a database like Postgresql or Oracle, the above would render a
+        statement like::
+
+            SELECT table.a, table.b FROM table FOR UPDATE NOWAIT
+
+        on other backends, the ``nowait`` option is ignored and instead
+        would produce::
+
+            SELECT table.a, table.b FROM table FOR UPDATE
+
+        When called with no arguments, the statement will render with
+        the suffix ``FOR UPDATE``.   Additional arguments can then be
+        provided which allow for common database-specific
         variants.
 
         :param nowait: boolean; will render ``FOR UPDATE NOWAIT`` on Oracle and
@@ -1289,12 +1303,12 @@ class SelectBase(Executable, FromClause):
          ``FOR SHARE`` on Postgresql.  On Postgresql, when combined with
          ``nowait``, will render ``FOR SHARE NOWAIT``.
 
-        :param of: SQL expression or list of SQL expression elements which
+        :param of: SQL expression or list of SQL expression elements
+         (typically :class:`.Column` objects or a compatible expression) which
          will render into a ``FOR UPDATE OF`` clause; supported by PostgreSQL
-         and Oracle.    May render as a table or as a column depending on
+         and Oracle.  May render as a table or as a column depending on
          backend.
 
-
         .. versionadded:: 0.9.0b2
 
         """
@@ -1943,27 +1957,22 @@ class Select(HasPrefixes, SelectBase):
           resulting statement.
 
           .. deprecated:: 0.9.0 - use :meth:`.SelectBase.with_for_update`
-             to specify for update arguments.
+             to specify the structure of the ``FOR UPDATE`` clause.
 
-          Additional values are accepted here, including:
+          ``for_update`` accepts various string values interpreted by
+          specific backends, including:
 
-            ``None`` - translates to no lockmode
+          * ``"read"`` - on MySQL, translates to ``LOCK IN SHARE MODE``;
+            on Postgresql, translates to ``FOR SHARE``.
+          * ``"nowait"`` - on Postgresql and Oracle, translates to
+            ``FOR UPDATE NOWAIT``.
+          * ``"read_nowait"`` - on Postgresql, translates to
+            ``FOR SHARE NOWAIT``.
 
-            ``'update'`` - translates to ``FOR UPDATE``
-            (standard SQL, supported by most dialects)
-
-            ``'update_nowait'`` - translates to ``FOR UPDATE NOWAIT``
-            (supported by Oracle, PostgreSQL 8.1 upwards)
-
-            ``'read'`` - translates to ``LOCK IN SHARE MODE`` (for MySQL),
-            and ``FOR SHARE`` (for PostgreSQL)
-
-            ``'read_nowait'`` - translates to ``FOR SHARE NOWAIT``
-            (supported by PostgreSQL). ``FOR SHARE`` and
-            ``FOR SHARE NOWAIT`` (PostgreSQL).
+         .. seealso::
 
-          The :meth:`.SelectBase.with_for_update` method should be preferred as
-          a means to specify FOR UPDATE more simply.
+            :meth:`.SelectBase.with_for_update` - improved API for
+            specifying the ``FOR UPDATE`` clause.
 
         :param group_by:
           a list of :class:`.ClauseElement` objects which will comprise the
diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py
index ce3a292ac..0e12963ce 100644
--- a/lib/sqlalchemy/testing/assertions.py
+++ b/lib/sqlalchemy/testing/assertions.py
@@ -8,7 +8,7 @@ from __future__ import absolute_import
 
 from . import util as testutil
 from sqlalchemy import pool, orm, util
-from sqlalchemy.engine import default, create_engine
+from sqlalchemy.engine import default, create_engine, url
 from sqlalchemy import exc as sa_exc
 from sqlalchemy.util import decorator
 from sqlalchemy import types as sqltypes, schema
@@ -208,7 +208,7 @@ class AssertsCompiledSQL(object):
             elif dialect == 'default':
                 dialect = default.DefaultDialect()
             elif isinstance(dialect, util.string_types):
-                dialect = create_engine("%s://" % dialect).dialect
+                dialect = url.URL(dialect).get_dialect()()
 
 
         kw = {}