PYTHON-4992 Support retrying KMS requests (#920)

Author: ShaneHarvey

bindings/python/pymongocrypt/binding.py

@@ -312,6 +312,17 @@ def _parse_version(version):
  */
 bool mongocrypt_setopt_log_handler(mongocrypt_t *crypt, mongocrypt_log_fn_t log_fn, void *log_ctx);
 
+/**
+ * Enable or disable KMS retry behavior.
+ *
+ * @param[in] crypt The @ref mongocrypt_t object.
+ * @param[in] enable A boolean indicating whether to retry operations.
+ * @pre @ref mongocrypt_init has not been called on @p crypt.
+ * @returns A boolean indicating success. If false, an error status is set.
+ * Retrieve it with @ref mongocrypt_ctx_status
+ */
+bool mongocrypt_setopt_retry_kms(mongocrypt_t *crypt, bool enable);
+
 /**
  * Configure an AWS KMS provider on the @ref mongocrypt_t object.
  *
@@ -1042,6 +1053,8 @@ def _parse_version(version):
  * If KMS handles are being handled synchronously, the driver can reuse the same
  * TLS socket to send HTTP requests and receive responses.
  *
+ * The returned KMS handle does not outlive `ctx`.
+ *
  * @param[in] ctx A @ref mongocrypt_ctx_t.
  * @returns a new @ref mongocrypt_kms_ctx_t or NULL.
  */
@@ -1086,6 +1099,14 @@ def _parse_version(version):
  */
 uint32_t mongocrypt_kms_ctx_bytes_needed(mongocrypt_kms_ctx_t *kms);
 
+/**
+ * Indicates how long to sleep before sending this request.
+ *
+ * @param[in] kms The @ref mongocrypt_kms_ctx_t.
+ * @returns How long to sleep in microseconds.
+ */
+int64_t mongocrypt_kms_ctx_usleep(mongocrypt_kms_ctx_t *kms);
+
 /**
  * Feed bytes from the HTTP response.
  *
@@ -1100,6 +1121,14 @@ def _parse_version(version):
  */
 bool mongocrypt_kms_ctx_feed(mongocrypt_kms_ctx_t *kms, mongocrypt_binary_t *bytes);
 
+/**
+ * Indicate a network-level failure.
+ *
+ * @param[in] kms The @ref mongocrypt_kms_ctx_t.
+ * @return A boolean indicating whether the failed request may be retried.
+ */
+bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms);
+
 /**
  * Get the status associated with a @ref mongocrypt_kms_ctx_t object.
  *
@@ -1413,7 +1442,7 @@ def _parse_version(version):
  * {
  *    "min": Optional<BSON value>,
  *    "max": Optional<BSON value>,
- *    "sparsity": Int64,
+ *    "sparsity": Optional<Int64>,
  *    "precision": Optional<Int32>,
  *    "trimFactor": Optional<Int32>
  * }
@@ -1426,6 +1455,15 @@ def _parse_version(version):
  */
 bool mongocrypt_ctx_setopt_algorithm_range(mongocrypt_ctx_t *ctx, mongocrypt_binary_t *opts);
 
+/**
+ * Set the expiration time for the data encryption key cache. Defaults to 60 seconds if not set.
+ *
+ * @param[in] ctx The @ref mongocrypt_ctx_t object.
+ * @param[in] cache_expiration_ms The cache expiration time in milliseconds. If zero, the cache
+ * never expires.
+ */
+bool mongocrypt_setopt_key_expiration(mongocrypt_t *crypt, uint64_t cache_expiration_ms);
+
 /// String constants for setopt_query_type
 // DEPRECATED: Support "rangePreview" has been removed in favor of "range".
 """

bindings/python/pymongocrypt/mongocrypt.py

@@ -149,6 +149,14 @@ 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,
+        try:
+            if not lib.mongocrypt_setopt_retry_kms(self.__crypt, True):
+                self.__raise_from_status()
+        except AttributeError:
+            # libmongocrypt < 1.12
+            pass
+
         if not lib.mongocrypt_init(self.__crypt):
             self.__raise_from_status()
 
@@ -670,6 +678,30 @@ def feed(self, data):
             if not lib.mongocrypt_kms_ctx_feed(self.__ctx, binary.bin):
                 self.__raise_from_status()
 
+    @property
+    def usleep(self):
+        """Indicates how long to sleep in microseconds before sending this request.
+
+        .. versionadded:: 1.12
+        """
+        try:
+            return lib.mongocrypt_kms_ctx_usleep(self.__ctx)
+        except AttributeError:
+            # libmongocrypt < 1.12
+            return 0
+
+    def fail(self):
+        """Indicate a network-level failure.
+
+        .. versionadded:: 1.12
+        """
+        try:
+            if not lib.mongocrypt_kms_ctx_fail(self.__ctx):
+                self.__raise_from_status()
+        except AttributeError:
+            # libmongocrypt < 1.12
+            pass
+
     def __raise_from_status(self):
         status = lib.mongocrypt_status_new()
         try: