fix: ibm watsonx s3 retry connection on error (#474)

* IBM watsonx s3 destination retry on connection error

* update version and changelog; ibm retry on connection error fix

Author: mpolomdeepsense

CHANGELOG.md

@@ -1,8 +1,9 @@
-## 1.0.12-dev0
+## 1.0.13
 
 ### Fixes
 
 * **Fix Notion connector database property missing 'description' attribute error**
+* **Retry IBM watsonx S3 upload on connection error**
 
 ## 1.0.12
 

test/unit/connectors/ibm_watsonx/test_ibm_watsonx_s3.py

@@ -46,6 +46,7 @@ def connection_config(access_config: IbmWatsonxAccessConfig):
         object_storage_endpoint="test_object_storage_endpoint/",
         object_storage_region="test_region",
         catalog="test_catalog",
+        max_retries_connection=2,
     )
 
 
@@ -248,7 +249,7 @@ def test_ibm_watsonx_connection_config_get_catalog_success(
 def test_ibm_watsonx_connection_config_get_catalog_failure(
     mocker: MockerFixture, connection_config: IbmWatsonxConnectionConfig
 ):
-    mocker.patch(
+    mock_load_catalog = mocker.patch(
         "pyiceberg.catalog.load_catalog",
         side_effect=Exception("Connection error"),
     )
@@ -258,7 +259,23 @@ def test_ibm_watsonx_connection_config_get_catalog_failure(
         new="test_bearer_token",
     )
     with pytest.raises(ProviderError), connection_config.get_catalog():
-        pass
+        mock_load_catalog.assert_called_once()
+
+
+def test_ibm_watsonx_connection_config_get_catalog_failure_retries(
+    mocker: MockerFixture, connection_config: IbmWatsonxConnectionConfig
+):
+    mock_load_catalog = mocker.patch(
+        "pyiceberg.catalog.load_catalog",
+        side_effect=RESTError("Connection error"),
+    )
+    mocker.patch.object(
+        IbmWatsonxConnectionConfig,
+        "bearer_token",
+        new="test_bearer_token",
+    )
+    with pytest.raises(ProviderError), connection_config.get_catalog():
+        assert mock_load_catalog.call_count == 2
 
 
 def test_ibm_watsonx_uploader_precheck_namespace_exists_table_exists(

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.0.12-dev0"  # pragma: no cover
+__version__ = "1.0.13"  # pragma: no cover

unstructured_ingest/processes/connectors/ibm_watsonx/ibm_watsonx_s3.py

@@ -56,6 +56,12 @@ class IbmWatsonxConnectionConfig(ConnectionConfig):
     object_storage_endpoint: str = Field(description="Cloud Object Storage public endpoint")
     object_storage_region: str = Field(description="Cloud Object Storage region")
     catalog: str = Field(description="Catalog name")
+    max_retries_connection: int = Field(
+        default=10,
+        description="Maximum number of retries in case of a connection error (RESTError)",
+        ge=2,
+        le=100,
+    )
 
     _bearer_token: Optional[dict[str, Any]] = None
 
@@ -145,10 +151,29 @@ def get_catalog_config(self) -> dict[str, Any]:
     @contextmanager
     def get_catalog(self) -> Generator["RestCatalog", None, None]:
         from pyiceberg.catalog import load_catalog
+        from pyiceberg.exceptions import RESTError
+        from tenacity import (
+            before_log,
+            retry,
+            retry_if_exception_type,
+            stop_after_attempt,
+            wait_exponential,
+        )
+
+        # Retry connection in case of a connection error
+        @retry(
+            stop=stop_after_attempt(self.max_retries_connection),
+            wait=wait_exponential(exp_base=2, multiplier=1, min=2, max=10),
+            retry=retry_if_exception_type(RESTError),
+            before=before_log(logger, logging.DEBUG),
+            reraise=True,
+        )
+        def _get_catalog(catalog_config: dict[str, Any]) -> "RestCatalog":
+            return load_catalog(**catalog_config)
 
         try:
             catalog_config = self.get_catalog_config()
-            catalog = load_catalog(**catalog_config)
+            catalog = _get_catalog(catalog_config)
         except Exception as e:
             logger.error(f"Failed to connect to catalog '{self.catalog}': {e}", exc_info=True)
             raise ProviderError(f"Failed to connect to catalog '{self.catalog}': {e}")
@@ -172,7 +197,10 @@ class IbmWatsonxUploaderConfig(UploaderConfig):
     namespace: str = Field(description="Namespace name")
     table: str = Field(description="Table name")
     max_retries: int = Field(
-        default=5, description="Maximum number of retries to upload data", ge=2, le=500
+        default=5,
+        description="Maximum number of retries to upload data (CommitFailedException)",
+        ge=2,
+        le=500,
     )
     record_id_key: str = Field(
         default=RECORD_ID_LABEL,
@@ -287,15 +315,15 @@ def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
             retry,
             retry_if_exception_type,
             stop_after_attempt,
-            wait_random,
+            wait_exponential,
         )
 
         data_table = self._df_to_arrow_table(df)
 
-        # Retry connection in case of connection error
+        # Retry connection in case of a connection error or token expiration
         @retry(
-            stop=stop_after_attempt(2),
-            wait=wait_random(),
+            stop=stop_after_attempt(self.connection_config.max_retries_connection),
+            wait=wait_exponential(exp_base=2, multiplier=1, min=2, max=10),
             retry=retry_if_exception_type(RESTError),
             before=before_log(logger, logging.DEBUG),
             reraise=True,