PYTHON-5191 Add key_expiration_ms option for DEK cache lifetime (#984)

Author: ShaneHarvey

bindings/python/CHANGELOG.rst

@@ -1,6 +1,11 @@
 Changelog
 =========
 
+Changes in Version 1.13.0
+-------------------------
+
+- Add support for the key_expiration_ms option to MongoCryptOptions.
+
 Changes in Version 1.12.0
 -------------------------
 

bindings/python/pymongocrypt/mongocrypt.py

@@ -149,10 +149,15 @@ def __init(self):
         if any([on_demand_aws, on_demand_gcp, on_demand_azure]):
             lib.mongocrypt_setopt_use_need_kms_credentials_state(self.__crypt)
 
-        # Enable KMS retry when available, libmongocrypt >= 1.12.0,
+        # Enable KMS retry and key_expiration_ms when available, libmongocrypt >= 1.12.0,
         try:
             if not lib.mongocrypt_setopt_retry_kms(self.__crypt, True):
                 self.__raise_from_status()
+            if self.__opts.key_expiration_ms is not None:
+                if not lib.mongocrypt_setopt_key_expiration(
+                    self.__crypt, self.__opts.key_expiration_ms
+                ):
+                    self.__raise_from_status()
         except AttributeError:
             # libmongocrypt < 1.12
             pass

bindings/python/pymongocrypt/options.py

@@ -11,6 +11,7 @@ def __init__(
         crypt_shared_lib_path=None,
         crypt_shared_lib_required=False,
         bypass_encryption=False,
+        key_expiration_ms=None,
     ):
         """Options for :class:`MongoCrypt`.
 
@@ -53,6 +54,11 @@ def __init__(
           - `crypt_shared_lib_required`: Whether to require a crypt_shared
             library.
           - `bypass_encryption`: Whether to bypass encryption.
+          - `key_expiration_ms` (int): The cache expiration time for data
+            encryption keys. Defaults to 60000. 0 means keys never expire.
+
+        .. versionadded:: 1.13
+           Added the ``key_expiration_ms`` parameter.
 
         .. versionremoved:: 1.11
            Removed the ``enable_range_v2`` parameter.
@@ -136,6 +142,11 @@ def __init__(
             encrypted_fields_map, bytes
         ):
             raise TypeError("encrypted_fields_map must be bytes or None")
+        if key_expiration_ms is not None:
+            if not isinstance(key_expiration_ms, int):
+                raise TypeError("key_expiration_ms must be int or None")
+            if key_expiration_ms < 0:
+                raise ValueError("key_expiration_ms must be >=0 or None")
 
         self.kms_providers = kms_providers
         self.schema_map = schema_map
@@ -144,6 +155,7 @@ def __init__(
         self.crypt_shared_lib_path = crypt_shared_lib_path
         self.crypt_shared_lib_required = crypt_shared_lib_required
         self.bypass_encryption = bypass_encryption
+        self.key_expiration_ms = key_expiration_ms
 
 
 class ExplicitEncryptOpts:

bindings/python/test/test_mongocrypt.py

@@ -143,6 +143,11 @@ def test_mongocrypt_options(self):
         )
         self.assertEqual(opts.encrypted_fields_map, encrypted_fields_map)
         self.assertTrue(opts.bypass_query_analysis)
+        for expiration in [0, 1, 1000000]:
+            opts = MongoCryptOptions(
+                valid[0][0], schema_map, key_expiration_ms=expiration
+            )
+            self.assertEqual(opts.key_expiration_ms, expiration)
 
     def test_mongocrypt_options_validation(self):
         with self.assertRaisesRegex(
@@ -192,6 +197,12 @@ def test_mongocrypt_options_validation(self):
             TypeError, "encrypted_fields_map must be bytes or None"
         ):
             MongoCryptOptions(valid_kms, encrypted_fields_map={})
+        with self.assertRaisesRegex(TypeError, "key_expiration_ms must be int or None"):
+            MongoCryptOptions(valid_kms, key_expiration_ms="123")
+        with self.assertRaisesRegex(
+            ValueError, "key_expiration_ms must be >=0 or None"
+        ):
+            MongoCryptOptions(valid_kms, key_expiration_ms=-1)
 
 
 class TestMongoCrypt(unittest.TestCase):