Add support for EncryptedEmbeddedModelArrayField

Author: aclark4life

django_mongodb_backend/fields/__init__.py

@@ -13,6 +13,7 @@
     EncryptedDecimalField,
     EncryptedDurationField,
     EncryptedEmailField,
+    EncryptedEmbeddedModelArrayField,
     EncryptedEmbeddedModelField,
     EncryptedFieldMixin,
     EncryptedFloatField,
@@ -44,6 +45,7 @@
     "EncryptedDecimalField",
     "EncryptedDurationField",
     "EncryptedEmailField",
+    "EncryptedEmbeddedModelArrayField",
     "EncryptedEmbeddedModelField",
     "EncryptedFieldMixin",
     "EncryptedFloatField",

django_mongodb_backend/fields/encryption.py

@@ -1,6 +1,10 @@
 from django.db import models
 
-from django_mongodb_backend.fields import EmbeddedModelField
+from django_mongodb_backend.fields import EmbeddedModelArrayField, EmbeddedModelField
+
+
+class EncryptedEmbeddedModelArrayField(EmbeddedModelArrayField):
+    encrypted = True
 
 
 class EncryptedEmbeddedModelField(EmbeddedModelField):

django_mongodb_backend/schema.py

@@ -8,7 +8,7 @@
 
 from django_mongodb_backend.indexes import SearchIndex
 
-from .fields import EmbeddedModelField
+from .fields import EmbeddedModelArrayField, EmbeddedModelField
 from .gis.schema import GISSchemaEditor
 from .query import wrap_database_errors
 from .utils import OperationCollector, model_has_encrypted_fields
@@ -551,6 +551,7 @@ def _get_encrypted_fields(
             new_key_alt_name = f"{key_alt_name}.{field.column}"
             path = f"{path_prefix}.{field.column}" if path_prefix else field.column
 
+            # --- Embedded Single Document ---
             if isinstance(field, EmbeddedModelField):
                 if getattr(field, "encrypted", False):
                     # Entire embedded object encrypted
@@ -582,7 +583,39 @@ def _get_encrypted_fields(
                         field_list.extend(embedded_result["fields"])
                 continue
 
-            # Leaf encrypted field
+            # --- Array of Embedded Documents ---
+            if isinstance(field, EmbeddedModelArrayField):
+                if getattr(field, "encrypted", False):
+                    # Entire array contents encrypted - flat entry
+                    data_key = self._get_data_key(
+                        client_encryption,
+                        key_vault_collection,
+                        create_data_keys,
+                        kms_provider,
+                        master_key,
+                        new_key_alt_name,
+                    )
+                    field_dict = {
+                        "bsonType": "array",
+                        "path": path,
+                        "keyId": data_key,
+                    }
+                    if getattr(field, "queries", False):
+                        field_dict["queries"] = field.queries
+                    field_list.append(field_dict)
+                else:
+                    # Recurse into embedded model for fields inside array elements
+                    embedded_result = self._get_encrypted_fields(
+                        field.embedded_model,
+                        create_data_keys=create_data_keys,
+                        key_alt_name=new_key_alt_name,
+                        path_prefix=path,  # array prefix in path
+                    )
+                    if embedded_result and embedded_result.get("fields"):
+                        field_list.extend(embedded_result["fields"])
+                continue
+
+            # --- Leaf encrypted field ---
             if getattr(field, "encrypted", False):
                 data_key = self._get_data_key(
                     client_encryption,

tests/encryption_/models.py

@@ -11,6 +11,7 @@
     EncryptedDecimalField,
     EncryptedDurationField,
     EncryptedEmailField,
+    EncryptedEmbeddedModelArrayField,
     EncryptedEmbeddedModelField,
     EncryptedFloatField,
     EncryptedGenericIPAddressField,
@@ -32,6 +33,7 @@ class Meta:
         required_db_features = {"supports_queryable_encryption"}
 
 
+# Embedded models
 class Patient(EncryptedTestModel):
     patient_name = models.CharField(max_length=255)
     patient_id = models.BigIntegerField()
@@ -52,7 +54,23 @@ class Billing(EmbeddedModel):
     cc_number = models.CharField(max_length=20)
 
 
-# Equality-queryable fields
+# Embedded array models
+class Actor(EmbeddedModel):
+    name = models.CharField(max_length=100)
+
+
+class Movie(EncryptedTestModel):
+    title = models.CharField(max_length=200)
+    plot = models.TextField(blank=True)
+    runtime = models.IntegerField(default=0)
+    released = models.DateTimeField("release date", null=True, blank=True)
+    cast = EncryptedEmbeddedModelArrayField(Actor, null=True, blank=True)
+
+    def __str__(self):
+        return self.title
+
+
+# Equality-queryable field models
 class BinaryModel(EncryptedTestModel):
     value = EncryptedBinaryField(queries={"queryType": "equality"})
 
@@ -81,7 +99,7 @@ class URLModel(EncryptedTestModel):
     value = EncryptedURLField(max_length=500, queries={"queryType": "equality"})
 
 
-# Range-queryable fields (also support equality)
+# Range-queryable field models
 class BigIntegerModel(EncryptedTestModel):
     value = EncryptedBigIntegerField(queries={"queryType": "range"})
 

tests/encryption_/test_fields.py

@@ -4,6 +4,7 @@
 from django_mongodb_backend.fields import EncryptedCharField
 
 from .models import (
+    Actor,
     BigIntegerModel,
     Billing,
     BinaryModel,
@@ -17,6 +18,7 @@
     FloatModel,
     GenericIPAddressModel,
     IntegerModel,
+    Movie,
     Patient,
     PatientRecord,
     PositiveBigIntegerModel,
@@ -30,7 +32,7 @@
 from .test_base import EncryptionTestCase
 
 
-class PatientModelTests(EncryptionTestCase):
+class EncryptedEmbeddedModelTests(EncryptionTestCase):
     def setUp(self):
         self.billing = Billing(cc_type="Visa", cc_number="4111111111111111")
         self.patient_record = PatientRecord(ssn="123-45-6789", billing=self.billing)
@@ -45,6 +47,21 @@ def test_patient(self):
         self.assertEqual(patient.patient_record.billing.cc_number, "4111111111111111")
 
 
+class EncryptedEmbeddedModelArrayTests(EncryptionTestCase):
+    def setUp(self):
+        self.actor1 = Actor(name="Actor One")
+        self.actor2 = Actor(name="Actor Two")
+        self.movie = Movie.objects.create(
+            title="Sample Movie",
+            cast=[self.actor1, self.actor2],
+        )
+
+    def test_movie_actors(self):
+        self.assertEqual(len(self.movie.cast), 2)
+        self.assertEqual(self.movie.cast[0].name, "Actor One")
+        self.assertEqual(self.movie.cast[1].name, "Actor Two")
+
+
 class EncryptedFieldTests(EncryptionTestCase):
     def assertEquality(self, model_cls, val):
         model_cls.objects.create(value=val)