Propose using RETURNING for bulk updates, deletes

This patch makes several improvements in the area of
bulk updates and deletes as well as the new session mechanics.

RETURNING is now used for an UPDATE or DELETE statement
emitted for a diaelct that supports "full returning"
in order to satisfy the "fetch" strategy; this currently
includes PostgreSQL and SQL Server.  The Oracle dialect
does not support RETURNING for more than one row,
so a new dialect capability "full_returning" is added
in addition to the existing "implicit_returning", indicating
this dialect supports RETURNING for zero or more rows,
not just a single identity row.

The "fetch" strategy will gracefully degrade to
the previous SELECT mechanics for dialects that do not
support RETURNING.

Additionally, the "fetch" strategy will attempt to use
evaluation for the VALUES that were UPDATEd, rather
than just expiring the updated attributes.   Values should
be evalutable in all cases where the value is not
a SQL expression.

The new approach also incurs some changes in the
session.execute mechanics, where do_orm_execute() event
handlers can now be chained to each return results;
this is in turn used by the handler to detect on a
per-bind basis if the fetch strategy needs to
do a SELECT or if it can do RETURNING.  A test suite is
added to test_horizontal_shard that breaks up a single
UPDATE or DELETE operation among multiple backends
where some are SQLite and don't support RETURNING and
others are PostgreSQL and do.

The session event mechanics are corrected
in terms of the "orm pre execute" hook, which now
receives a flag "is_reentrant" so that the two
ORM implementations for this can skip on their work
if they are being called inside of ORMExecuteState.invoke(),
where previously bulk update/delete were calling its
SELECT a second time.

In order for "fetch" to get the correct identity when
called as pre-execute, it also requests the identity_token
for each mapped instance which is now added as an optional
capability of a SELECT for ORM columns.   the identity_token
that's placed by horizontal_sharding is now made available
within each result row, so that even when fetching a
merged result of plain rows we can tell which row belongs
to which identity token.

The evaluator that takes place within the ORM bulk update and delete for
synchronize_session="evaluate" now supports the IN and NOT IN operators.
Tuple IN is also supported.

Fixes: #1653

Change-Id: I2292b56ae004b997cef0ba4d3fc350ae1dd5efc1

6 files changed, 257 insertions, 20 deletions

diff --git a/test/aaa_profiling/test_orm.py b/test/aaa_profiling/test_orm.py
index 8f06220e2..13e92f5c4 100644
--- a/test/aaa_profiling/test_orm.py
+++ b/test/aaa_profiling/test_orm.py
@@ -878,6 +878,7 @@ class JoinedEagerLoadTest(NoCache, fixtures.MappedTest):
                     {},
                     exec_opts,
                     bind_arguments,
+                    is_reentrant_invoke=False,
                 )
 
                 r = sess.connection().execute(
@@ -888,7 +889,7 @@ class JoinedEagerLoadTest(NoCache, fixtures.MappedTest):
 
                 r.context.compiled.compile_state = compile_state
                 obj = ORMCompileState.orm_setup_cursor_result(
-                    sess, compile_state.statement, exec_opts, {}, r
+                    sess, compile_state.statement, exec_opts, {}, r,
                 )
                 list(obj)
                 sess.close()
diff --git a/test/engine/test_reflection.py b/test/engine/test_reflection.py
index 0fea029fe..f1b54cb8f 100644
--- a/test/engine/test_reflection.py
+++ b/test/engine/test_reflection.py
@@ -1197,7 +1197,6 @@ class ReflectionTest(fixtures.TestBase, ComparesTables):
     @testing.provide_metadata
     def test_reflect_all(self):
         existing = inspect(testing.db).get_table_names()
-
         names = ["rt_%s" % name for name in ("a", "b", "c", "d", "e")]
         nameset = set(names)
         for name in names:
diff --git a/test/ext/test_horizontal_shard.py b/test/ext/test_horizontal_shard.py
index c0029fbb6..9855cd5ab 100644
--- a/test/ext/test_horizontal_shard.py
+++ b/test/ext/test_horizontal_shard.py
@@ -35,8 +35,6 @@ from sqlalchemy.testing import provision
 from sqlalchemy.testing.engines import testing_engine
 from sqlalchemy.testing.engines import testing_reaper
 
-# TODO: ShardTest can be turned into a base for further subclasses
-
 
 class ShardTest(object):
     __skip_if__ = (lambda: util.win32,)
@@ -47,9 +45,9 @@ class ShardTest(object):
     def setUp(self):
         global db1, db2, db3, db4, weather_locations, weather_reports
 
-        db1, db2, db3, db4 = self._init_dbs()
+        db1, db2, db3, db4 = self._dbs = self._init_dbs()
 
-        meta = MetaData()
+        meta = self.metadata = MetaData()
         ids = Table("ids", meta, Column("nextid", Integer, nullable=False))
 
         def id_generator(ctx):
@@ -578,9 +576,11 @@ class ShardTest(object):
         temps = sess.execute(future_select(Report)).scalars().all()
         eq_(set(t.temperature for t in temps), {80.0, 75.0, 85.0})
 
+        # MARKMARK
+        # omitting the criteria so that the UPDATE affects three out of
+        # four shards
         sess.execute(
             update(Report)
-            .filter(Report.temperature >= 80)
             .values({"temperature": Report.temperature + 6},)
             .execution_options(synchronize_session="fetch")
         )
@@ -590,11 +590,11 @@ class ShardTest(object):
                 row.temperature
                 for row in sess.execute(future_select(Report.temperature))
             ),
-            {86.0, 75.0, 91.0},
+            {86.0, 81.0, 91.0},
         )
 
         # test synchronize session as well
-        eq_(set(t.temperature for t in temps), {86.0, 75.0, 91.0})
+        eq_(set(t.temperature for t in temps), {86.0, 81.0, 91.0})
 
     def test_bulk_delete_future_synchronize_evaluate(self):
         sess = self._fixture_data()
@@ -711,9 +711,8 @@ class TableNameConventionShardTest(ShardTest, fixtures.TestBase):
     This used to be called "AttachedFileShardTest" but I didn't see any
     ATTACH going on.
 
-    The approach taken by this test is awkward and I wouldn't recommend  using
-    this pattern in a real situation.  I'm not sure of the history of this test
-    but it likely predates when we knew how to use real ATTACH in SQLite.
+    A more modern approach here would be to use the schema_translate_map
+    option.
 
     """
 
@@ -742,6 +741,49 @@ class TableNameConventionShardTest(ShardTest, fixtures.TestBase):
         return db1, db2, db3, db4
 
 
+class MultipleDialectShardTest(ShardTest, fixtures.TestBase):
+    __only_on__ = "postgresql"
+
+    schema = "changeme"
+
+    def _init_dbs(self):
+        e1 = testing_engine("sqlite://")
+        with e1.connect() as conn:
+            for i in [1, 3]:
+                conn.exec_driver_sql(
+                    'ATTACH DATABASE "shard%s_%s.db" AS shard%s'
+                    % (i, provision.FOLLOWER_IDENT, i)
+                )
+
+        e2 = testing_engine()
+        with e2.connect() as conn:
+            for i in [2, 4]:
+                conn.exec_driver_sql(
+                    "CREATE SCHEMA IF NOT EXISTS shard%s" % (i,)
+                )
+
+        db1 = e1.execution_options(schema_translate_map={"changeme": "shard1"})
+        db2 = e2.execution_options(schema_translate_map={"changeme": "shard2"})
+        db3 = e1.execution_options(schema_translate_map={"changeme": "shard3"})
+        db4 = e2.execution_options(schema_translate_map={"changeme": "shard4"})
+
+        self.sqlite_engine = e1
+        self.postgresql_engine = e2
+        return db1, db2, db3, db4
+
+    def teardown(self):
+        clear_mappers()
+
+        self.sqlite_engine.connect().invalidate()
+        for i in [1, 3]:
+            os.remove("shard%d_%s.db" % (i, provision.FOLLOWER_IDENT))
+
+        with self.postgresql_engine.connect() as conn:
+            self.metadata.drop_all(conn)
+            for i in [2, 4]:
+                conn.exec_driver_sql("DROP SCHEMA shard%s CASCADE" % (i,))
+
+
 class SelectinloadRegressionTest(fixtures.DeclarativeMappedTest):
     """test #4175
     """
diff --git a/test/orm/test_evaluator.py b/test/orm/test_evaluator.py
index 5bc054486..20577d8e6 100644
--- a/test/orm/test_evaluator.py
+++ b/test/orm/test_evaluator.py
@@ -8,6 +8,7 @@ from sqlalchemy import Integer
 from sqlalchemy import not_
 from sqlalchemy import or_
 from sqlalchemy import String
+from sqlalchemy import tuple_
 from sqlalchemy.orm import evaluator
 from sqlalchemy.orm import mapper
 from sqlalchemy.orm import relationship
@@ -19,7 +20,6 @@ from sqlalchemy.testing import is_
 from sqlalchemy.testing.schema import Column
 from sqlalchemy.testing.schema import Table
 
-
 compiler = evaluator.EvaluatorCompiler()
 
 
@@ -191,6 +191,54 @@ class EvaluateTest(fixtures.MappedTest):
             ],
         )
 
+    def test_in(self):
+        User = self.classes.User
+
+        eval_eq(
+            User.name.in_(["foo", "bar"]),
+            testcases=[
+                (User(id=1, name="foo"), True),
+                (User(id=2, name="bat"), False),
+                (User(id=1, name="bar"), True),
+                (User(id=1, name=None), None),
+            ],
+        )
+
+        eval_eq(
+            User.name.notin_(["foo", "bar"]),
+            testcases=[
+                (User(id=1, name="foo"), False),
+                (User(id=2, name="bat"), True),
+                (User(id=1, name="bar"), False),
+                (User(id=1, name=None), None),
+            ],
+        )
+
+    def test_in_tuples(self):
+        User = self.classes.User
+
+        eval_eq(
+            tuple_(User.id, User.name).in_([(1, "foo"), (2, "bar")]),
+            testcases=[
+                (User(id=1, name="foo"), True),
+                (User(id=2, name="bat"), False),
+                (User(id=1, name="bar"), False),
+                (User(id=2, name="bar"), True),
+                (User(id=1, name=None), None),
+            ],
+        )
+
+        eval_eq(
+            tuple_(User.id, User.name).notin_([(1, "foo"), (2, "bar")]),
+            testcases=[
+                (User(id=1, name="foo"), False),
+                (User(id=2, name="bat"), True),
+                (User(id=1, name="bar"), True),
+                (User(id=2, name="bar"), False),
+                (User(id=1, name=None), None),
+            ],
+        )
+
     def test_null_propagation(self):
         User = self.classes.User
 
diff --git a/test/orm/test_update_delete.py b/test/orm/test_update_delete.py
index 12a8417ba..310b17047 100644
--- a/test/orm/test_update_delete.py
+++ b/test/orm/test_update_delete.py
@@ -23,6 +23,9 @@ from sqlalchemy.testing import assert_raises
 from sqlalchemy.testing import assert_raises_message
 from sqlalchemy.testing import eq_
 from sqlalchemy.testing import fixtures
+from sqlalchemy.testing import in_
+from sqlalchemy.testing import not_in_
+from sqlalchemy.testing.assertsql import CompiledSQL
 from sqlalchemy.testing.schema import Column
 from sqlalchemy.testing.schema import Table
 
@@ -144,6 +147,50 @@ class UpdateDeleteTest(fixtures.MappedTest):
                 q.delete,
             )
 
+    def test_update_w_unevaluatable_value_evaluate(self):
+        """test that the "evaluate" strategy falls back to 'expire' for an
+        update SET that is not evaluable in Python."""
+
+        User = self.classes.User
+
+        s = Session()
+
+        jill = s.query(User).filter(User.name == "jill").one()
+
+        s.execute(
+            update(User)
+            .filter(User.name == "jill")
+            .values({"name": User.name + User.name}),
+            execution_options={"synchronize_session": "evaluate"},
+        )
+
+        eq_(jill.name, "jilljill")
+
+    def test_update_w_unevaluatable_value_fetch(self):
+        """test that the "fetch" strategy falls back to 'expire' for an
+        update SET that is not evaluable in Python.
+
+        Prior to 1.4 the "fetch" strategy used expire for everything
+        but now it tries to evaluate a SET clause to avoid a round
+        trip.
+
+        """
+
+        User = self.classes.User
+
+        s = Session()
+
+        jill = s.query(User).filter(User.name == "jill").one()
+
+        s.execute(
+            update(User)
+            .filter(User.name == "jill")
+            .values({"name": User.name + User.name}),
+            execution_options={"synchronize_session": "fetch"},
+        )
+
+        eq_(jill.name, "jilljill")
+
     def test_evaluate_clauseelement(self):
         User = self.classes.User
 
@@ -479,6 +526,87 @@ class UpdateDeleteTest(fixtures.MappedTest):
             list(zip([25, 37, 29, 27])),
         )
 
+    def test_update_fetch_returning(self):
+        User = self.classes.User
+
+        sess = Session()
+
+        john, jack, jill, jane = sess.query(User).order_by(User.id).all()
+
+        with self.sql_execution_asserter() as asserter:
+            sess.query(User).filter(User.age > 29).update(
+                {"age": User.age - 10}, synchronize_session="fetch"
+            )
+
+            # these are simple values, these are now evaluated even with
+            # the "fetch" strategy, new in 1.4, so there is no expiry
+            eq_([john.age, jack.age, jill.age, jane.age], [25, 37, 29, 27])
+
+        if testing.db.dialect.full_returning:
+            asserter.assert_(
+                CompiledSQL(
+                    "UPDATE users SET age_int=(users.age_int - %(age_int_1)s) "
+                    "WHERE users.age_int > %(age_int_2)s RETURNING users.id",
+                    [{"age_int_1": 10, "age_int_2": 29}],
+                    dialect="postgresql",
+                ),
+            )
+        else:
+            asserter.assert_(
+                CompiledSQL(
+                    "SELECT users.id FROM users "
+                    "WHERE users.age_int > :age_int_1",
+                    [{"age_int_1": 29}],
+                ),
+                CompiledSQL(
+                    "UPDATE users SET age_int=(users.age_int - :age_int_1) "
+                    "WHERE users.age_int > :age_int_2",
+                    [{"age_int_1": 10, "age_int_2": 29}],
+                ),
+            )
+
+    def test_delete_fetch_returning(self):
+        User = self.classes.User
+
+        sess = Session()
+
+        john, jack, jill, jane = sess.query(User).order_by(User.id).all()
+
+        in_(john, sess)
+        in_(jack, sess)
+
+        with self.sql_execution_asserter() as asserter:
+            sess.query(User).filter(User.age > 29).delete(
+                synchronize_session="fetch"
+            )
+
+        if testing.db.dialect.full_returning:
+            asserter.assert_(
+                CompiledSQL(
+                    "DELETE FROM users WHERE users.age_int > %(age_int_1)s "
+                    "RETURNING users.id",
+                    [{"age_int_1": 29}],
+                    dialect="postgresql",
+                ),
+            )
+        else:
+            asserter.assert_(
+                CompiledSQL(
+                    "SELECT users.id FROM users "
+                    "WHERE users.age_int > :age_int_1",
+                    [{"age_int_1": 29}],
+                ),
+                CompiledSQL(
+                    "DELETE FROM users WHERE users.age_int > :age_int_1",
+                    [{"age_int_1": 29}],
+                ),
+            )
+
+        in_(john, sess)
+        not_in_(jack, sess)
+        in_(jill, sess)
+        not_in_(jane, sess)
+
     def test_update_without_load(self):
         User = self.classes.User
 
diff --git a/test/sql/test_returning.py b/test/sql/test_returning.py
index f856c15a4..90c21ed45 100644
--- a/test/sql/test_returning.py
+++ b/test/sql/test_returning.py
@@ -129,6 +129,32 @@ class ReturningTest(fixtures.TestBase, AssertsExecutionResults):
         )
         eq_(result2.fetchall(), [(1, True), (2, False)])
 
+    @testing.requires.full_returning
+    def test_update_full_returning(self, connection):
+        connection.execute(
+            table.insert(),
+            [{"persons": 5, "full": False}, {"persons": 3, "full": False}],
+        )
+
+        result = connection.execute(
+            table.update(table.c.persons > 2)
+            .values(full=True)
+            .returning(table.c.id, table.c.full)
+        )
+        eq_(result.fetchall(), [(1, True), (2, True)])
+
+    @testing.requires.full_returning
+    def test_delete_full_returning(self, connection):
+        connection.execute(
+            table.insert(),
+            [{"persons": 5, "full": False}, {"persons": 3, "full": False}],
+        )
+
+        result = connection.execute(
+            table.delete().returning(table.c.id, table.c.full)
+        )
+        eq_(result.fetchall(), [(1, False), (2, False)])
+
     def test_insert_returning(self, connection):
         result = connection.execute(
             table.insert().returning(table.c.id), {"persons": 1, "full": False}
@@ -474,13 +500,6 @@ class ImplicitReturningFlag(fixtures.TestBase):
         testing.requires.returning(go)()
         e = engines.testing_engine()
 
-        # starts as False.  This is because all of Firebird,
-        # PostgreSQL, Oracle, SQL Server started supporting RETURNING
-        # as of a certain version, and the flag is not set until
-        # version detection occurs.  If some DB comes along that has
-        # RETURNING in all cases, this test can be adjusted.
-        assert e.dialect.implicit_returning is False
-
         # version detection on connect sets it
         c = e.connect()
         c.close()