INTPYTHON-348 aggregate via raw_mql

Author: aclark4life

django_mongodb/manager.py

@@ -3,4 +3,5 @@
 
 
 class MongoManager(BaseManager.from_queryset(MongoQuerySet)):
-    pass
+    def get_query_set(self):
+        return MongoQuerySet(self.model, using=self._db)

django_mongodb/query.py

@@ -12,6 +12,7 @@
 from django.db.models.sql.constants import INNER, GET_ITERATOR_CHUNK_SIZE
 from django.db.models.sql.datastructures import Join
 from django.db.models.sql.where import AND, OR, XOR, ExtraWhere, NothingNode, WhereNode
+from django.db.models.sql import Query
 from django.utils.functional import cached_property
 from pymongo.errors import BulkWriteError, DuplicateKeyError, PyMongoError
 
@@ -314,280 +315,12 @@ def register_nodes():
 
 
 class MongoQuerySet(QuerySet):
-    def raw_mql(self, raw_query, params=(), translations=None, using=None):
-        if using is None:
-            using = self.db
-        qs = RawQuerySet(
-            raw_query,
-            model=self.model,
-            params=params,
-            translations=translations,
-            using=using,
-        )
-        return qs
+    def raw_mql(self, raw_query):
+        return QuerySet(self.model, RawQuery(self.model, raw_query))
 
 
-class RawQuerySet:
-    """
-    Provide an iterator which converts the results of raw SQL queries into
-    annotated model instances.
-    """
+class RawQuery(Query):
 
-    def __init__(
-        self,
-        raw_query,
-        model=None,
-        query=None,
-        params=(),
-        translations=None,
-        using=None,
-        hints=None,
-    ):
+    def __init__(self, model, raw_query):
+        super(RawQuery, self).__init__(model)
         self.raw_query = raw_query
-        self.model = model
-        self._db = using
-        self._hints = hints or {}
-        self.query = query or RawQuery(sql=raw_query, using=self.db, params=params)
-        self.params = params
-        self.translations = translations or {}
-        self._result_cache = None
-        self._prefetch_related_lookups = ()
-        self._prefetch_done = False
-
-    def resolve_model_init_order(self):
-        """Resolve the init field names and value positions."""
-        converter = connections[self.db].introspection.identifier_converter
-        model_init_fields = [
-            f for f in self.model._meta.fields if converter(f.column) in self.columns
-        ]
-        annotation_fields = [
-            (column, pos)
-            for pos, column in enumerate(self.columns)
-            if column not in self.model_fields
-        ]
-        model_init_order = [self.columns.index(converter(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
-
-    def prefetch_related(self, *lookups):
-        """Same as QuerySet.prefetch_related()"""
-        clone = self._clone()
-        if lookups == (None,):
-            clone._prefetch_related_lookups = ()
-        else:
-            clone._prefetch_related_lookups = clone._prefetch_related_lookups + lookups
-        return clone
-
-    def _prefetch_related_objects(self):
-        prefetch_related_objects(self._result_cache, *self._prefetch_related_lookups)
-        self._prefetch_done = True
-
-    def _clone(self):
-        """Same as QuerySet._clone()"""
-        c = self.__class__(
-            self.raw_query,
-            model=self.model,
-            query=self.query,
-            params=self.params,
-            translations=self.translations,
-            using=self._db,
-            hints=self._hints,
-        )
-        c._prefetch_related_lookups = self._prefetch_related_lookups[:]
-        return c
-
-    def _fetch_all(self):
-        if self._result_cache is None:
-            self._result_cache = list(self.iterator())
-        if self._prefetch_related_lookups and not self._prefetch_done:
-            self._prefetch_related_objects()
-
-    def __len__(self):
-        self._fetch_all()
-        return len(self._result_cache)
-
-    def __bool__(self):
-        self._fetch_all()
-        return bool(self._result_cache)
-
-    def __iter__(self):
-        self._fetch_all()
-        return iter(self._result_cache)
-
-    def __aiter__(self):
-        # Remember, __aiter__ itself is synchronous, it's the thing it returns
-        # that is async!
-        async def generator():
-            await sync_to_async(self._fetch_all)()
-            for item in self._result_cache:
-                yield item
-
-        return generator()
-
-    def iterator(self):
-        yield from RawModelIterable(self)
-
-    def __repr__(self):
-        return "<%s: %s>" % (self.__class__.__name__, self.query)
-
-    def __getitem__(self, k):
-        return list(self)[k]
-
-    @property
-    def db(self):
-        """Return the database used if this query is executed now."""
-        return self._db or router.db_for_read(self.model, **self._hints)
-
-    def using(self, alias):
-        """Select the database this RawQuerySet should execute against."""
-        return RawQuerySet(
-            self.raw_query,
-            model=self.model,
-            query=self.query.chain(using=alias),
-            params=self.params,
-            translations=self.translations,
-            using=alias,
-        )
-
-    @cached_property
-    def columns(self):
-        """
-        A list of model field names in the order they'll appear in the
-        query results.
-        """
-        columns = self.query.get_columns()
-        # Adjust any column names which don't match field names
-        for query_name, model_name in self.translations.items():
-            # Ignore translations for nonexistent column names
-            try:
-                index = columns.index(query_name)
-            except ValueError:
-                pass
-            else:
-                columns[index] = model_name
-        return columns
-
-    @cached_property
-    def model_fields(self):
-        """A dict mapping column names to model field names."""
-        converter = connections[self.db].introspection.identifier_converter
-        model_fields = {}
-        for field in self.model._meta.fields:
-            name, column = field.get_attname_column()
-            model_fields[converter(column)] = field
-        return model_fields
-
-
-class RawQuery:
-    """A single raw SQL query."""
-
-    def __init__(self, sql, using, params=()):
-        self.params = params
-        self.sql = sql
-        self.using = using
-        self.cursor = None
-
-        # Mirror some properties of a normal query so that
-        # the compiler can be used to process results.
-        self.low_mark, self.high_mark = 0, None  # Used for offset/limit
-        self.extra_select = {}
-        self.annotation_select = {}
-
-    def chain(self, using):
-        return self.clone(using)
-
-    def clone(self, using):
-        return RawQuery(self.sql, using, params=self.params)
-
-    def get_columns(self):
-        if self.cursor is None:
-            self._execute_query()
-        converter = connections[self.using].introspection.identifier_converter
-        return [converter(column_meta[0]) for column_meta in self.cursor.description]
-
-    def __iter__(self):
-        # Always execute a new query for a new iterator.
-        # This could be optimized with a cache at the expense of RAM.
-        self._execute_query()
-        if not connections[self.using].features.can_use_chunked_reads:
-            # If the database can't use chunked reads we need to make sure we
-            # evaluate the entire query up front.
-            result = list(self.cursor)
-        else:
-            result = self.cursor
-        return iter(result)
-
-    def __repr__(self):
-        return "<%s: %s>" % (self.__class__.__name__, self)
-
-    @property
-    def params_type(self):
-        if self.params is None:
-            return None
-        return dict if isinstance(self.params, Mapping) else tuple
-
-    def __str__(self):
-        if self.params_type is None:
-            return self.sql
-        return self.sql % self.params_type(self.params)
-
-    def _execute_query(self):
-        connection = connections[self.using]
-
-        # Adapt parameters to the database, as much as possible considering
-        # that the target type isn't known. See #17755.
-        params_type = self.params_type
-        adapter = connection.ops.adapt_unknown_value
-        if params_type is tuple:
-            params = tuple(adapter(val) for val in self.params)
-        elif params_type is dict:
-            params = {key: adapter(val) for key, val in self.params.items()}
-        elif params_type is None:
-            params = None
-        else:
-            raise RuntimeError("Unexpected params type: %s" % params_type)
-
-        self.cursor = connection.cursor()
-        self.cursor.execute(self.sql, params)
-
-
-class RawModelIterable(BaseIterable):
-    """
-    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.
-        db = self.queryset.db
-        query = self.queryset.query
-        connection = connections[db]
-        compiler = connection.ops.compiler("SQLCompiler")(query, connection, db)
-        query_iterator = iter(query)
-
-        try:
-            (
-                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 exceptions.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:
-            # Done iterating the Query. If it has its own cursor, close it.
-            if hasattr(query, "cursor") and query.cursor:
-                query.cursor.close()