tidy code and remove unncessary overrides

Author: timgraham

django_mongodb/query.py

@@ -315,92 +315,62 @@ def raw_mql(self, pipeline, using=None):
 
 class MongoRawQuery(RawQuery):
     def __init__(self, pipeline, using, model):
+        self.pipeline = pipeline
         super().__init__(sql=None, using=using)
         self.model = model
-        self.pipeline = pipeline
-
-    def __iter__(self):
-        self._execute_query()
-        return self.cursor
 
     def _execute_query(self):
         connection = connections[self.using]
         collection = connection.get_collection(self.model._meta.db_table)
         self.cursor = collection.aggregate(self.pipeline)
 
-    def get_columns(self):
-        return [f.column for f in self.model._meta.fields]
-
     def __str__(self):
-        return "%s" % self.pipeline
+        return str(self.pipeline)
 
 
 class MongoRawQuerySet(RawQuerySet):
-    def __init__(
-        self,
-        pipeline,
-        model=None,
-        query=None,
-        translations=None,
-        using=None,
-        hints=None,
-    ):
-        super().__init__(
-            pipeline,
-            model=model,
-            query=query,
-            using=using,
-            hints=hints,
-            translations=translations,
-        )
-        self.query = query or MongoRawQuery(pipeline, using=self.db, model=self.model)
+    def __init__(self, pipeline, model=None, using=None):
+        super().__init__(pipeline, model=model, using=using)
+        self.query = MongoRawQuery(pipeline, using=self.db, model=self.model)
+        # Override the superclass's columns property which relies on PEP 249's
+        # cursor.description. Instead, RawModelIterable will set the columns
+        # based on the keys in the first result.
+        self.columns = None
 
     def iterator(self):
         yield from MongoRawModelIterable(self)
 
-    def resolve_model_init_order(self, columns):
-        """Resolve the init field names and value positions."""
-        model_init_fields = [f for f in self.model._meta.fields if f.column in columns]
-        annotation_fields = [
-            (column, pos) for pos, column in enumerate(columns) if column not in self.model_fields
-        ]
-        model_init_order = [columns.index(f.column) for f in model_init_fields]
-        model_init_names = [f.attname for f in model_init_fields]
-        return model_init_names, model_init_order, annotation_fields
-
 
 class MongoRawModelIterable(RawModelIterable):
-    """
-    Iterable that yields a model instance for each row from a raw queryset.
-    """
-
     def __iter__(self):
-        # Cache some things for performance reasons outside the loop.
+        """
+        This is mostly copied from the superclass except for the part that
+        sets self.queryset.columns from the first document.
+        """
         db = self.queryset.db
         query = self.queryset.query
         connection = connections[db]
         compiler = connection.ops.compiler("SQLCompiler")(query, connection, db)
         query_iterator = iter(query)
-        # Get the columns from the first result.
-        try:
-            first_item = next(query_iterator)
-        except StopIteration:
-            # No results.
-            query.cursor.close()
-            return
-        columns = list(first_item.keys())
-        # Reset the iterator to include the first item.
-        query_iterator = self._make_result(chain([first_item], query_iterator))
         try:
+            # Get the columns from the first result.
+            try:
+                first_item = next(query_iterator)
+            except StopIteration:
+                # No results.
+                return
+            self.queryset.columns = list(first_item.keys())
+            # Reset the iterator to include the first item.
+            query_iterator = self._make_result(chain([first_item], query_iterator))
             (
                 model_init_names,
                 model_init_pos,
                 annotation_fields,
-            ) = self.queryset.resolve_model_init_order(columns)
+            ) = self.queryset.resolve_model_init_order()
             model_cls = self.queryset.model
             if model_cls._meta.pk.attname not in model_init_names:
                 raise FieldDoesNotExist("Raw query must include the primary key")
-            fields = [self.queryset.model_fields.get(c) for c in columns]
+            fields = [self.queryset.model_fields.get(c) for c in self.queryset.columns]
             converters = compiler.get_converters(
                 [f.get_col(f.model._meta.db_table) if f else None for f in fields]
             )
@@ -415,9 +385,7 @@ def __iter__(self):
                         setattr(instance, column, values[pos])
                 yield instance
         finally:
-            # Done iterating the Query. If it has its own cursor, close it.
-            if hasattr(query, "cursor") and query.cursor:
-                query.cursor.close()
+            query.cursor.close()
 
     def _make_result(self, query):
         for result in query: