From cbd661e0cdfdba98663ae542c6af2863fc30ae09 Mon Sep 17 00:00:00 2001
From: Alessandro Cucci <alessandro.cucci@gmail.com>
Date: Sat, 25 Aug 2018 09:14:22 -0400
Subject: [PATCH] Add option to sort into inserts/updates to bulk_save_objects

Added new flag :paramref:`.Session.bulk_save_objects.preserve_order` to the
:meth:`.Session.bulk_save_objects` method, which defaults to True. When set
to False, the given mappings will be grouped into inserts and updates per
each object type, to allow for greater opportunities to batch common
operations together.  Pull request courtesy Alessandro Cucci.

Change-Id: I0d041f7696cf733655a74beeceee3fa80640efd7
Pull-request: https://bitbucket.org/zzzeek/sqlalchemy/pull-requests/6
---
 .../unreleased_13/preserve_order.rst          |  8 +++
 lib/sqlalchemy/orm/session.py                 | 20 ++++++--
 test/orm/test_bulk.py                         | 51 +++++++++++++++++++
 3 files changed, 76 insertions(+), 3 deletions(-)
 create mode 100644 doc/build/changelog/unreleased_13/preserve_order.rst

diff --git a/doc/build/changelog/unreleased_13/preserve_order.rst b/doc/build/changelog/unreleased_13/preserve_order.rst
new file mode 100644
index 0000000000..3265da3030
--- /dev/null
+++ b/doc/build/changelog/unreleased_13/preserve_order.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: feature, orm
+
+    Added new flag :paramref:`.Session.bulk_save_objects.preserve_order` to the
+    :meth:`.Session.bulk_save_objects` method, which defaults to True. When set
+    to False, the given mappings will be grouped into inserts and updates per
+    each object type, to allow for greater opportunities to batch common
+    operations together.  Pull request courtesy Alessandro Cucci.
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index 9ac529aeb7..2c7fd86d2c 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -2380,7 +2380,8 @@ class Session(_SessionClassMethods):
                 transaction.rollback(_capture_exception=True)
 
     def bulk_save_objects(
-            self, objects, return_defaults=False, update_changed_only=True):
+            self, objects, return_defaults=False, update_changed_only=True,
+            preserve_order=True):
         """Perform a bulk save of the given list of objects.
 
         The bulk save feature allows mapped objects to be used as the
@@ -2443,6 +2444,13 @@ class Session(_SessionClassMethods):
          When False, all attributes present are rendered into the SET clause
          with the exception of primary key attributes.
 
+        :param preserve_order: when True, the order of inserts and updates
+         matches exactly the order in which the objects are given.   When
+         False, common types of objects are grouped into inserts
+         and updates, to allow for more batching opportunities.
+
+         .. versionadded:: 1.3
+
         .. seealso::
 
             :ref:`bulk_operations`
@@ -2452,9 +2460,15 @@ class Session(_SessionClassMethods):
             :meth:`.Session.bulk_update_mappings`
 
         """
+        def key(state):
+            return (state.mapper, state.key is not None)
+
+        obj_states = tuple(attributes.instance_state(obj) for obj in objects)
+        if not preserve_order:
+            obj_states = sorted(obj_states, key=key)
+
         for (mapper, isupdate), states in itertools.groupby(
-            (attributes.instance_state(obj) for obj in objects),
-            lambda state: (state.mapper, state.key is not None)
+            obj_states, key
         ):
             self._bulk_save_mappings(
                 mapper, states, isupdate, True,
diff --git a/test/orm/test_bulk.py b/test/orm/test_bulk.py
index 0763fe70cd..9d0a000380 100644
--- a/test/orm/test_bulk.py
+++ b/test/orm/test_bulk.py
@@ -1,6 +1,7 @@
 from sqlalchemy import testing
 from sqlalchemy.testing import eq_
 from sqlalchemy.testing.schema import Table, Column
+from sqlalchemy.testing import mock
 from sqlalchemy.testing import fixtures
 from sqlalchemy import Integer, String, ForeignKey, FetchedValue
 from sqlalchemy.orm import mapper, Session
@@ -107,6 +108,56 @@ class BulkInsertUpdateTest(BulkTest, _fixtures.FixtureTest):
         )
         eq_(objects[0].__dict__['id'], 1)
 
+    def test_bulk_save_mappings_preserve_order(self):
+        User, = self.classes("User", )
+
+        s = Session()
+
+        # commit some object into db
+        user1 = User(name="i1")
+        user2 = User(name="i2")
+        s.add(user1)
+        s.add(user2)
+        s.commit()
+
+        # make some changes
+        user1.name = "u1"
+        user3 = User(name="i3")
+        s.add(user3)
+        user2.name = "u2"
+
+        objects = [user1, user3, user2]
+
+        from sqlalchemy import inspect
+
+        def _bulk_save_mappings(
+                mapper, mappings, isupdate, isstates,
+                return_defaults, update_changed_only, render_nulls):
+            mock_method(list(mappings), isupdate)
+
+        mock_method = mock.Mock()
+        with mock.patch.object(s, '_bulk_save_mappings', _bulk_save_mappings):
+            s.bulk_save_objects(objects)
+            eq_(
+                mock_method.mock_calls,
+                [
+                    mock.call([inspect(user1)], True),
+                    mock.call([inspect(user3)], False),
+                    mock.call([inspect(user2)], True),
+                ]
+            )
+
+        mock_method = mock.Mock()
+        with mock.patch.object(s, '_bulk_save_mappings', _bulk_save_mappings):
+            s.bulk_save_objects(objects, preserve_order=False)
+            eq_(
+                mock_method.mock_calls,
+                [
+                    mock.call([inspect(user3)], False),
+                    mock.call([inspect(user1), inspect(user2)], True),
+                ]
+            )
+
     def test_bulk_save_no_defaults(self):
         User, = self.classes("User",)
 
-- 
2.47.2