diff --git a/pymongo/encryption.py b/pymongo/encryption.py index 274fce2b1..384d4f5e7 100644 --- a/pymongo/encryption.py +++ b/pymongo/encryption.py @@ -18,6 +18,7 @@ from __future__ import annotations import contextlib import enum import socket +import uuid import weakref from copy import deepcopy from typing import ( @@ -30,6 +31,7 @@ from typing import ( MutableMapping, Optional, Sequence, + Union, cast, ) @@ -759,7 +761,7 @@ class ClientEncryption(Generic[_DocumentType]): 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, @@ -767,6 +769,8 @@ class ClientEncryption(Generic[_DocumentType]): 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 @@ class ClientEncryption(Generic[_DocumentType]): 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 @@ class ClientEncryption(Generic[_DocumentType]): :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 @@ class ClientEncryption(Generic[_DocumentType]): 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 @@ class ClientEncryption(Generic[_DocumentType]): :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 index 04982fa9c..fb297854f 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -461,6 +461,14 @@ class TestExplicitSimple(EncryptionIntegrationTest): ) 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 @@ class TestExplicitSimple(EncryptionIntegrationTest): 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"))