Factor _get_data_key from _get_encrypted_fields

Author: aclark4life

django_mongodb_backend/schema.py

@@ -490,15 +490,42 @@ def _create_collection(self, model):
             # Unencrypted path
             db.create_collection(db_table)
 
+    def _get_data_key(
+        self,
+        client_encryption,
+        key_vault_collection,
+        create_data_keys,
+        kms_provider,
+        master_key,
+        key_alt_name,
+    ):
+        """Return an existing or newly-created data key ID for a field."""
+        if create_data_keys:
+            if not client_encryption:
+                raise ImproperlyConfigured("client_encryption is not configured.")
+            return client_encryption.create_data_key(
+                kms_provider=kms_provider,
+                master_key=master_key,
+                key_alt_names=[key_alt_name],
+            )
+        if key_vault_collection is None:
+            raise ImproperlyConfigured(
+                f"Encrypted field {key_alt_name} detected but no key vault configured"
+            )
+        key_doc = key_vault_collection.find_one({"keyAltNames": key_alt_name})
+        if not key_doc:
+            raise ValueError(
+                f"No key found in keyvault for keyAltName={key_alt_name}. "
+                "Run with '--create-data-keys' to create missing keys."
+            )
+        return key_doc["_id"]
+
     def _get_encrypted_fields(
         self, model, create_data_keys=False, key_alt_name=None, path_prefix=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 used for keyAltNames.
-        path_prefix is the dot-notated path inside the document for schema mapping.
         """
         connection = self.connection
         client = connection.connection
@@ -524,40 +551,24 @@ 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
 
-            # --- EmbeddedModelField ---
             if isinstance(field, EmbeddedModelField):
                 if getattr(field, "encrypted", False):
-                    # Entire sub-object 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_key_alt_name],
-                        )
-                    else:
-                        if key_vault_collection is None:
-                            raise ImproperlyConfigured(
-                                f"Encrypted field {new_key_alt_name} detected "
-                                "but no key vault configured"
-                            )
-                        key_doc = key_vault_collection.find_one({"keyAltNames": new_key_alt_name})
-                        if not key_doc:
-                            raise ValueError(
-                                f"No key found in keyvault for keyAltName={new_key_alt_name}. "
-                                "Run with '--create-data-keys' to create missing keys."
-                            )
-                        data_key = key_doc["_id"]
-
+                    # Entire embedded object encrypted
+                    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": "object",
                         "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
@@ -571,38 +582,23 @@ def _get_encrypted_fields(
                         field_list.extend(embedded_result["fields"])
                 continue
 
-            # --- Leaf encrypted field ---
+            # Leaf encrypted field
             if getattr(field, "encrypted", False):
-                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_key_alt_name],
-                    )
-                else:
-                    if key_vault_collection is None:
-                        raise ImproperlyConfigured(
-                            f"Encrypted field {new_key_alt_name} detected "
-                            "but no key vault configured"
-                        )
-                    key_doc = key_vault_collection.find_one({"keyAltNames": new_key_alt_name})
-                    if not key_doc:
-                        raise ValueError(
-                            f"No key found in keyvault for keyAltName={new_key_alt_name}. "
-                            "Run with '--create-data-keys' to create missing keys."
-                        )
-                    data_key = key_doc["_id"]
-
+                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": field.db_type(connection),
                     "path": path,
                     "keyId": data_key,
                 }
                 if getattr(field, "queries", False):
                     field_dict["queries"] = field.queries
-
                 field_list.append(field_dict)
 
         return {"fields": field_list} if field_list else None