Merge pull request pytest-dev#16 from cypreess/master

Implement async `prefetch()`, close pytest-dev#11. Thanks @cypreess 👍 !

Author: rudyryk

docs/peewee_async/api.rst

@@ -11,6 +11,7 @@ Select, update, delete
 .. autofunction:: peewee_async.create_object
 .. autofunction:: peewee_async.delete_object
 .. autofunction:: peewee_async.update_object
+.. autofunction:: peewee_async.prefetch
 
 Transactions
 ------------

peewee_async.py

@@ -720,3 +720,44 @@ def _compose_dsn(dbname, **kwargs):
         if v:
             dsn += ' %s=%s' % (k, v)
     return dsn, kwargs
+
+
+@asyncio.coroutine
+def prefetch(sq, *subqueries):
+    """Asynchronous version of the prefetch function from peewee.
+
+    Returns Query that has already cached data.
+    """
+
+    # This code is copied from peewee.prefetch and adopted to use async execute
+
+    if not subqueries:
+        return sq
+    fixed_queries = peewee.prefetch_add_subquery(sq, subqueries)
+
+    deps = {}
+    rel_map = {}
+    for prefetch_result in reversed(fixed_queries):
+        query_model = prefetch_result.model
+        if prefetch_result.fields:
+            for rel_model in prefetch_result.rel_models:
+                rel_map.setdefault(rel_model, [])
+                rel_map[rel_model].append(prefetch_result)
+
+        deps[query_model] = {}
+        id_map = deps[query_model]
+        has_relations = bool(rel_map.get(query_model))
+
+        # This is hack, because peewee async execute do a copy of query and do not change state of query
+        # comparing to what real peewee is doing when execute method is called
+        prefetch_result.query._qr = yield from execute(prefetch_result.query)
+        prefetch_result.query._dirty = False
+
+        for instance in prefetch_result.query._qr:
+            if prefetch_result.fields:
+                prefetch_result.store_instance(instance, id_map)
+            if has_relations:
+                for rel in rel_map[query_model]:
+                    rel.populate_instance(instance, deps[rel.model])
+
+    return prefetch_result.query

tests/__init__.py

@@ -108,6 +108,29 @@ class Meta:
         database = database
 
 
+class TestModelAlpha(peewee.Model):
+    text = peewee.CharField()
+
+    class Meta:
+        database = database
+
+
+class TestModelBeta(peewee.Model):
+    alpha = peewee.ForeignKeyField(TestModelAlpha, related_name='betas')
+    text = peewee.CharField()
+
+    class Meta:
+        database = database
+
+
+class TestModelGamma(peewee.Model):
+    text = peewee.CharField()
+    beta = peewee.ForeignKeyField(TestModelBeta, related_name='gammas')
+
+    class Meta:
+        database = database
+
+
 class UUIDTestModel(peewee.Model):
     id = peewee.UUIDField(primary_key=True, default=uuid.uuid4)
     text = peewee.CharField()
@@ -117,6 +140,8 @@ class Meta:
 
 
 class BaseAsyncPostgresTestCase(unittest.TestCase):
+    db_tables = [TestModel, UUIDTestModel, TestModelAlpha, TestModelBeta, TestModelGamma]
+
     @classmethod
     def setUpClass(cls, *args, **kwargs):
         # Sync connect 
@@ -129,22 +154,38 @@ def test():
             yield from database.connect_async(loop=cls.loop)
         cls.loop.run_until_complete(test())
 
-        # Clean up after possible errors
-        TestModel.drop_table(True)
-        UUIDTestModel.drop_table(True)
+        for table in reversed(cls.db_tables):
+            # Clean up after possible errors
+            table.drop_table(True, cascade=True)
 
-        # Create table with sync connection
-        TestModel.create_table()
-        UUIDTestModel.create_table()
+        for table in cls.db_tables:
+            # Create table with sync connection
+            table.create_table()
 
         # Create at least one object per model
         cls.obj = TestModel.create(text='[sync] Hello!')
         cls.uuid_obj = UUIDTestModel.create(text='[sync] Hello!')
 
+        cls.alpha_1 = TestModelAlpha.create(text='Alpha 1')
+        cls.alpha_2 = TestModelAlpha.create(text='Alpha 2')
+
+        cls.beta_11 = TestModelBeta.create(text='Beta 1', alpha=cls.alpha_1)
+        cls.beta_12 = TestModelBeta.create(text='Beta 2', alpha=cls.alpha_1)
+
+        cls.beta_21 = TestModelBeta.create(text='Beta 1', alpha=cls.alpha_2)
+        cls.beta_22 = TestModelBeta.create(text='Beta 2', alpha=cls.alpha_2)
+
+        cls.gamma_111 = TestModelGamma.create(text='Gamma 1', beta=cls.beta_11)
+        cls.gamma_112 = TestModelGamma.create(text='Gamma 2', beta=cls.beta_11)
+
+        cls.gamma_121 = TestModelGamma.create(text='Gamma 1', beta=cls.beta_12)
+
+
     @classmethod
     def tearDownClass(cls, *args, **kwargs):
+        for table in reversed(cls.db_tables):
         # Finally, clean up
-        TestModel.drop_table()
+            table.drop_table()
 
         # Close database
         database.close()
@@ -351,6 +392,29 @@ def test():
 
         self.run_until_complete(test())
 
+    def test_prefetch(self):
+        # Async prefetch
+        @asyncio.coroutine
+        def test():
+            with sync_unwanted(database):
+                result = yield from peewee_async.prefetch(TestModelAlpha.select(), TestModelBeta.select(),
+                                                          TestModelGamma.select())
+
+                result = list(result)   # this should NOT fire any call (will read it from query cache)
+
+                # Check if we have here both alpha items in specific order
+                self.assertEqual(result, [self.alpha_1, self.alpha_2])
+
+                alpha_1 = result[0]
+                self.assertEqual(alpha_1.betas_prefetch, [self.beta_11, self.beta_12])
+
+                beta_11 = alpha_1.betas_prefetch[0]
+                self.assertEqual(beta_11, self.beta_11)
+
+                self.assertEqual(beta_11.gammas_prefetch, [self.gamma_111, self.gamma_112])
+
+        self.run_until_complete(test())
+
 
 if sys.version_info >= (3, 5):
     from .tests_py35 import *