Code review and other fixes

- Add patient schema to management and schema tests
- Use results from recursion
- Reorder patient models and add unencrypted bill amount field
- Specify databases for management and schema tests

Author: aclark4life

django_mongodb_backend/schema.py

@@ -483,52 +483,111 @@ def _create_collection(self, model):
         else:
             db.create_collection(db_table)
 
-    def _get_encrypted_fields(self, model, client, create_data_keys=False):
+    def _get_encrypted_fields(
+        self, model, client, create_data_keys=False, key_alt_name=None, client_encryption=None
+    ):
+        """
+        Recursively collect encryption schema data for fields in a model.
+
+        key_alt_name is the base path for this level, typically model._meta.db_table
+        """
         connection = self.connection
         fields = model._meta.fields
-        options = client._options
-        auto_encryption_opts = options.auto_encryption_opts
-        kms_provider = router.kms_provider(model)
-        master_key = self.connection.settings_dict.get("KMS_CREDENTIALS", {}).get(kms_provider)
-        client_encryption = ClientEncryption(
-            auto_encryption_opts._kms_providers,
-            auto_encryption_opts._key_vault_namespace,
-            client,
-            client.codec_options,
-        )
-        key_vault_db, key_vault_coll = auto_encryption_opts._key_vault_namespace.split(".", 1)
-        key_vault_collection = client[key_vault_db][key_vault_coll]
-        db_table = model._meta.db_table
+        key_alt_name = key_alt_name or model._meta.db_table
+
+        # Initialize ClientEncryption once
+        if client_encryption is None:
+            options = client._options
+            auto_encryption_opts = options.auto_encryption_opts
+            kms_provider = router.kms_provider(model)
+            master_key = connection.settings_dict.get("KMS_CREDENTIALS", {}).get(kms_provider)
+            client_encryption = ClientEncryption(
+                auto_encryption_opts._kms_providers,
+                auto_encryption_opts._key_vault_namespace,
+                client,
+                client.codec_options,
+            )
+            key_vault_db, key_vault_coll = auto_encryption_opts._key_vault_namespace.split(".", 1)
+            key_vault_collection = client[key_vault_db][key_vault_coll]
+        else:
+            auto_encryption_opts = client._options.auto_encryption_opts
+            key_vault_db, key_vault_coll = auto_encryption_opts._key_vault_namespace.split(".", 1)
+            key_vault_collection = client[key_vault_db][key_vault_coll]
+            kms_provider = router.kms_provider(model)
+            master_key = connection.settings_dict.get("KMS_CREDENTIALS", {}).get(kms_provider)
+
         field_list = []
+
         for field in fields:
+            new_path = f"{key_alt_name}.{field.column}"
+
+            # --- EmbeddedModelField case ---
             if isinstance(field, EmbeddedModelField):
-                # Recursively get encrypted fields for the embedded model.
-                self._get_encrypted_fields(field.embedded_model, client, create_data_keys)
+                field_dict = {"bsonType": "object", "path": field.column}
+
+                if getattr(field, "encrypted", False):
+                    if create_data_keys:
+                        data_key = client_encryption.create_data_key(
+                            kms_provider=kms_provider,
+                            master_key=master_key,
+                            key_alt_names=[new_path],
+                        )
+                    else:
+                        key_doc = key_vault_collection.find_one({"keyAltNames": new_path})
+                        if not key_doc:
+                            raise ValueError(
+                                f"No key found in keyvault for keyAltName={new_path}. "
+                                "Run with '--create-data-keys' to create missing keys."
+                            )
+                        data_key = key_doc["_id"]
+
+                    field_dict["keyId"] = data_key
+
+                    if getattr(field, "queries", False):
+                        field_dict["queries"] = field.queries
+
+                    field_list.append(field_dict)
+                    continue
+
+                # Not encrypting whole object — add object entry and recurse
+                field_list.append(field_dict)
+                embedded_result = self._get_encrypted_fields(
+                    field.embedded_model,
+                    client,
+                    create_data_keys=create_data_keys,
+                    key_alt_name=new_path,
+                    client_encryption=client_encryption,
+                )
+                field_list.extend(embedded_result["fields"])
+                continue
+
+            # --- Leaf encrypted field case ---
             if getattr(field, "encrypted", False):
-                key_alt_name = f"{db_table}.{field.column}"
                 if create_data_keys:
                     data_key = client_encryption.create_data_key(
                         kms_provider=kms_provider,
                         master_key=master_key,
-                        key_alt_names=[key_alt_name],
+                        key_alt_names=[new_path],  # distinct per field
                     )
                 else:
-                    key_doc = key_vault_collection.find_one({"keyAltNames": key_alt_name})
+                    key_doc = key_vault_collection.find_one({"keyAltNames": new_path})
                     if not key_doc:
                         raise ValueError(
-                            f"No key found in keyvault for keyAltName={key_alt_name}. "
-                            "You may need to run the management command with "
-                            "'--create-data-keys' to create missing keys."
+                            f"No key found in keyvault for keyAltName={new_path}. "
+                            "Run with '--create-data-keys' to create missing keys."
                         )
                     data_key = key_doc["_id"]
+
                 field_dict = {
                     "bsonType": field.db_type(connection),
                     "path": field.column,
                     "keyId": data_key,
                 }
                 if getattr(field, "queries", False):
                     field_dict["queries"] = field.queries
+
                 field_list.append(field_dict)
+
         return {"fields": field_list}
 
 

tests/encryption_/models.py

@@ -26,26 +26,27 @@
 from django_mongodb_backend.models import EmbeddedModel
 
 
-class Billing(EmbeddedModel):
-    cc_type = models.CharField(max_length=50)
-    cc_number = models.CharField(max_length=20)
-
-
-class PatientRecord(EmbeddedModel):
-    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 = EmbeddedModelField("PatientRecord")
 
     def __str__(self):
         return f"{self.patient_name} ({self.patient_id})"
 
 
-class EncryptedModel(models.Model):
+class PatientRecord(EmbeddedModel):
+    ssn = EncryptedCharField(max_length=11, queries={"queryType": "equality"})
+    billing = EncryptedEmbeddedModelField("Billing")
+    bill_amount = models.DecimalField(max_digits=10, decimal_places=2)
+
+
+class Billing(EmbeddedModel):
+    cc_type = models.CharField(max_length=50)
+    cc_number = models.CharField(max_length=20)
+
+
+class EncryptedModelBase(models.Model):
     """
     Abstract base model for all Encrypted models
     that require the 'supports_queryable_encryption' DB feature.
@@ -57,78 +58,78 @@ class Meta:
 
 
 # Equality-queryable fields
-class EncryptedBinaryTest(EncryptedModel):
+class EncryptedBinaryTest(EncryptedModelBase):
     value = EncryptedBinaryField(queries={"queryType": "equality"})
 
 
-class EncryptedBooleanTest(EncryptedModel):
+class EncryptedBooleanTest(EncryptedModelBase):
     value = EncryptedBooleanField(queries={"queryType": "equality"})
 
 
-class EncryptedCharTest(EncryptedModel):
+class EncryptedCharTest(EncryptedModelBase):
     value = EncryptedCharField(max_length=255, queries={"queryType": "equality"})
 
 
-class EncryptedEmailTest(EncryptedModel):
+class EncryptedEmailTest(EncryptedModelBase):
     value = EncryptedEmailField(max_length=255, queries={"queryType": "equality"})
 
 
-class EncryptedGenericIPAddressTest(EncryptedModel):
+class EncryptedGenericIPAddressTest(EncryptedModelBase):
     value = EncryptedGenericIPAddressField(queries={"queryType": "equality"})
 
 
-class EncryptedTextTest(EncryptedModel):
+class EncryptedTextTest(EncryptedModelBase):
     value = EncryptedTextField(queries={"queryType": "equality"})
 
 
-class EncryptedURLTest(EncryptedModel):
+class EncryptedURLTest(EncryptedModelBase):
     value = EncryptedURLField(max_length=500, queries={"queryType": "equality"})
 
 
 # Range-queryable fields (also support equality)
-class EncryptedBigIntegerTest(EncryptedModel):
+class EncryptedBigIntegerTest(EncryptedModelBase):
     value = EncryptedBigIntegerField(queries={"queryType": "range"})
 
 
-class EncryptedDateTest(EncryptedModel):
+class EncryptedDateTest(EncryptedModelBase):
     value = EncryptedDateField(queries={"queryType": "range"})
 
 
-class EncryptedDateTimeTest(EncryptedModel):
+class EncryptedDateTimeTest(EncryptedModelBase):
     value = EncryptedDateTimeField(queries={"queryType": "range"})
 
 
-class EncryptedDecimalTest(EncryptedModel):
+class EncryptedDecimalTest(EncryptedModelBase):
     value = EncryptedDecimalField(max_digits=10, decimal_places=2, queries={"queryType": "range"})
 
 
-class EncryptedDurationTest(EncryptedModel):
+class EncryptedDurationTest(EncryptedModelBase):
     value = EncryptedDurationField(queries={"queryType": "range"})
 
 
-class EncryptedFloatTest(EncryptedModel):
+class EncryptedFloatTest(EncryptedModelBase):
     value = EncryptedFloatField(queries={"queryType": "range"})
 
 
-class EncryptedIntegerTest(EncryptedModel):
+class EncryptedIntegerTest(EncryptedModelBase):
     value = EncryptedIntegerField(queries={"queryType": "range"})
 
 
-class EncryptedPositiveBigIntegerTest(EncryptedModel):
+class EncryptedPositiveBigIntegerTest(EncryptedModelBase):
     value = EncryptedPositiveBigIntegerField(queries={"queryType": "range"})
 
 
-class EncryptedPositiveIntegerTest(EncryptedModel):
+class EncryptedPositiveIntegerTest(EncryptedModelBase):
     value = EncryptedPositiveIntegerField(queries={"queryType": "range"})
 
 
-class EncryptedPositiveSmallIntegerTest(EncryptedModel):
+class EncryptedPositiveSmallIntegerTest(EncryptedModelBase):
     value = EncryptedPositiveSmallIntegerField(queries={"queryType": "range"})
 
 
-class EncryptedSmallIntegerTest(EncryptedModel):
+class EncryptedSmallIntegerTest(EncryptedModelBase):
     value = EncryptedSmallIntegerField(queries={"queryType": "range"})
 
 
-class EncryptedTimeTest(EncryptedModel):
+class EncryptedTimeTest(EncryptedModelBase):
     value = EncryptedTimeField(queries={"queryType": "range"})

tests/encryption_/test_management.py

@@ -8,10 +8,17 @@
 @skipUnlessDBFeature("supports_queryable_encryption")
 @modify_settings(INSTALLED_APPS={"prepend": "django_mongodb_backend"})
 class CommandTests(TestCase):
+    databases = {"default", "encrypted"}
     maxDiff = None
 
     # Expected encrypted field maps for all Encrypted* models
     expected_maps = {
+        "encryption__patientrecord": {
+            "fields": [
+                {"bsonType": "string", "path": "ssn", "queries": {"queryType": "equality"}},
+                {"bsonType": "object", "path": "billing"},
+            ]
+        },
         # Equality-queryable fields
         "encryption__encryptedbinarytest": {
             "fields": [

tests/encryption_/test_schema.py

@@ -1,14 +1,21 @@
 from django.db import connections
 from django.test import TestCase
 
-from . import models  # your encryption_ models file with Encrypted*Test classes
+from . import models
 
 
 class SchemaTests(TestCase):
+    databases = {"default", "encrypted"}
     maxDiff = None
 
     # Expected encrypted fields map per model
     expected_map = {
+        "PatientRecord": {
+            "fields": [
+                {"bsonType": "string", "path": "ssn", "queries": {"queryType": "equality"}},
+                {"bsonType": "object", "path": "billing"},
+            ]
+        },
         "EncryptedBinaryTest": {
             "fields": [
                 {"bsonType": "binData", "path": "value", "queries": {"queryType": "equality"}}