Merge lp:~wgrant/launchpad/bulk-insert into lp:launchpad

[William Grant]

[r=sinzui][no-qa] Introduce lp.services.database.bulk.create().

This branch adds lp.services.database.bulk.create(), a helper to insert multiple DB rows in a single statement.

It touches some of Storm's privates in inappropriate ways in order to handle References in a pretty (and typesafe) manner. Reference columns and values have to be flattened into (potentially compound) primary keys, and the methods to do that are private.

[Curtis Hovey (sinzui)]

Change the assert on line 82. It should be explicitly raised if the method has an integrity issue. Is this an AssertionError? I read "The Storm columns to insert values into. Must be from a single class", which implies this is a ValueError. Maybe?
    if len(clses) == 1:
        raise ValueError('The Storm columns to insert values into must be from a single class.')

review: Needs Fixing (code)

=== modified file 'lib/lp/services/database/bulk.py'
--- lib/lp/services/database/bulk.py	2012-02-20 07:11:22 +0000
+++ lib/lp/services/database/bulk.py	2012-02-21 22:17:21 +0000
@@ -5,6 +5,7 @@
 
 __metaclass__ = type
 __all__ = [
+    'create',
     'load',
     'load_referencing',
     'load_related',
@@ -13,14 +14,24 @@
 
 
 from collections import defaultdict
+from itertools import chain
 from functools import partial
-from operator import attrgetter
+from operator import (
+    attrgetter,
+    itemgetter,
+    )
 
+from storm.databases.postgres import Returning
 from storm.expr import (
     And,
+    Insert,
     Or,
     )
-from storm.info import get_cls_info
+from storm.info import (
+    get_cls_info,
+    get_obj_info,
+    )
+from storm.references import Reference
 from storm.store import Store
 from zope.security.proxy import removeSecurityProxy
 
@@ -153,3 +164,64 @@
     for owning_object in owning_objects:
         keys.update(map(partial(getattr, owning_object), foreign_keys))
     return load(object_type, keys)
+
+
+def _dbify_value(col, val):
+    """Convert a value into a form that Storm can compile directly."""
+    if isinstance(col, Reference):
+        # References are mainly meant to be used as descriptors, so we
+        # have to perform a bit of evil here to turn the (potentially
+        # None) value into a sequence of primary key values.
+        if val is None:
+            return (None,) * len(col._relation._get_local_columns(col._cls))
+        else:
+            return col._relation.get_remote_variables(
+                get_obj_info(val).get_obj())
+    else:
+        return (col.variable_factory(value=val),)
+
+
+def _dbify_column(col):
+    """Convert a column into a form that Storm can compile directly."""
+    if isinstance(col, Reference):
+        # References are mainly meant to be used as descriptors, so we
+        # haver to perform a bit of evil here to turn the column into
+        # a sequence of primary key columns.
+        return col._relation._get_local_columns(col._cls)
+    else:
+        return (col,)
+
+
+def create(columns, values):
+    """Create a large number of objects efficiently.
+
+    :param cols: The Storm columns to insert values into. Must be from a
+        single class.
+    :param values: A list of lists of values for the columns.
+    :return: A list of the created objects.
+    """
+    # Flatten Reference faux-columns into their primary keys.
+    db_cols = list(chain.from_iterable(map(_dbify_column, columns)))
+    clses = set(col.cls for col in db_cols)
+    if len(clses) != 1:
+        raise ValueError(
+            "The Storm columns to insert values into must be from a single "
+            "class.")
+    [cls] = clses
+    primary_key = get_cls_info(cls).primary_key
+
+    # Mangle our value list into compilable values. Normal columns just
+    # get passed through the variable factory, while References get
+    # squashed into primary key variables.
+    db_values = [
+        list(chain.from_iterable(
+            _dbify_value(col, val) for col, val in zip(columns, value)))
+        for value in values]
+
+    result = IStore(cls).execute(
+        Returning(Insert(
+            db_cols, expr=db_values, primary_columns=primary_key)))
+    if len(primary_key) == 1:
+        return load(cls, map(itemgetter(0), result))
+    else:
+        return load(cls, result)
 
=== modified file 'lib/lp/services/database/tests/test_bulk.py'
--- lib/lp/services/database/tests/test_bulk.py	2012-02-20 07:11:22 +0000
+++ lib/lp/services/database/tests/test_bulk.py	2012-02-21 22:17:21 +0000
@@ -5,16 +5,26 @@
 
 __metaclass__ = type
 
+import datetime
+
+from pytz import UTC
 from storm.exceptions import ClassInfoError
 from storm.info import get_obj_info
 from storm.store import Store
+from testtools.matchers import Equals
 import transaction
 from zope.security import (
     checker,
     proxy,
     )
 
+from lp.bugs.enum import BugNotificationLevel
 from lp.bugs.model.bug import BugAffectsPerson
+from lp.bugs.model.bugsubscription import BugSubscription
+from lp.code.model.branchjob import (
+    BranchJob,
+    BranchJobType,
+    )
 from lp.code.model.branchsubscription import BranchSubscription
 from lp.registry.model.person import Person
 from lp.services.database import bulk
@@ -27,12 +37,15 @@
     FeatureFlag,
     getFeatureStore,
     )
+from lp.services.job.model.job import Job
 from lp.soyuz.model.component import Component
 from lp.testing import (
+    StormStatementRecorder,
     TestCase,
     TestCaseWithFactory,
     )
 from lp.testing.layers import DatabaseFunctionalLayer
+from lp.testing.matchers import HasQueryCount
 
 
 object_is_key = lambda thing: thing
@@ -219,9 +232,60 @@
             self.factory.makeBranch(),
             self.factory.makeBranch(),
             ]
-        expected = set(list(owned_objects[0].subscriptions) + 
+        expected = set(list(owned_objects[0].subscriptions) +
             list(owned_objects[1].subscriptions))
         self.assertNotEqual(0, len(expected))
         self.assertEqual(expected,
             set(bulk.load_referencing(BranchSubscription, owned_objects,
                 ['branchID'])))
+
+
+class TestCreate(TestCaseWithFactory):
+
+    layer = DatabaseFunctionalLayer
+
+    def test_references_and_enums(self):
+        # create() correctly compiles plain types, enums and references.
+        bug = self.factory.makeBug()
+        people = [self.factory.makePerson() for i in range(5)]
+
+        wanted = [
+            (bug, person, person, datetime.datetime.now(UTC),
+             BugNotificationLevel.LIFECYCLE)
+            for person in people]
+
+        with StormStatementRecorder() as recorder:
+            subs = bulk.create(
+                (BugSubscription.bug, BugSubscription.person,
+                 BugSubscription.subscribed_by, BugSubscription.date_created,
+                 BugSubscription.bug_notification_level),
+                wanted)
+
+        self.assertThat(recorder, HasQueryCount(Equals(2)))
+        self.assertContentEqual(
+            wanted,
+            ((sub.bug, sub.person, sub.subscribed_by, sub.date_created,
+              sub.bug_notification_level) for sub in subs))
+
+    def test_null_reference(self):
+        # create() handles None as a Reference value.
+        job = Job()
+        IStore(Job).add(job)
+        wanted = [(None, job, BranchJobType.RECLAIM_BRANCH_SPACE)]
+        [branchjob] = bulk.create(
+            (BranchJob.branch, BranchJob.job, BranchJob.job_type),
+            wanted)
+        self.assertEqual(
+            wanted, [(branchjob.branch, branchjob.job, branchjob.job_type)])
+
+    def test_fails_on_multiple_classes(self):
+        # create() only inserts into columns on a single class.
+        self.assertRaises(
+            ValueError,
+            bulk.create, (BugSubscription.bug, BranchSubscription.branch), [])
+
+    def test_fails_on_reference_mismatch(self):
+        # create() handles Reference columns in a typesafe manner.
+        self.assertRaisesWithContent(
+            RuntimeError, "Property used in an unknown class",
+            bulk.create, (BugSubscription.bug,), [[self.factory.makeBranch()]])