1079 Add batch to raw queries

docs/src/piccolo/query_clauses/batch.rst

@@ -6,6 +6,7 @@ batch
 You can use ``batch`` clauses with the following queries:
 
 * :ref:`Objects`
+* :ref:`Raw`
 * :ref:`Select`
 
 Example

docs/src/piccolo/query_types/raw.rst

@@ -7,15 +7,26 @@ Should you need to, you can execute raw SQL.
 
 .. code-block:: python
 
-    >>> await Band.raw('select name from band')
+    >>> await Band.raw('SELECT name FROM band')
     [{'name': 'Pythonistas'}]
 
 It's recommended that you parameterise any values. Use curly braces ``{}`` as
 placeholders:
 
 .. code-block:: python
 
-    >>> await Band.raw('select * from band where name = {}', 'Pythonistas')
+    >>> await Band.raw('SELECT * FROM band WHERE name = {}', 'Pythonistas')
     [{'name': 'Pythonistas', 'manager': 1, 'popularity': 1000, 'id': 1}]
 
 .. warning:: Be careful to avoid SQL injection attacks. Don't add any user submitted data into your SQL strings, unless it's parameterised.
+
+
+-------------------------------------------------------------------------------
+
+Query clauses
+-------------
+
+batch
+~~~~~
+
+See :ref:`batch`.

piccolo/query/methods/raw.py

@@ -2,6 +2,7 @@
 
 import typing as t
 
+from piccolo.engine.base import BaseBatch
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
 
@@ -21,6 +22,18 @@ def __init__(
         super().__init__(table, **kwargs)
         self.querystring = querystring
 
+    async def batch(
+        self,
+        batch_size: t.Optional[int] = None,
+        node: t.Optional[str] = None,
+        **kwargs,
+    ) -> BaseBatch:
+        if batch_size:
+            kwargs.update(batch_size=batch_size)
+        if node:
+            kwargs.update(node=node)
+        return await self.table._meta.db.batch(self, **kwargs)
+
     @property
     def default_querystrings(self) -> t.Sequence[QueryString]:
         return [self.querystring]

tests/table/test_batch.py

@@ -95,6 +95,47 @@ def test_batch(self):
         self.assertEqual(iterations, _iterations)
 
 
+class TestBatchRaw(DBTestCase):
+    def _check_results(self, batch):
+        """
+        Make sure the data is returned in the correct format.
+        """
+        self.assertEqual(type(batch), list)
+        if len(batch) > 0:
+            row = batch[0]
+            self.assertIsInstance(row, Manager)
+
+    async def run_batch(self, batch_size):
+        row_count = 0
+        iterations = 0
+
+        async with await Manager.raw("SELECT * FROM manager").batch(
+            batch_size=batch_size
+        ) as batch:
+            async for _batch in batch:
+                self._check_results(_batch)
+                _row_count = len(_batch)
+                row_count += _row_count
+                iterations += 1
+
+        return row_count, iterations
+
+    async def test_batch(self):
+        row_count = 1000
+        self.insert_many_rows(row_count)
+
+        batch_size = 10
+
+        _row_count, iterations = asyncio.run(
+            self.run_batch(batch_size=batch_size), debug=True
+        )
+
+        _iterations = math.ceil(row_count / batch_size)
+
+        self.assertEqual(_row_count, row_count)
+        self.assertEqual(iterations, _iterations)
+
+
 @engines_only("postgres", "cockroach")
 class TestBatchNodeArg(TestCase):
     def test_batch_extra_node(self):