@@ -368,30 +368,48 @@ def iterator(self):
368
368
369
369
370
370
class MongoRawModelIterable (RawModelIterable ):
371
+ """
372
+ Iterable that yields a model instance for each row from a raw queryset.
373
+ """
371
374
372
375
def __iter__ (self ):
376
+ # Cache some things for performance reasons outside the loop.
373
377
db = self .queryset .db
374
378
query = self .queryset .query
375
379
connection = connections [db ]
376
380
compiler = connection .ops .compiler ("SQLCompiler" )(query , connection , db )
377
- query_iterator = iter (query )
381
+ query_iterator = self ._make_result (query )
382
+
378
383
try :
379
384
(
380
385
model_init_names ,
381
386
model_init_pos ,
382
387
annotation_fields ,
383
388
) = self .queryset .resolve_model_init_order ()
384
389
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
+ )
385
394
fields = [self .queryset .model_fields .get (c ) for c in self .queryset .columns ]
386
395
converters = compiler .get_converters (
387
396
[f .get_col (f .model ._meta .db_table ) if f else None for f in fields ]
388
397
)
389
398
if converters :
390
399
query_iterator = compiler .apply_converters (query_iterator , converters )
391
400
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 ]
393
403
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 ])
394
407
yield instance
395
408
finally :
409
+ # Done iterating the Query. If it has its own cursor, close it.
396
410
if hasattr (query , "cursor" ) and query .cursor :
397
411
query .cursor .close ()
412
+
413
+ def _make_result (self , query ):
414
+ for result in query :
415
+ yield list (result .values ())
0 commit comments