make QuerySet.extra(select=... and where=...) raise NotSupportedError

- PR review fixes

Author: aclark4life

django_mongodb/expressions.py

@@ -22,7 +22,6 @@
     When,
 )
 from django.db.models.sql import Query
-from django.db.models.sql.where import ExtraWhere
 
 
 def case(self, compiler, connection):
@@ -68,10 +67,6 @@ def expression_wrapper(self, compiler, connection):
     return self.expression.as_mql(compiler, connection)
 
 
-def extra_where(self, compiler, connection):  # noqa: ARG001
-    raise NotSupportedError("QuerySet extras are not supported on MongoDB.")
-
-
 def f(self, compiler, connection):  # noqa: ARG001
     return f"${self.name}"
 
@@ -87,8 +82,10 @@ def order_by(self, compiler, connection):
 def query(self, compiler, connection):  # noqa: ARG001
     raise NotSupportedError("Using a QuerySet in annotate() is not supported on MongoDB.")
 
+
 def raw_sql(self, compiler, connection):  # noqa: ARG001
-    raise NotSupportedError("QuerySet extras are not supported on MongoDB.")
+    raise NotSupportedError("RawSQL is not supported on MongoDB.")
+
 
 def ref(self, compiler, connection):  # noqa: ARG001
     prefix = (
@@ -133,7 +130,6 @@ def register_expressions():
     Case.as_mql = case
     Col.as_mql = col
     CombinedExpression.as_mql = combined_expression
-    ExtraWhere.as_mql = extra_where
     ExpressionWrapper.as_mql = expression_wrapper
     F.as_mql = f
     NegatedExpression.as_mql = negated_expression

django_mongodb/features.py

@@ -493,6 +493,11 @@ def django_test_expected_failures(self):
             "queries.tests.ValuesQuerysetTests.test_named_values_list_with_fields",
             "queries.tests.ValuesQuerysetTests.test_named_values_list_without_fields",
             "select_related.tests.SelectRelatedTests.test_select_related_with_extra",
+            "aggregation_regress.tests.AggregationTests.test_annotate_with_extra",
+            "aggregation_regress.tests.AggregationTests.test_annotation",
+            "aggregation_regress.tests.AggregationTests.test_more_more3",
+            "aggregation_regress.tests.AggregationTests.test_more_more_more3",
+            "many_to_one.tests.ManyToOneTests.test_selects",
         },
         "Test inspects query for SQL": {
             "aggregation.tests.AggregateAnnotationPruningTests.test_non_aggregate_annotation_pruned",
@@ -664,17 +669,6 @@ def django_test_expected_failures(self):
             "backends.base.test_base.DatabaseWrapperLoggingTests",
             "migrations.test_operations.OperationTests.test_run_python_atomic",
         },
-        "QuerySet extras are not supported on MongoDB": {
-            # QuerySet.extra(select=...) raises NotSupportedError instead of:
-            # 'RawSQL' object has no attribute 'as_mql'.
-            "aggregation_regress.tests.AggregationTests.test_annotate_with_extra",
-            "aggregation_regress.tests.AggregationTests.test_annotation",
-            "aggregation_regress.tests.AggregationTests.test_more_more3",
-            "aggregation_regress.tests.AggregationTests.test_more_more_more3",
-            # QuerySet.extra(where=...) should raise NotSupportedError instead of:
-            # 'ExtraWhere' object has no attribute 'as_mql'.
-            "many_to_one.tests.ManyToOneTests.test_selects",
-        },
     }
 
     @cached_property

django_mongodb/query.py

@@ -2,13 +2,13 @@
 from operator import add as add_operator
 
 from django.core.exceptions import EmptyResultSet, FullResultSet
-from django.db import DatabaseError, IntegrityError
+from django.db import DatabaseError, IntegrityError, NotSupportedError
 from django.db.models.expressions import Case, When
 from django.db.models.functions import Mod
 from django.db.models.lookups import Exact
 from django.db.models.sql.constants import INNER
 from django.db.models.sql.datastructures import Join
-from django.db.models.sql.where import AND, OR, XOR, NothingNode, WhereNode
+from django.db.models.sql.where import AND, OR, XOR, ExtraWhere, NothingNode, WhereNode
 from pymongo.errors import DuplicateKeyError, PyMongoError
 
 
@@ -231,7 +231,12 @@ def where_node(self, compiler, connection):
     return mql
 
 
+def extra_where(self, compiler, connection):  # noqa: ARG001
+    raise NotSupportedError("QuerySet.extra() is not supported on MongoDB.")
+
+
 def register_nodes():
     Join.as_mql = join
     NothingNode.as_mql = NothingNode.as_sql
     WhereNode.as_mql = where_node
+    ExtraWhere.as_mql = extra_where