PYTHON-2267: Allow UUID key_id to be passed to ClientEncryption.encry…

…pt (mongodb#1494)
blink1073 · Feb 2, 2024 · da2bf9d · da2bf9d
1 parent c2af3df
commit da2bf9d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/pymongo/encryption.py b/pymongo/encryption.py
@@ -18,6 +18,7 @@
 import contextlib
 import enum
 import socket
+import uuid
 import weakref
 from copy import deepcopy
 from typing import (
@@ -30,6 +31,7 @@
     MutableMapping,
     Optional,
     Sequence,
+    Union,
     cast,
 )
 
@@ -759,14 +761,16 @@ def _encrypt_helper(
         self,
         value: Any,
         algorithm: str,
-        key_id: Optional[Binary] = None,
+        key_id: Optional[Union[Binary, uuid.UUID]] = None,
         key_alt_name: Optional[str] = None,
         query_type: Optional[str] = None,
         contention_factor: Optional[int] = None,
         range_opts: Optional[RangeOpts] = None,
         is_expression: bool = False,
     ) -> Any:
         self._check_closed()
+        if isinstance(key_id, uuid.UUID):
+            key_id = Binary.from_uuid(key_id)
         if key_id is not None and not (
             isinstance(key_id, Binary) and key_id.subtype == UUID_SUBTYPE
         ):
@@ -799,7 +803,7 @@ def encrypt(
         self,
         value: Any,
         algorithm: str,
-        key_id: Optional[Binary] = None,
+        key_id: Optional[Union[Binary, uuid.UUID]] = None,
         key_alt_name: Optional[str] = None,
         query_type: Optional[str] = None,
         contention_factor: Optional[int] = None,
@@ -826,6 +830,9 @@ def encrypt(
 
         :return: The encrypted value, a :class:`~bson.binary.Binary` with subtype 6.
 
+        .. versionchanged:: 4.7
+            ``key_id`` can now be passed in as a :class:`uuid.UUID`.
+
         .. versionchanged:: 4.2
            Added the `query_type` and `contention_factor` parameters.
         """
@@ -847,7 +854,7 @@ def encrypt_expression(
         self,
         expression: Mapping[str, Any],
         algorithm: str,
-        key_id: Optional[Binary] = None,
+        key_id: Optional[Union[Binary, uuid.UUID]] = None,
         key_alt_name: Optional[str] = None,
         query_type: Optional[str] = None,
         contention_factor: Optional[int] = None,
@@ -875,6 +882,9 @@ def encrypt_expression(
 
         :return: The encrypted expression, a :class:`~bson.RawBSONDocument`.
 
+        .. versionchanged:: 4.7
+            ``key_id`` can now be passed in as a :class:`uuid.UUID`.
+
         .. versionadded:: 4.4
         """
         return cast(

diff --git a/test/test_encryption.py b/test/test_encryption.py
@@ -461,6 +461,14 @@ def test_encrypt_decrypt(self):
         )
         self.assertEqual(encrypted_ssn, encrypted_ssn2)
 
+        # Test encryption via UUID
+        encrypted_ssn3 = client_encryption.encrypt(
+            doc["ssn"],
+            Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+            key_id=key_id.as_uuid(),
+        )
+        self.assertEqual(encrypted_ssn, encrypted_ssn3)
+
         # Test decryption.
         decrypted_ssn = client_encryption.decrypt(encrypted_ssn)
         self.assertEqual(decrypted_ssn, doc["ssn"])
@@ -479,9 +487,6 @@ def test_validation(self):
 
         msg = "key_id must be a bson.binary.Binary with subtype 4"
         algo = Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic
-        uid = uuid.uuid4()
-        with self.assertRaisesRegex(TypeError, msg):
-            client_encryption.encrypt("str", algo, key_id=uid)  # type: ignore[arg-type]
         with self.assertRaisesRegex(TypeError, msg):
             client_encryption.encrypt("str", algo, key_id=Binary(b"123"))