Commits

Vladimir Mihailenco committed 1333ab0

Refactoring

Comments (0)

Files changed (4)

 class DatabaseFeatures(NonrelDatabaseFeatures):
     allows_primary_key_0 = True
     supports_dicts = True
+    supports_batches = True
 
 class DatabaseOperations(NonrelDatabaseOperations):
     compiler_module = __name__.rsplit('.', 1)[0] + '.compiler'
 
     def get_instance_data(self, instance, raw=False):
         data = {}
-        for field in instance.get_fields():
+        for field in instance._meta.local_fields:
             if raw:
                 value = getattr(self, field.attname)
             else:
 class SQLDeleteCompiler(NonrelDeleteCompiler, SQLCompiler):
     pass
 
+
 class SQLBatchCompiler(SQLCompiler):
-    def __init__(self, query, connection, using, raw=False):
+    def __init__(self, query, connection, using, raw=False, datastore_deadline=15, **kwargs):
         super(SQLBatchCompiler, self).__init__(query, connection, using)
         self.raw = raw
-        self.pool = MutationPool(self)
+        self.datastore_deadline = datastore_deadline
+
+    def _create_rpc(self):
+        return datastore.CreateRPC(deadline=self.datastore_deadline)
 
     @safe_call
-    def save(self, instance):
-        data = self.get_instance_data(instance, raw=self.raw)
-        entity = self.create_entity(instance.__class__, data)
+    def batch_save(self, instances):
+        items = []
+        for instance in instances:
+            data = self.get_instance_data(instance, raw=self.raw)
+            items.append(self.create_entity(instance.__class__, data))
 
-        self.pool.put(entity)
+        keys = datastore.Put(items, rpc=self._create_rpc())
+        pk_vals = [key.id_or_name() for key in keys]
 
-    def on_batch_save(self, keys):
-        pk_vals = [key.id_or_name() for key in keys]
-        self.query.on_batch_save(pk_vals)
+        return pk_vals
 
     @safe_call
-    def delete(self, instance):
-        key = create_key(instance._meta.db_table, instance.pk)
-        self.pool.delete(key)
+    def batch_delete(self, instances):
+        keys = []
+        for instance in instances:
+            keys.append(create_key(instance._meta.db_table, instance.pk))
 
-    def on_batch_delete(self):
-        self.query.on_batch_delete()
-
-    def flush(self):
-        self.pool.flush()
+        datastore.Delete(keys, rpc=self._create_rpc())
 
 def to_datetime(value):
     """Convert a time or date to a datetime for datastore storage.

db/mutation_pool.py

-from google.appengine.api import datastore
-from google.appengine.ext.mapreduce.context import MutationPool as BaseMutationPool
-
-
-class MutationPool(BaseMutationPool):
-    def __init__(self, batch_operation=None, *args, **kwargs):
-        super(MutationPool, self).__init__(*args, **kwargs)
-        self.batch_operation = batch_operation
-
-    def __flush_puts(self):
-        """Flush all puts to datastore."""
-        if self.puts.length:
-            keys = datastore.Put(self.puts.items, rpc=self.__create_rpc())
-            self.batch_operation.on_batch_save(keys)
-        self.puts.clear()
-
-    def __flush_deletes(self):
-        """Flush all deletes to datastore."""
-        if self.deletes.length:
-            datastore.Delete(self.deletes.items, rpc=self.__create_rpc())
-            self.batch_operation.on_batch_delete()
-        self.deletes.clear()
 from django.db.models.signals import post_save
 from django.db.models.signals import pre_delete
 from django.db.models.signals import post_delete
-from django.db.models import BatchQuery
+from django.db.models import batch_query
 from django.test import TestCase
 
 from google.appengine.api import datastore
         datastore.Delete = Delete
 
     def test_batch_save(self):
-        with BatchQuery() as q:
+        with batch_query() as q:
             categories = []
             for name in ['appengine', 'django']:
                 c = BatchCategory(name=name)
         self.assertEqual(BatchCategory.objects.count(), 2)
         self.assertEqual(mem.put_count, 1)
 
-        with BatchQuery() as q:
+        with batch_query() as q:
             categories = BatchCategory.objects.all()
             for c in categories:
                 q.delete(c)
         self.assertEqual(mem.delete_count, 1)
 
     def test_plain_operation_in_batch(self):
-        with BatchQuery() as q:
+        with batch_query() as q:
             appengine = BatchCategory(name='appengine')
             q.save(appengine)
             self.assertEqual(BatchCategory.objects.count(), 0)
 
         appengine = BatchCategory.objects.get(name='appengine')
         django = BatchCategory.objects.get(name='django')
-        with BatchQuery() as q:
+        with batch_query() as q:
             q.delete(appengine)
             self.assertEqual(BatchCategory.objects.count(), 2)