fix peek to get columns

Author: timgraham

django_mongodb/query.py

@@ -1,4 +1,5 @@
 from functools import reduce, wraps
+from itertools import chain
 from operator import add as add_operator
 
 from django.core.exceptions import EmptyResultSet, FieldDoesNotExist, FullResultSet
@@ -387,11 +388,17 @@ def __iter__(self):
         query = self.queryset.query
         connection = connections[db]
         compiler = connection.ops.compiler("SQLCompiler")(query, connection, db)
-        query_iterator = self._make_result(query)
-        # FIXME Get columns from first row of the query without iterating over
-        # the query.
-        columns = self.queryset.columns
-
+        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:
             (
                 model_init_names,
@@ -401,7 +408,7 @@ def __iter__(self):
             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 self.queryset.columns]
+            fields = [self.queryset.model_fields.get(c) for c in columns]
             converters = compiler.get_converters(
                 [f.get_col(f.model._meta.db_table) if f else None for f in fields]
             )