add encryption-compatible aggregation wrap

Author: timgraham

django_mongodb_backend/base.py

@@ -244,13 +244,14 @@ def get_database(self):
 
     @cached_property
     def client_encryption(self):
-        auto_encryption_opts = self.connection._options.auto_encryption_opts
-        return ClientEncryption(
-            auto_encryption_opts._kms_providers,
-            auto_encryption_opts._key_vault_namespace,
-            self.connection,
-            self.connection.codec_options,
-        )
+        if auto_encryption_opts := self.connection._options.auto_encryption_opts:
+            return ClientEncryption(
+                auto_encryption_opts._kms_providers,
+                auto_encryption_opts._key_vault_namespace,
+                self.connection,
+                self.connection.codec_options,
+            )
+        return None
 
     @cached_property
     def database(self):

django_mongodb_backend/query.py

@@ -93,24 +93,41 @@ def get_pipeline(self):
         if self.aggregation_pipeline:
             pipeline.extend(self.aggregation_pipeline)
         if self.needs_wrap_aggregation:
-            # Add the aggregation stage for queries without a GROUP BY.
-            # e.g. SQL equivalent of SELECT avg(col) FROM table
-            pipeline.extend(
-                [
-                    # Workaround for https://jira.mongodb.org/browse/SERVER-114196:
-                    # $$NOW becomes unavailable after $unionWith, so it must be
-                    # stored beforehand to ensure it remains accessible later
-                    # in the pipeline.
-                    {"$addFields": {"__now": "$$NOW"}},
-                    # Add an empty extra document to handle default values on
-                    # empty results.
-                    {"$unionWith": {"pipeline": [{"$documents": [{}]}]}},
-                    # Limiting to one document ensures the original result
-                    # takes precedence when present, otherwise the injected
-                    # empty document is used.
-                    {"$limit": 1},
+            if self.compiler.connection.client_encryption:
+                pipeline = [
+                    {"$collStats": {}},
+                    {
+                        "$lookup": {
+                            "from": self.compiler.collection_name,
+                            "as": "wrapped",
+                            "pipeline": pipeline,
+                        }
+                    },
+                    {
+                        "$replaceWith": {
+                            "$cond": [{"$eq": ["$wrapped", []]}, {}, {"$first": "$wrapped"}]
+                        }
+                    },
                 ]
-            )
+            else:
+                # Add the aggregation stage for queries without a GROUP BY.
+                # e.g. SQL equivalent of SELECT avg(col) FROM table
+                pipeline.extend(
+                    [
+                        # Workaround for https://jira.mongodb.org/browse/SERVER-114196:
+                        # $$NOW becomes unavailable after $unionWith, so it
+                        # must be stored beforehand to ensure it remains
+                        # accessible later in the pipeline.
+                        {"$addFields": {"__now": "$$NOW"}},
+                        # Add an empty extra document to handle default values
+                        # on empty results.
+                        {"$unionWith": {"pipeline": [{"$documents": [{}]}]}},
+                        # Limiting to one document ensures the original result
+                        # takes precedence when present, otherwise the injected
+                        # empty document is used.
+                        {"$limit": 1},
+                    ]
+                )
         if self.project_fields:
             pipeline.append({"$project": self.project_fields})
         if self.combinator_pipeline: