Merge lp:~cjwatson/launchpad/better-gitref-bulk-loading into lp:launchpad

[Colin Watson]

Use a more compact query to load objects with compound primary keys.

Mentioned in passing in bug 1790047, although it's not really essential to that bug.

[William Grant (wgrant)]

review: Approve (code)

=== modified file 'lib/lp/services/database/bulk.py'
--- lib/lp/services/database/bulk.py	2013-08-23 05:36:46 +0000
+++ lib/lp/services/database/bulk.py	2018-09-12 09:02:58 +0000
@@ -1,4 +1,4 @@
-# Copyright 2010 Canonical Ltd.  This software is licensed under the
+# Copyright 2010-2018 Canonical Ltd.  This software is licensed under the
 # GNU Affero General Public License version 3 (see the file LICENSE).
 
 """Optimized bulk operations against the database."""
@@ -16,7 +16,10 @@
 
 from collections import defaultdict
 from functools import partial
-from itertools import chain
+from itertools import (
+    chain,
+    groupby,
+    )
 from operator import (
     attrgetter,
     itemgetter,
@@ -99,6 +102,32 @@
         return primary_key
 
 
+def _make_compound_load_clause(primary_key, values_list):
+    """Construct a bulk-loading query clause for a compound primary key.
+
+    When the primary key is compound, we expect that in practice we will
+    often want to load objects with common leading elements of the key: for
+    example, we often want to load many `GitRef` objects from the same
+    repository.  This helper returns a query clause optimised to be compact
+    in this case.
+
+    :param primary_key: The primary key columns.
+    :param values_list: A sequence of values of the primary key whose
+        corresponding rows should be loaded.  This must be sorted by the
+        caller (because this function calls itself recursively, and it's
+        usually more efficient to sort the whole sequence in one go).
+    """
+    if len(primary_key) > 1:
+        return Or(*(
+            And(
+                primary_key[0] == leading_value,
+                _make_compound_load_clause(
+                    primary_key[1:], [values[1:] for values in group]))
+            for leading_value, group in groupby(values_list, itemgetter(0))))
+    else:
+        return primary_key[0].is_in([values[0] for values in values_list])
+
+
 def load(object_type, primary_keys, store=None):
     """Load a large number of objects efficiently."""
     primary_key = _primary_key(object_type, allow_compound=True)
@@ -107,9 +136,8 @@
     if not primary_keys:
         return []
     if isinstance(primary_key, tuple):
-        condition = Or(*(
-            And(*(key == value for (key, value) in zip(primary_key, values)))
-            for values in primary_keys))
+        condition = _make_compound_load_clause(
+            primary_key, sorted(primary_keys))
     else:
         condition = primary_key.is_in(primary_keys)
     if store is None:
 
=== modified file 'lib/lp/services/database/tests/test_bulk.py'
--- lib/lp/services/database/tests/test_bulk.py	2018-04-19 00:02:19 +0000
+++ lib/lp/services/database/tests/test_bulk.py	2018-09-12 09:02:58 +0000
@@ -1,4 +1,4 @@
-# Copyright 2010 Canonical Ltd.  This software is licensed under the
+# Copyright 2010-2018 Canonical Ltd.  This software is licensed under the
 # GNU Affero General Public License version 3 (see the file LICENSE).
 
 """Test the bulk database functions."""
@@ -35,7 +35,10 @@
     ISlaveStore,
     IStore,
     )
-from lp.services.database.sqlbase import get_transaction_timestamp
+from lp.services.database.sqlbase import (
+    convert_storm_clause_to_string,
+    get_transaction_timestamp,
+    )
 from lp.services.features.model import (
     FeatureFlag,
     getFeatureStore,
@@ -172,6 +175,27 @@
         bulk.reload([db_object])
         self.assertIsNone(db_object_info.get('invalidated'))
 
+    def test__make_compound_load_clause(self):
+        # The query constructed by _make_compound_load_clause has the
+        # correct structure.
+        clause = bulk._make_compound_load_clause(
+            # This primary key is fictional, but lets us do a more elaborate
+            # test.
+            (FeatureFlag.scope, FeatureFlag.priority, FeatureFlag.flag),
+            sorted(
+                [(u'foo', 0, u'bar'), (u'foo', 0, u'baz'),
+                 (u'foo', 1, u'bar'), (u'foo', 1, u'quux'),
+                 (u'bar', 0, u'foo')]))
+        self.assertEqual(
+            "FeatureFlag.scope = E'bar' AND ("
+            "FeatureFlag.priority = 0 AND FeatureFlag.flag IN (E'foo')) OR "
+            "FeatureFlag.scope = E'foo' AND ("
+            "FeatureFlag.priority = 0 AND "
+            "FeatureFlag.flag IN (E'bar', E'baz') OR "
+            "FeatureFlag.priority = 1 AND "
+            "FeatureFlag.flag IN (E'bar', E'quux'))",
+            convert_storm_clause_to_string(clause))
+
     def test_load(self):
         # load() loads objects of the given type by their primary keys.
         db_objects = [
@@ -189,7 +213,7 @@
             ClassInfoError, bulk.load, str, [])
 
     def test_load_with_compound_primary_keys(self):
-        # load() does not like compound primary keys.
+        # load() can load objects with compound primary keys.
         flags = [
             FeatureFlag(u'foo', 0, u'bar', u'true'),
             FeatureFlag(u'foo', 0, u'baz', u'false'),