Skip to content

Commit 832a6a9

Browse files
committed
INTPYTHON-348 add support for QuerySet.raw_mql()
- Subclass RawModelIterable
1 parent f9fbef8 commit 832a6a9

File tree

1 file changed

+36
-1
lines changed

1 file changed

+36
-1
lines changed

django_mongodb/query.py

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from django.db.models.expressions import Case, Col, When
88
from django.db.models.functions import Mod
99
from django.db.models.lookups import Exact
10-
from django.db.models.query import RawQuerySet
10+
from django.db.models.query import RawQuerySet, RawModelIterable
1111
from django.db.models.sql.constants import INNER
1212
from django.db.models.sql.datastructures import Join
1313
from django.db.models.sql.query import RawQuery
@@ -308,6 +308,7 @@ def register_nodes():
308308

309309

310310
class MongoQuerySet(QuerySet):
311+
311312
def raw_mql(self, raw_query, params=(), translations=None, using=None):
312313
return MongoRawQuerySet(
313314
raw_query,
@@ -360,6 +361,7 @@ def __init__(
360361
using=None,
361362
hints=None,
362363
):
364+
self._prefetch_related_lookups = None
363365
self.model = model
364366
self._db = using
365367
self._hints = hints or {}
@@ -382,3 +384,36 @@ def resolve_model_init_order(self):
382384
model_init_order = [self.columns.index(converter(f)) for f in model_init_fields]
383385
model_init_names = [f.attname for f in model_init_fields]
384386
return model_init_names, model_init_order, annotation_fields
387+
388+
def iterator(self):
389+
yield from MongoRawModelIterable(self)
390+
391+
392+
class MongoRawModelIterable(RawModelIterable):
393+
394+
def __iter__(self):
395+
db = self.queryset.db
396+
query = self.queryset.query
397+
connection = connections[db]
398+
compiler = connection.ops.compiler("SQLCompiler")(query, connection, db)
399+
query_iterator = iter(query)
400+
try:
401+
(
402+
model_init_names,
403+
model_init_pos,
404+
annotation_fields,
405+
) = self.queryset.resolve_model_init_order()
406+
model_cls = self.queryset.model
407+
fields = [self.queryset.model_fields.get(c) for c in self.queryset.columns]
408+
converters = compiler.get_converters(
409+
[f.get_col(f.model._meta.db_table) if f else None for f in fields]
410+
)
411+
if converters:
412+
query_iterator = compiler.apply_converters(query_iterator, converters)
413+
for values in query_iterator:
414+
model_init_values = values.values()
415+
instance = model_cls.from_db(db, model_init_names, model_init_values)
416+
yield instance
417+
finally:
418+
if hasattr(query, "cursor") and query.cursor:
419+
query.cursor.close()

0 commit comments

Comments
 (0)