rename module/move code

Author: timgraham

django_mongodb/features.py

@@ -412,9 +412,6 @@ def django_test_expected_failures(self):
             "model_fields.test_jsonfield.TestQuerying.test_nested_key_transform_raw_expression",
             "multiple_database.tests.QueryTestCase.test_raw",
             "prefetch_related.tests.RawQuerySetTests",
-            "raw_query.tests.RawQueryTests.test_params",
-            "raw_query.tests.RawQueryTests.test_params_none",
-            "raw_query.tests.RawQueryTests.test_escaped_percent",
             "queries.tests.Queries1Tests.test_order_by_rawsql",
             "schema.test_logging.SchemaLoggerTests.test_extra_args",
             "schema.tests.SchemaTests.test_remove_constraints_capital_letters",

django_mongodb/manager.py renamed to django_mongodb/managers.py

@@ -1,6 +1,6 @@
 from django.db.models.manager import BaseManager
 
-from .query import MongoQuerySet
+from .queryset import MongoQuerySet
 
 
 class MongoManager(BaseManager.from_queryset(MongoQuerySet)):

django_mongodb/query.py

@@ -1,17 +1,13 @@
 from functools import reduce, wraps
-from itertools import chain
 from operator import add as add_operator
 
-from django.core.exceptions import EmptyResultSet, FieldDoesNotExist, FullResultSet
-from django.db import DatabaseError, IntegrityError, NotSupportedError, connections
-from django.db.models import QuerySet
+from django.core.exceptions import EmptyResultSet, FullResultSet
+from django.db import DatabaseError, IntegrityError, NotSupportedError
 from django.db.models.expressions import Case, Col, When
 from django.db.models.functions import Mod
 from django.db.models.lookups import Exact
-from django.db.models.query import RawModelIterable, RawQuerySet
 from django.db.models.sql.constants import INNER
 from django.db.models.sql.datastructures import Join
-from django.db.models.sql.query import RawQuery
 from django.db.models.sql.where import AND, OR, XOR, ExtraWhere, NothingNode, WhereNode
 from pymongo.errors import BulkWriteError, DuplicateKeyError, PyMongoError
 
@@ -306,87 +302,3 @@ def register_nodes():
     Join.as_mql = join
     NothingNode.as_mql = NothingNode.as_sql
     WhereNode.as_mql = where_node
-
-
-class MongoQuerySet(QuerySet):
-    def raw_mql(self, pipeline, using=None):
-        return MongoRawQuerySet(pipeline, model=self.model, using=using)
-
-
-class MongoRawQuery(RawQuery):
-    def __init__(self, pipeline, using, model):
-        self.pipeline = pipeline
-        super().__init__(sql=None, using=using)
-        self.model = model
-
-    def _execute_query(self):
-        connection = connections[self.using]
-        collection = connection.get_collection(self.model._meta.db_table)
-        self.cursor = collection.aggregate(self.pipeline)
-
-    def __str__(self):
-        return str(self.pipeline)
-
-
-class MongoRawQuerySet(RawQuerySet):
-    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)
-
-
-class MongoRawModelIterable(RawModelIterable):
-    def __iter__(self):
-        """
-        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)
-        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()
-            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]
-            converters = compiler.get_converters(
-                [f.get_col(f.model._meta.db_table) if f else None for f in fields]
-            )
-            if converters:
-                query_iterator = compiler.apply_converters(query_iterator, converters)
-            for values in query_iterator:
-                # Associate fields to values
-                model_init_values = [values[pos] for pos in model_init_pos]
-                instance = model_cls.from_db(db, model_init_names, model_init_values)
-                if annotation_fields:
-                    for column, pos in annotation_fields:
-                        setattr(instance, column, values[pos])
-                yield instance
-        finally:
-            query.cursor.close()
-
-    def _make_result(self, query):
-        for result in query:
-            yield list(result.values())

django_mongodb/queryset.py

@@ -0,0 +1,92 @@
+from itertools import chain
+
+from django.core.exceptions import FieldDoesNotExist
+from django.db import connections
+from django.db.models import QuerySet
+from django.db.models.query import RawModelIterable as BaseRawModelIterable
+from django.db.models.query import RawQuerySet as BaseRawQuerySet
+from django.db.models.sql.query import RawQuery as BaseRawQuery
+
+
+class MongoQuerySet(QuerySet):
+    def raw_mql(self, pipeline, using=None):
+        return RawQuerySet(pipeline, model=self.model, using=using)
+
+
+class RawQuery(BaseRawQuery):
+    def __init__(self, pipeline, using, model):
+        self.pipeline = pipeline
+        super().__init__(sql=None, using=using)
+        self.model = model
+
+    def _execute_query(self):
+        connection = connections[self.using]
+        collection = connection.get_collection(self.model._meta.db_table)
+        self.cursor = collection.aggregate(self.pipeline)
+
+    def __str__(self):
+        return str(self.pipeline)
+
+
+class RawQuerySet(BaseRawQuerySet):
+    def __init__(self, pipeline, model=None, using=None):
+        super().__init__(pipeline, model=model, using=using)
+        self.query = RawQuery(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 RawModelIterable(self)
+
+
+class RawModelIterable(BaseRawModelIterable):
+    def __iter__(self):
+        """
+        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)
+        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()
+            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]
+            converters = compiler.get_converters(
+                [f.get_col(f.model._meta.db_table) if f else None for f in fields]
+            )
+            if converters:
+                query_iterator = compiler.apply_converters(query_iterator, converters)
+            for values in query_iterator:
+                # Associate fields to values
+                model_init_values = [values[pos] for pos in model_init_pos]
+                instance = model_cls.from_db(db, model_init_names, model_init_values)
+                if annotation_fields:
+                    for column, pos in annotation_fields:
+                        setattr(instance, column, values[pos])
+                yield instance
+        finally:
+            query.cursor.close()
+
+    def _make_result(self, query):
+        for result in query:
+            yield list(result.values())