Refs #33646 -- Made QuerySet.raw() async-compatible.

This commit is contained in:
Mariusz Felisiak 2022-04-27 11:30:43 +02:00
parent 77926176b2
commit 271a8e73ee
2 changed files with 21 additions and 0 deletions

View File

@ -2025,6 +2025,12 @@ class RawQuerySet:
if self._prefetch_related_lookups and not self._prefetch_done: if self._prefetch_related_lookups and not self._prefetch_done:
self._prefetch_related_objects() self._prefetch_related_objects()
async def _async_fetch_all(self):
if self._result_cache is None:
self._result_cache = [result async for result in RawModelIterable(self)]
if self._prefetch_related_lookups and not self._prefetch_done:
sync_to_async(self._prefetch_related_objects)()
def __len__(self): def __len__(self):
self._fetch_all() self._fetch_all()
return len(self._result_cache) return len(self._result_cache)
@ -2037,6 +2043,16 @@ class RawQuerySet:
self._fetch_all() self._fetch_all()
return iter(self._result_cache) return iter(self._result_cache)
def __aiter__(self):
# Remember, __aiter__ itself is synchronous, it's the thing it returns
# that is async!
async def generator():
await self._async_fetch_all()
for item in self._result_cache:
yield item
return generator()
def iterator(self): def iterator(self):
yield from RawModelIterable(self) yield from RawModelIterable(self)

View File

@ -225,3 +225,8 @@ class AsyncQuerySetTest(TestCase):
json.loads(result) json.loads(result)
except json.JSONDecodeError as e: except json.JSONDecodeError as e:
self.fail(f"QuerySet.aexplain() result is not valid JSON: {e}") self.fail(f"QuerySet.aexplain() result is not valid JSON: {e}")
async def test_raw(self):
sql = "SELECT id, field FROM async_queryset_simplemodel WHERE created=%s"
qs = SimpleModel.objects.raw(sql, [self.s1.created])
self.assertEqual([o async for o in qs], [self.s1])