fix QuerySet.update() with multi-table inheritance

Author: timgraham

django_mongodb/compiler.py

@@ -1,10 +1,12 @@
+from itertools import chain
+
 from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import DatabaseError, IntegrityError, NotSupportedError
 from django.db.models import Count, Expression
 from django.db.models.aggregates import Aggregate
 from django.db.models.expressions import OrderBy
 from django.db.models.sql import compiler
-from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, ORDER_DIR
+from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, ORDER_DIR, SINGLE
 from django.utils.functional import cached_property
 
 from .base import Cursor
@@ -33,11 +35,23 @@ def execute_sql(
         except EmptyResultSet:
             return iter([]) if result_type == MULTI else None
 
-        return (
-            (self._make_result(row, columns) for row in query.fetch())
-            if result_type == MULTI
-            else self._make_result(next(query.fetch()), columns)
-        )
+        cursor = query.get_cursor()
+
+        if result_type == SINGLE:
+            return self._make_result(next(cursor))
+
+        if result_type == MULTI:
+            result = self.cursor_iter(cursor, chunk_size, columns)
+
+            if not chunked_fetch or not self.connection.features.can_use_chunked_reads:
+                # If using non-chunked reads, return the same data structure as
+                # normally, but ensure it is all read into memory before going
+                # any further. Use chunked_fetch if requested,
+                # unless the database doesn't support it.
+                return list(result)
+            return result
+
+        raise ValueError("shouldn't get here")
 
     def results_iter(
         self,
@@ -49,14 +63,16 @@ def results_iter(
         """
         Return an iterator over the results from executing query given
         to this compiler. Called by QuerySet methods.
+
+        XXX: Can this method be removed?
         """
         if results is None:
             # QuerySet.values() or values_list()
             results = self.execute_sql(MULTI, chunked_fetch=chunked_fetch, chunk_size=chunk_size)
 
         fields = [s[0] for s in self.select[0 : self.col_count]]
         converters = self.get_converters(fields)
-        rows = results
+        rows = chain.from_iterable(results)
         if converters:
             rows = self.apply_converters(rows, converters)
         if tuple_expected:
@@ -86,6 +102,15 @@ def _make_result(self, entity, columns):
             result.append(obj.get(name))
         return result
 
+    def cursor_iter(self, cursor, chunk_size, columns):
+        """Generator to yield chunks from cursor."""
+        chunk = []
+        for i, row in enumerate(cursor):
+            if i % chunk_size == 0 and i > 0:
+                yield chunk
+            chunk.append(self._make_result(row, columns))
+        yield chunk
+
     def check_query(self):
         """Check if the current query is supported by the database."""
         if self.query.is_empty():
@@ -293,8 +318,14 @@ def check_query(self):
             )
 
 
-class SQLUpdateCompiler(SQLCompiler):
+class SQLUpdateCompiler(compiler.SQLUpdateCompiler, SQLCompiler):
     def execute_sql(self, result_type):
+        """
+        Execute the specified update. Return the number of rows affected by
+        the primary update query. The "primary update query" is the first
+        non-empty query that is executed. Row counts for any subsequent,
+        related queries are not available.
+        """
         self.pre_sql_setup()
         values = []
         for field, _, value in self.query.values:
@@ -309,7 +340,17 @@ def execute_sql(self, result_type):
                     )
             prepared = field.get_db_prep_save(value, connection=self.connection)
             values.append((field, prepared))
-        return self.update(values)
+        # if is_empty := not bool(values):
+        #     rows = 0
+        # else:
+        #     rows = self.update(values)
+        rows = 0 if (is_empty := not bool(values)) else self.update(values)
+        for query in self.query.get_related_updates():
+            aux_rows = query.get_compiler(self.using).execute_sql(result_type)
+            if is_empty and aux_rows:
+                rows = aux_rows
+                is_empty = False
+        return rows
 
     def update(self, values):
         spec = {}

django_mongodb/features.py

@@ -51,9 +51,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "update.tests.AdvancedTests.test_update_ordered_by_inline_m2m_annotation",
         "update.tests.AdvancedTests.test_update_ordered_by_m2m_annotation",
         "update.tests.AdvancedTests.test_update_ordered_by_m2m_annotation_desc",
-        # pymongo: ValueError: update cannot be empty
-        "update.tests.SimpleTest.test_empty_update_with_inheritance",
-        "update.tests.SimpleTest.test_nonempty_update_with_inheritance",
         # Pattern lookups that use regexMatch don't work on JSONField:
         # Unsupported conversion from array to string in $convert
         "model_fields.test_jsonfield.TestQuerying.test_icontains",

django_mongodb/query.py

@@ -53,9 +53,9 @@ def __init__(self, compiler, columns):
     def __repr__(self):
         return f"<MongoQuery: {self.mongo_query!r} ORDER {self.ordering!r}>"
 
-    def fetch(self):
-        """Return an iterator over the query results."""
-        yield from self.get_cursor()
+    #    def fetch(self):
+    #        """Return an iterator over the query results."""
+    #        yield from self.get_cursor()
 
     @wrap_database_errors
     def count(self, limit=None, skip=None):