INTPYTHON-527 Add Queryable Encryption support

django_mongodb_backend/schema.py

@@ -8,7 +8,7 @@
 
 from django_mongodb_backend.indexes import SearchIndex
 
-from .fields import EmbeddedModelArrayField, EmbeddedModelField
+from .fields import EmbeddedModelField
 from .gis.schema import GISSchemaEditor
 from .query import wrap_database_errors
 from .utils import OperationCollector, model_has_encrypted_fields
@@ -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
     ):
@@ -532,26 +502,22 @@ def _get_encrypted_fields(
         path_prefix = path_prefix or ""
 
         options = client._options
-        auto_encryption_opts = getattr(options, "auto_encryption_opts", None)
+        auto_encryption_opts = options.auto_encryption_opts
 
-        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]
+        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 = connection.settings_dict.get("KMS_CREDENTIALS", {}).get(kms_provider)
-        client_encryption = getattr(self.connection, "client_encryption", None)
+        client_encryption = self.connection.client_encryption
 
         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
 
-            if isinstance(field, (EmbeddedModelField, EmbeddedModelArrayField)) and not getattr(
-                field, "encrypted", False
-            ):
+            if isinstance(field, EmbeddedModelField) and not getattr(field, "encrypted", False):
                 embedded_result = self._get_encrypted_fields(
                     field.embedded_model,
                     create_data_keys=create_data_keys,
@@ -564,14 +530,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,

tests/encryption_/test_fields.py

@@ -32,22 +32,22 @@
 from .test_base import EncryptionTestCase
 
 
-class EncryptedEmbeddedModelTests(EncryptionTestCase):
+class EmbeddedModelTests(EncryptionTestCase):
     def setUp(self):
         self.billing = Billing(cc_type="Visa", cc_number="4111111111111111")
         self.patient_record = PatientRecord(ssn="123-45-6789", billing=self.billing)
         self.patient = Patient.objects.create(
             patient_name="John Doe", patient_id=123456789, patient_record=self.patient_record
         )
 
-    def test_patient(self):
+    def test_object(self):
         patient = Patient.objects.get(id=self.patient.id)
         self.assertEqual(patient.patient_record.ssn, "123-45-6789")
         self.assertEqual(patient.patient_record.billing.cc_type, "Visa")
         self.assertEqual(patient.patient_record.billing.cc_number, "4111111111111111")
 
 
-class EncryptedEmbeddedModelArrayTests(EncryptionTestCase):
+class EmbeddedModelArrayTests(EncryptionTestCase):
     def setUp(self):
         self.actor1 = Actor(name="Actor One")
         self.actor2 = Actor(name="Actor Two")
@@ -56,13 +56,13 @@ def setUp(self):
             cast=[self.actor1, self.actor2],
         )
 
-    def test_movie_actors(self):
+    def test_array(self):
         self.assertEqual(len(self.movie.cast), 2)
         self.assertEqual(self.movie.cast[0].name, "Actor One")
         self.assertEqual(self.movie.cast[1].name, "Actor Two")
 
 
-class EncryptedFieldTests(EncryptionTestCase):
+class FieldTests(EncryptionTestCase):
     def assertEquality(self, model_cls, val):
         model_cls.objects.create(value=val)
         fetched = model_cls.objects.get(value=val)
@@ -162,7 +162,7 @@ def test_time(self):
         )
 
 
-class EncryptedFieldMixinTests(EncryptionTestCase):
+class FieldMixinTests(EncryptionTestCase):
     def test_null_true_raises_error(self):
         with self.assertRaisesMessage(
             ValueError, "'null=True' is not supported for encrypted fields."