Skip to content

Commit a2008bf

Browse files
committed
INTPYTHON-348 Add support for QuerySet.raw()
- Fix MongoRawModelIterable
1 parent 1d56969 commit a2008bf

File tree

1 file changed

+20
-2
lines changed

1 file changed

+20
-2
lines changed

django_mongodb/query.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -368,30 +368,48 @@ def iterator(self):
368368

369369

370370
class MongoRawModelIterable(RawModelIterable):
371+
"""
372+
Iterable that yields a model instance for each row from a raw queryset.
373+
"""
371374

372375
def __iter__(self):
376+
# Cache some things for performance reasons outside the loop.
373377
db = self.queryset.db
374378
query = self.queryset.query
375379
connection = connections[db]
376380
compiler = connection.ops.compiler("SQLCompiler")(query, connection, db)
377-
query_iterator = iter(query)
381+
query_iterator = self._make_result(query)
382+
378383
try:
379384
(
380385
model_init_names,
381386
model_init_pos,
382387
annotation_fields,
383388
) = self.queryset.resolve_model_init_order()
384389
model_cls = self.queryset.model
390+
if model_cls._meta.pk.attname not in model_init_names:
391+
raise exceptions.FieldDoesNotExist(
392+
"Raw query must include the primary key"
393+
)
385394
fields = [self.queryset.model_fields.get(c) for c in self.queryset.columns]
386395
converters = compiler.get_converters(
387396
[f.get_col(f.model._meta.db_table) if f else None for f in fields]
388397
)
389398
if converters:
390399
query_iterator = compiler.apply_converters(query_iterator, converters)
391400
for values in query_iterator:
392-
model_init_values = values.values()
401+
# Associate fields to values
402+
model_init_values = [values[pos] for pos in model_init_pos]
393403
instance = model_cls.from_db(db, model_init_names, model_init_values)
404+
if annotation_fields:
405+
for column, pos in annotation_fields:
406+
setattr(instance, column, values[pos])
394407
yield instance
395408
finally:
409+
# Done iterating the Query. If it has its own cursor, close it.
396410
if hasattr(query, "cursor") and query.cursor:
397411
query.cursor.close()
412+
413+
def _make_result(self, query):
414+
for result in query:
415+
yield list(result.values())

0 commit comments

Comments
 (0)