Code review fixes

- assertion methods now use camelCase
- document Django QuerySet limitations
- low, high, threshold args are now required kwargs
- Add $ to all shell prompts
- Add patient schema to management and schema tests
- Add recursive fields to schema
- Reorder patient models
- Add bill_amount field
- Set databases for management and schema tests
- Remove client arg from _get_encrypted_fields
- Move client_encryption to cached property in base
- Re-add test_base with EncryptionTestCase
- Factor out common attributes into base test case
- Factor out skipUnlessDBFeature into base test case
- Remove short variable assignments from test assertions
- Fix recursion

  - model check now recurses into embedded models to look for encrypted
    fields
  - fields list is now conditionally extended based on the result of the
    recursive call

- Test Patient not PatientRecord
  - Something is still wrong here. Both models pass the test …
- Test encryption_patient not encryption_patientrecord
  - still not right, but will follow up to clarify/fix

Author: aclark4life

django_mongodb_backend/base.py

@@ -10,6 +10,7 @@
 from django.utils.functional import cached_property
 from pymongo.collection import Collection
 from pymongo.driver_info import DriverInfo
+from pymongo.encryption import ClientEncryption
 from pymongo.mongo_client import MongoClient
 from pymongo.uri_parser import parse_uri
 
@@ -182,6 +183,17 @@ def get_database(self):
             return OperationDebugWrapper(self)
         return self.database
 
+    @cached_property
+    def client_encryption(self):
+        # Initialize ClientEncryption once
+        auto_encryption_opts = getattr(self.connection._options, "auto_encryption_opts", None)
+        return ClientEncryption(
+            auto_encryption_opts._kms_providers,
+            auto_encryption_opts._key_vault_namespace,
+            self.connection,
+            self.connection.codec_options,
+        )
+
     @cached_property
     def database(self):
         """Connect to the database the first time it's accessed."""

django_mongodb_backend/management/commands/showencryptedfieldsmap.py

@@ -35,14 +35,13 @@ def handle(self, *args, **options):
         db = options["database"]
         create_data_keys = options.get("create_data_keys", False)
         connection = connections[db]
-        client = connection.connection
         encrypted_fields_map = {}
         with connection.schema_editor() as editor:
             for app_config in apps.get_app_configs():
                 for model in router.get_migratable_models(app_config, db):
                     if model_has_encrypted_fields(model):
                         fields = editor._get_encrypted_fields(
-                            model, client, create_data_keys=create_data_keys
+                            model, create_data_keys=create_data_keys
                         )
                         encrypted_fields_map[model._meta.db_table] = fields
         self.stdout.write(json_util.dumps(encrypted_fields_map, indent=2))

django_mongodb_backend/schema.py

@@ -4,7 +4,6 @@
 from django.db import router
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 from django.db.models import Index, UniqueConstraint
-from pymongo.encryption import ClientEncryption
 from pymongo.operations import SearchIndexModel
 
 from django_mongodb_backend.indexes import SearchIndex
@@ -457,79 +456,147 @@ def wait_until_index_dropped(collection, index_name, timeout=60, interval=0.5):
 
     def _create_collection(self, model):
         """
-        Create a collection for the model with the encrypted fields. If
-        provided, use the `_encrypted_fields_map` in the client's
-        `auto_encryption_opts`. Otherwise, create the encrypted fields map
-        with `_get_encrypted_fields`.
+        Create a collection for the model.
+        If the model has encrypted fields, build (or retrieve) the encrypted_fields schema.
         """
         db = self.get_database()
         db_table = model._meta.db_table
+
         if model_has_encrypted_fields(model):
+            # Encrypted path
             client = self.connection.connection
             auto_encryption_opts = getattr(client._options, "auto_encryption_opts", None)
             if not auto_encryption_opts:
                 raise ImproperlyConfigured(
                     f"Encrypted fields found but DATABASES['{self.connection.alias}']['OPTIONS'] "
                     "is missing auto_encryption_opts."
                 )
+
             encrypted_fields_map = getattr(auto_encryption_opts, "_encrypted_fields_map", None)
+
             if not encrypted_fields_map:
-                encrypted_fields = self._get_encrypted_fields(model, client, create_data_keys=True)
+                encrypted_fields = self._get_encrypted_fields(model, create_data_keys=True)
             else:
-                # If the encrypted fields map is provided, get the encrypted fields for the
-                # specific collection.
                 encrypted_fields = encrypted_fields_map.get(db_table)
-            db.create_collection(db_table, encryptedFields=encrypted_fields)
+
+            if encrypted_fields and encrypted_fields.get("fields"):
+                db.create_collection(db_table, encryptedFields=encrypted_fields)
+            else:
+                db.create_collection(db_table)
+
         else:
+            # Unencrypted path
             db.create_collection(db_table)
 
-    def _get_encrypted_fields(self, model, client, create_data_keys=False):
+    def _get_encrypted_fields(self, model, create_data_keys=False, key_alt_name=None):
+        """
+        Recursively collect encryption schema data for only encrypted fields in a model.
+        Returns None if no encrypted fields are found anywhere in the model hierarchy.
+
+        key_alt_name is the base path for this level, typically model._meta.db_table.
+        """
         connection = self.connection
+        client = connection.connection
         fields = model._meta.fields
+        key_alt_name = key_alt_name or model._meta.db_table
+
         options = client._options
-        auto_encryption_opts = options.auto_encryption_opts
+        auto_encryption_opts = getattr(options, "auto_encryption_opts", None)
+
+        key_vault_collection = None
+        if 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 = 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
+        master_key = connection.settings_dict.get("KMS_CREDENTIALS", {}).get(kms_provider)
+        client_encryption = getattr(self.connection, "client_encryption", None)
+
         field_list = []
+
         for field in fields:
+            new_path = f"{key_alt_name}.{field.column}"
+
+            # --- EmbeddedModelField ---
             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):
+                    # Entire sub-object is encrypted
+                    if create_data_keys:
+                        if not client_encryption:
+                            raise ImproperlyConfigured("client_encryption is not configured.")
+                        data_key = client_encryption.create_data_key(
+                            kms_provider=kms_provider,
+                            master_key=master_key,
+                            key_alt_names=[new_path],
+                        )
+                    else:
+                        if key_vault_collection is None:
+                            raise ImproperlyConfigured(
+                                f"Encrypted field {new_path} detected but no key vault configured"
+                            )
+                        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 = {
+                        "bsonType": "object",
+                        "path": field.column,
+                        "keyId": data_key,
+                    }
+                    if getattr(field, "queries", False):
+                        field_dict["queries"] = field.queries
+
+                    field_list.append(field_dict)
+                else:
+                    # Not encrypting whole object — recurse first then
+                    # conditionally extend field list
+                    embedded_result = self._get_encrypted_fields(
+                        field.embedded_model,
+                        create_data_keys=create_data_keys,
+                        key_alt_name=new_path,
+                    )
+                    if embedded_result and embedded_result.get("fields"):
+                        field_list.extend(embedded_result["fields"])
+                continue
+
+            # --- Leaf encrypted field ---
             if getattr(field, "encrypted", False):
-                key_alt_name = f"{db_table}.{field.column}"
                 if create_data_keys:
+                    if not client_encryption:
+                        raise ImproperlyConfigured("client_encryption is not configured.")
                     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],
                     )
                 else:
-                    key_doc = key_vault_collection.find_one({"keyAltNames": key_alt_name})
+                    if key_vault_collection is None:
+                        raise ImproperlyConfigured(
+                            f"Encrypted field {new_path} detected but no key vault configured"
+                        )
+                    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}
+
+        return {"fields": field_list} if field_list else None
 
 
 # GISSchemaEditor extends some SchemaEditor methods.

django_mongodb_backend/utils.py

@@ -189,4 +189,20 @@ def wrapper(self, *args, **kwargs):
 
 
 def model_has_encrypted_fields(model):
-    return any(getattr(field, "encrypted", False) for field in model._meta.fields)
+    """
+    Recursively check if this model or any embedded models contain encrypted fields.
+    Returns True if encryption is found anywhere in the hierarchy.
+    """
+    from django_mongodb_backend.fields import EmbeddedModelField  # noqa: PLC0415
+
+    for field in model._meta.fields:
+        if getattr(field, "encrypted", False):
+            return True
+
+        # Recursively check embedded models.
+        if isinstance(field, EmbeddedModelField) and model_has_encrypted_fields(
+            field.embedded_model
+        ):
+            return True
+
+    return False

docs/howto/queryable-encryption.rst

@@ -185,7 +185,7 @@ the :djadmin:`showencryptedfieldsmap` command.
 To see the keys created by Django MongoDB Backend in the above scenario, you can
 run the following command::
 
-    python manage.py showencryptedfieldsmap --database encrypted
+    $ python manage.py showencryptedfieldsmap --database encrypted
 
 You can then use the output of the :djadmin:`showencryptedfieldsmap` command
 to set the ``encrypted_fields_map`` in
@@ -202,7 +202,7 @@ pre-defined encrypted fields map.
 If you do not want to use the data keys created by Django MongoDB Backend (when
 ``python manage.py migrate`` is run), you can generate new data keys with::
 
-    python manage.py showencryptedfieldsmap --database encrypted \
+    $ python manage.py showencryptedfieldsmap --database encrypted \
         --create-data-keys
 
 In this scenario, Django MongoDB Backend will use the newly created data keys