Move create data keys back to get encrypted fields

Author: aclark4life

django_mongodb_backend/schema.py

@@ -488,36 +488,6 @@ 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 = key_vault_collection.find_one({"keyAltNames": key_alt_name})
-        if not key:
-            raise ValueError(
-                f"No key found in keyvault for keyAltName={key_alt_name}. "
-                "Run with '--create-data-keys' to create missing keys."
-            )
-        return key["_id"]
-
     def _get_encrypted_fields(
         self, model, create_data_keys=False, key_alt_name=None, path_prefix=None
     ):
@@ -564,14 +534,15 @@ def _get_encrypted_fields(
 
             if getattr(field, "encrypted", False):
                 bson_type = field.db_type(connection)
-                data_key = self._get_data_key(
-                    client_encryption,
-                    key_vault_collection,
-                    create_data_keys,
-                    kms_provider,
-                    master_key,
-                    new_key_alt_name,
-                )
+                if create_data_keys:
+                    data_key = client_encryption.create_data_key(
+                        kms_provider=kms_provider,
+                        master_key=master_key,
+                        key_alt_names=[new_key_alt_name],
+                    )
+                else:
+                    key = key_vault_collection.find_one({"keyAltNames": new_key_alt_name})
+                    data_key = key["_id"]
                 field_dict = {
                     "bsonType": bson_type,
                     "path": path,