Refactor for DRY in _get_encrypted_fields method

Author: aclark4life

django_mongodb_backend/schema.py

@@ -545,34 +545,42 @@ def _get_encrypted_fields(
         master_key = connection.settings_dict.get("KMS_CREDENTIALS", {}).get(kms_provider)
         client_encryption = getattr(self.connection, "client_encryption", None)
 
+        def _field_dict(bson_type, path, new_key_alt_name, queries=None):
+            """Helper to generate a dictionary for an encrypted field.
+            Included in parent function's scope to avoid passing parameters.
+            """
+            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": bson_type,
+                "path": path,
+                "keyId": data_key,
+            }
+            if queries:
+                field_dict["queries"] = queries
+            return field_dict
+
         field_list = []
 
         for field in fields:
             new_key_alt_name = f"{key_alt_name}.{field.column}"
             path = f"{path_prefix}.{field.column}" if path_prefix else field.column
 
-            # --- Embedded Single Document ---
-            if isinstance(field, EmbeddedModelField):
+            if isinstance(field, (EmbeddedModelField, EmbeddedModelArrayField)):
                 if getattr(field, "encrypted", False):
-                    # 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,
+                    bson_type = "object" if isinstance(field, EmbeddedModelField) else "array"
+                    field_list.append(
+                        _field_dict(
+                            bson_type, path, new_key_alt_name, getattr(field, "queries", None)
+                        )
                     )
-                    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
                     embedded_result = self._get_encrypted_fields(
                         field.embedded_model,
                         create_data_keys=create_data_keys,
@@ -581,58 +589,11 @@ def _get_encrypted_fields(
                     )
                     if embedded_result and embedded_result.get("fields"):
                         field_list.extend(embedded_result["fields"])
-                continue
-
-            # --- Array of Embedded Documents ---
-            if isinstance(field, EmbeddedModelArrayField):
-                if getattr(field, "encrypted", False):
-                    # Entire array contents encrypted - flat entry
-                    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": "array",
-                        "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 for fields inside array elements
-                    embedded_result = self._get_encrypted_fields(
-                        field.embedded_model,
-                        create_data_keys=create_data_keys,
-                        key_alt_name=new_key_alt_name,
-                        path_prefix=path,  # array prefix in path
-                    )
-                    if embedded_result and embedded_result.get("fields"):
-                        field_list.extend(embedded_result["fields"])
-                continue
-
-            # --- Leaf encrypted field ---
-            if getattr(field, "encrypted", False):
-                data_key = self._get_data_key(
-                    client_encryption,
-                    key_vault_collection,
-                    create_data_keys,
-                    kms_provider,
-                    master_key,
-                    new_key_alt_name,
+            elif getattr(field, "encrypted", False):
+                bson_type = field.db_type(connection)
+                field_list.append(
+                    _field_dict(bson_type, path, new_key_alt_name, getattr(field, "queries", None))
                 )
-                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