Add embedding

- EncryptedEmbeddedModel for encrypted objects
- EmbeddedModel for models with encrypted fields

Author: aclark4life

django_mongodb_backend/fields/__init__.py

@@ -13,6 +13,7 @@
     EncryptedDecimalField,
     EncryptedDurationField,
     EncryptedEmailField,
+    EncryptedEmbeddedModelField,
     EncryptedFieldMixin,
     EncryptedFloatField,
     EncryptedGenericIPAddressField,
@@ -43,6 +44,7 @@
     "EncryptedDecimalField",
     "EncryptedDurationField",
     "EncryptedEmailField",
+    "EncryptedEmbeddedModelField",
     "EncryptedFieldMixin",
     "EncryptedFloatField",
     "EncryptedGenericIPAddressField",

django_mongodb_backend/fields/encryption.py

@@ -1,5 +1,14 @@
 from django.db import models
 
+from django_mongodb_backend.fields import EmbeddedModelField
+
+
+class EncryptedEmbeddedModelField(EmbeddedModelField):
+    encrypted = True
+
+    def db_type(self, connection):
+        return "object"
+
 
 class EncryptedFieldMixin:
     encrypted = True

django_mongodb_backend/management/commands/showencryptedfieldsmap.py

@@ -3,7 +3,7 @@
 from django.core.management.base import BaseCommand
 from django.db import DEFAULT_DB_ALIAS, connections, router
 
-from django_mongodb_backend.model_utils import model_has_encrypted_fields
+from django_mongodb_backend.utils import model_has_encrypted_fields
 
 
 class Command(BaseCommand):

django_mongodb_backend/model_utils.py

@@ -14,3 +14,11 @@ def delete(self, *args, **kwargs):
 
     def save(self, *args, **kwargs):
         raise NotSupportedError("EmbeddedModels cannot be saved.")
+
+
+class EncryptedEmbeddedModel(EmbeddedModel):
+    encrypted = True
+
+    class Meta:
+        abstract = True
+        required_db_features = {"supports_queryable_encryption"}

django_mongodb_backend/models.py

@@ -11,9 +11,8 @@
 
 from .fields import EmbeddedModelField
 from .gis.schema import GISSchemaEditor
-from .model_utils import model_has_encrypted_fields
 from .query import wrap_database_errors
-from .utils import OperationCollector
+from .utils import OperationCollector, model_has_encrypted_fields
 
 
 def ignore_embedded_models(func):
@@ -502,6 +501,9 @@ def _get_encrypted_fields(self, model, client, create_data_keys=False):
         db_table = model._meta.db_table
         field_list = []
         for field in fields:
+            if isinstance(field, EmbeddedModelField):
+                # Recursively get encrypted fields for the embedded model.
+                self._get_encrypted_fields(field.embedded_model, client, create_data_keys)
             if getattr(field, "encrypted", False):
                 key_alt_name = f"{db_table}.{field.column}"
                 if create_data_keys:
@@ -524,7 +526,7 @@ def _get_encrypted_fields(self, model, client, create_data_keys=False):
                     "path": field.column,
                     "keyId": data_key,
                 }
-                if field.queries:
+                if getattr(field, "queries", False):
                     field_dict["queries"] = field.queries
                 field_list.append(field_dict)
         return {"fields": field_list}

django_mongodb_backend/schema.py

@@ -186,3 +186,13 @@ def wrapper(self, *args, **kwargs):
             self.log(method, args, kwargs)
 
         return wrapper
+
+
+def model_has_encrypted_fields(model):
+    from django_mongodb_backend.models import EncryptedEmbeddedModel  # noqa: PLC0415
+
+    for field in model._meta.fields:
+        if getattr(field, "encrypted", False):
+            return True
+
+    return bool(issubclass(model, EncryptedEmbeddedModel))

django_mongodb_backend/utils.py

@@ -1,7 +1,6 @@
 from django.db import models
 
 from django_mongodb_backend.fields import (
-    EmbeddedModelField,
     EncryptedBigIntegerField,
     EncryptedBinaryField,
     EncryptedBooleanField,
@@ -11,6 +10,7 @@
     EncryptedDecimalField,
     EncryptedDurationField,
     EncryptedEmailField,
+    EncryptedEmbeddedModelField,
     EncryptedFloatField,
     EncryptedGenericIPAddressField,
     EncryptedIntegerField,
@@ -22,23 +22,23 @@
     EncryptedTimeField,
     EncryptedURLField,
 )
-from django_mongodb_backend.models import EmbeddedModel
+from django_mongodb_backend.models import EncryptedEmbeddedModel
 
 
-class Billing(EmbeddedModel):
+class Billing(EncryptedEmbeddedModel):
     cc_type = models.CharField(max_length=50)
     cc_number = models.CharField(max_length=20)
 
 
-class PatientRecord(EmbeddedModel):
-    ssn = models.CharField(max_length=11)
-    billing = EmbeddedModelField(Billing)
+class PatientRecord(EncryptedEmbeddedModel):
+    ssn = EncryptedCharField(max_length=11, queries={"queryType": "equality"})
+    billing = EncryptedEmbeddedModelField(Billing)
 
 
 class Patient(models.Model):
     patient_name = models.CharField(max_length=255)
     patient_id = models.BigIntegerField()
-    patient_record = EmbeddedModelField(PatientRecord)
+    patient_record = EncryptedEmbeddedModelField(PatientRecord)
 
     def __str__(self):
         return f"{self.patient_name} ({self.patient_id})"

tests/encryption_/models.py

@@ -29,22 +29,20 @@
 )
 
 
+@skipUnlessDBFeature("supports_queryable_encryption")
 class PatientModelTests(TestCase):
+    databases = {"default", "encrypted"}
+
     def setUp(self):
         self.billing = Billing(cc_type="Visa", cc_number="4111111111111111")
         self.patient_record = PatientRecord(ssn="123-45-6789", billing=self.billing)
         self.patient = Patient.objects.create(
             patient_name="John Doe", patient_id=123456789, patient_record=self.patient_record
         )
 
-    def test_patient_record_content(self):
-        """Embedded patient record data should be stored and retrieved correctly."""
+    def test_patient(self):
         patient = Patient.objects.get(id=self.patient.id)
         self.assertEqual(patient.patient_record.ssn, "123-45-6789")
-
-    def test_billing_information(self):
-        """Billing data inside the encrypted embedded model should be correct."""
-        patient = Patient.objects.get(id=self.patient.id)
         self.assertEqual(patient.patient_record.billing.cc_type, "Visa")
         self.assertEqual(patient.patient_record.billing.cc_number, "4111111111111111")