Add support for commit.retry.total-timeout-ms property

Author: MoIMC

pyiceberg/table/__init__.py

@@ -42,7 +42,7 @@
 
 from pydantic import Field
 from sortedcontainers import SortedList
-from tenacity import RetryCallState, retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential
+from tenacity import RetryCallState, retry, retry_if_exception_type, stop_after_attempt, stop_after_delay, wait_random_exponential
 
 import pyiceberg.expressions.parser as parser
 from pyiceberg.exceptions import CommitFailedException
@@ -91,6 +91,8 @@
     COMMIT_MIN_RETRY_WAIT_MS_DEFAULT,
     COMMIT_NUM_RETRIES,
     COMMIT_NUM_RETRIES_DEFAULT,
+    COMMIT_RETRY_TOTAL_TIMEOUT_MS,
+    COMMIT_RETRY_TOTAL_TIMEOUT_MS_DEFAULT,
     INITIAL_SEQUENCE_NUMBER,
     TableMetadata,
 )
@@ -785,6 +787,7 @@ def commit_transaction(self) -> Table:
         min_wait_ms = int(self.table_metadata.properties.get(COMMIT_MIN_RETRY_WAIT_MS, COMMIT_MIN_RETRY_WAIT_MS_DEFAULT))
         max_wait_ms = int(self.table_metadata.properties.get(COMMIT_MAX_RETRY_WAIT_MS, COMMIT_MAX_RETRY_WAIT_MS_DEFAULT))
         num_retries = int(self.table_metadata.properties.get(COMMIT_NUM_RETRIES, COMMIT_NUM_RETRIES_DEFAULT))
+        timeout_ms = int(self.table_metadata.properties.get(COMMIT_RETRY_TOTAL_TIMEOUT_MS, COMMIT_RETRY_TOTAL_TIMEOUT_MS_DEFAULT))
 
         def _before_attempt(state: RetryCallState):
             if state.attempt_number > 1:
@@ -814,7 +817,7 @@ def _error_callback(state: RetryCallState):
         @wraps(self.commit_transaction)
         @retry(
             wait=wait_random_exponential(min=min_wait_ms / 1000, max=max_wait_ms / 1000),
-            stop=stop_after_attempt(num_retries + 1),
+            stop=(stop_after_attempt(num_retries + 1) | stop_after_delay(timeout_ms / 1000)),
             retry=retry_if_exception_type(CommitFailedException),
             before=_before_attempt,
             retry_error_callback=_error_callback,

pyiceberg/table/metadata.py

@@ -70,6 +70,8 @@
 COMMIT_MAX_RETRY_WAIT_MS = "commit.retry.max-wait-ms"
 COMMIT_MAX_RETRY_WAIT_MS_DEFAULT = 5000  # 5 seconds
 
+COMMIT_RETRY_TOTAL_TIMEOUT_MS = "commit.retry.total-timeout-ms"
+COMMIT_RETRY_TOTAL_TIMEOUT_MS_DEFAULT = 1_800_000  # 30 mins
 
 INITIAL_SEQUENCE_NUMBER = 0
 INITIAL_SPEC_ID = 0

pyiceberg/table/update/__init__.py

@@ -25,7 +25,15 @@
 from typing import TYPE_CHECKING, Annotated, Any, Dict, Generic, List, Literal, Optional, Tuple, TypeVar, Union, cast
 
 from pydantic import Field, field_validator, model_validator
-from tenacity import RetryCallState, before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential
+from tenacity import (
+    RetryCallState,
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_random_exponential,
+)
 
 from pyiceberg.exceptions import CommitFailedException
 from pyiceberg.partitioning import PARTITION_FIELD_ID_START, PartitionSpec
@@ -37,6 +45,8 @@
     COMMIT_MIN_RETRY_WAIT_MS_DEFAULT,
     COMMIT_NUM_RETRIES,
     COMMIT_NUM_RETRIES_DEFAULT,
+    COMMIT_RETRY_TOTAL_TIMEOUT_MS,
+    COMMIT_RETRY_TOTAL_TIMEOUT_MS_DEFAULT,
     SUPPORTED_TABLE_FORMAT_VERSION,
     TableMetadata,
     TableMetadataUtil,
@@ -86,6 +96,9 @@ def commit(self) -> None:
             self._transaction.table_metadata.properties.get(COMMIT_MAX_RETRY_WAIT_MS, COMMIT_MAX_RETRY_WAIT_MS_DEFAULT)
         )
         num_retries = int(self._transaction.table_metadata.properties.get(COMMIT_NUM_RETRIES, COMMIT_NUM_RETRIES_DEFAULT))
+        timeout_ms = int(
+            self._transaction.table_metadata.properties.get(COMMIT_RETRY_TOTAL_TIMEOUT_MS, COMMIT_RETRY_TOTAL_TIMEOUT_MS_DEFAULT)
+        )
 
         def _before_commit_inner(state: RetryCallState) -> None:
             if state.attempt_number > 1:
@@ -94,7 +107,7 @@ def _before_commit_inner(state: RetryCallState) -> None:
         @wraps(self.commit)
         @retry(
             wait=wait_random_exponential(min=min_wait_ms / 1000, max=max_wait_ms / 1000),
-            stop=stop_after_attempt(num_retries),
+            stop=(stop_after_attempt(num_retries) | stop_after_delay(timeout_ms / 1000)),
             retry=retry_if_exception_type(CommitFailedException),
             before=_before_commit_inner,
             before_sleep=lambda state: logger.debug(