Move client_encryption to cached property in base

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/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
@@ -483,9 +482,7 @@ def _create_collection(self, model):
         else:
             db.create_collection(db_table)
 
-    def _get_encrypted_fields(
-        self, model, create_data_keys=False, key_alt_name=None, client_encryption=None
-    ):
+    def _get_encrypted_fields(self, model, create_data_keys=False, key_alt_name=None):
         """
         Recursively collect encryption schema data for fields in a model.
 
@@ -502,15 +499,7 @@ def _get_encrypted_fields(
         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)
-
-        # Initialize ClientEncryption once
-        if client_encryption is None:
-            client_encryption = ClientEncryption(
-                auto_encryption_opts._kms_providers,
-                auto_encryption_opts._key_vault_namespace,
-                client,
-                client.codec_options,
-            )
+        client_encryption = self.connection.client_encryption
 
         field_list = []
 
@@ -551,7 +540,6 @@ def _get_encrypted_fields(
                     field.embedded_model,
                     create_data_keys=create_data_keys,
                     key_alt_name=new_path,
-                    client_encryption=client_encryption,
                 )
                 field_list.extend(embedded_result["fields"])
                 continue