fix: ibm watsonx s3 expired token error (#470)

* Fix token expiration error in ibm watsonx connector

* version and changelog bump; ibm watsonx connector fix

* Rename unit test

* version and changelog bump; ibm watsonx connector fix

Author: mpolomdeepsense

CHANGELOG.md

@@ -1,8 +1,9 @@
-## 1.0.12-dev0
+## 1.0.12
 
 ### Fixes
 
 * **Replaced google drive connector's mechanism for file downloads.**
+* **Fix Token expiration error in IBM watsonx.data connector**
 
 ## 1.0.11
 

test/unit/connectors/ibm_watsonx/test_ibm_watsonx_s3.py

@@ -4,7 +4,7 @@
 import pandas as pd
 import pytest
 from pydantic import Secret
-from pyiceberg.exceptions import CommitFailedException
+from pyiceberg.exceptions import CommitFailedException, RESTError
 from pytest_mock import MockerFixture
 
 from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
@@ -202,7 +202,7 @@ def test_ibm_watsonx_connection_config_bearer_token_soon_to_expire_token(
 ):
     connection_config._bearer_token = {
         "access_token": "soon_to_expire_token",
-        "expiration": timestamp_now + 60,
+        "expiration": timestamp_now + (60 * 5),
     }
     mock_generate_bearer_token = mocker.patch.object(
         IbmWatsonxConnectionConfig,
@@ -332,6 +332,21 @@ def test_ibm_watsonx_uploader_upload_data_table_commit_exception(
     assert mock_table.refresh.call_count == 5
 
 
+def test_ibm_watsonx_uploader_upload_data_table_rest_error(
+    uploader: IbmWatsonxUploader,
+    mock_table: MagicMock,
+    mock_transaction: MagicMock,
+    mock_data_table: MagicMock,
+    mock_delete: MagicMock,
+    file_data: FileData,
+):
+    mock_transaction.append.side_effect = RESTError()
+
+    with pytest.raises(RESTError):
+        uploader.upload_data_table(mock_table, mock_data_table, file_data)
+    assert mock_table.refresh.call_count == 0
+
+
 def test_ibm_watsonx_uploader_upload_data_table_exception(
     uploader: IbmWatsonxUploader,
     mock_table: MagicMock,
@@ -428,6 +443,63 @@ def test_ibm_watsonx_uploader_upload_dataframe_success(
     mock_upload_data_table.assert_called_once_with(mock_table, mock_data_table, file_data)
 
 
+def test_ibm_watsonx_uploader_upload_dataframe_rest_error(
+    mocker: MockerFixture,
+    uploader: IbmWatsonxUploader,
+    test_df: pd.DataFrame,
+    mock_get_table: MagicMock,
+    mock_data_table: MagicMock,
+    file_data: FileData,
+):
+    mocker.patch.object(IbmWatsonxUploader, "_df_to_arrow_table", return_value=mock_data_table)
+    mock_upload_data_table = mocker.patch.object(
+        IbmWatsonxUploader, "upload_data_table", side_effect=RESTError()
+    )
+
+    with pytest.raises(ProviderError):
+        uploader.upload_dataframe(test_df, file_data)
+    assert mock_get_table.call_count == 2
+    assert mock_upload_data_table.call_count == 2
+
+
+def test_ibm_watsonx_uploader_upload_dataframe_provider_error(
+    mocker: MockerFixture,
+    uploader: IbmWatsonxUploader,
+    test_df: pd.DataFrame,
+    mock_get_table: MagicMock,
+    mock_data_table: MagicMock,
+    file_data: FileData,
+):
+    mocker.patch.object(IbmWatsonxUploader, "_df_to_arrow_table", return_value=mock_data_table)
+    mock_upload_data_table = mocker.patch.object(
+        IbmWatsonxUploader, "upload_data_table", side_effect=ProviderError()
+    )
+
+    with pytest.raises(ProviderError):
+        uploader.upload_dataframe(test_df, file_data)
+    mock_get_table.assert_called_once()
+    mock_upload_data_table.assert_called_once()
+
+
+def test_ibm_watsonx_uploader_upload_dataframe_exception(
+    mocker: MockerFixture,
+    uploader: IbmWatsonxUploader,
+    test_df: pd.DataFrame,
+    mock_get_table: MagicMock,
+    mock_data_table: MagicMock,
+    file_data: FileData,
+):
+    mocker.patch.object(IbmWatsonxUploader, "_df_to_arrow_table", return_value=mock_data_table)
+    mock_upload_data_table = mocker.patch.object(
+        IbmWatsonxUploader, "upload_data_table", side_effect=Exception()
+    )
+
+    with pytest.raises(ProviderError):
+        uploader.upload_dataframe(test_df, file_data)
+    mock_get_table.assert_called_once()
+    mock_upload_data_table.assert_called_once()
+
+
 def test_ibm_watsonx_uploader_delete_can_delete(
     mocker: MockerFixture,
     uploader: IbmWatsonxUploader,

unstructured_ingest/__version__.py

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

unstructured_ingest/processes/connectors/ibm_watsonx/ibm_watsonx_s3.py

@@ -69,8 +69,8 @@ def object_storage_url(self) -> str:
 
     @property
     def bearer_token(self) -> str:
-        # Add 60 seconds to deal with edge cases where the token expires before the request is made
-        timestamp = int(time.time()) + 60
+        # Add 5 minutes to deal with edge cases where the token expires before the request is made
+        timestamp = int(time.time()) + (60 * 5)
         if self._bearer_token is None or self._bearer_token.get("expiration", 0) <= timestamp:
             self._bearer_token = self.generate_bearer_token()
         return self._bearer_token["access_token"]
@@ -240,7 +240,7 @@ def _delete(self, transaction: "Transaction", identifier: str) -> None:
     def upload_data_table(
         self, table: "Table", data_table: "ArrowTable", file_data: FileData
     ) -> None:
-        from pyiceberg.exceptions import CommitFailedException
+        from pyiceberg.exceptions import CommitFailedException, RESTError
         from tenacity import (
             before_log,
             retry,
@@ -265,21 +265,51 @@ def _upload_data_table(table: "Table", data_table: "ArrowTable", file_data: File
                 table.refresh()
                 logger.debug(e)
                 raise IcebergCommitFailedException(e)
+            except RESTError:
+                raise
             except Exception as e:
                 raise ProviderError(f"Failed to upload data to table: {e}")
 
         try:
             return _upload_data_table(table, data_table, file_data)
+        except RESTError:
+            raise
         except ProviderError:
             raise
         except Exception as e:
             raise ProviderError(f"Failed to upload data to table: {e}")
 
+    @requires_dependencies(["pyiceberg", "tenacity"], extras="ibm-watsonx-s3")
     def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
+        from pyiceberg.exceptions import RESTError
+        from tenacity import (
+            before_log,
+            retry,
+            retry_if_exception_type,
+            stop_after_attempt,
+            wait_random,
+        )
+
         data_table = self._df_to_arrow_table(df)
 
-        with self.get_table() as table:
-            self.upload_data_table(table, data_table, file_data)
+        # Retry connection in case of connection error
+        @retry(
+            stop=stop_after_attempt(2),
+            wait=wait_random(),
+            retry=retry_if_exception_type(RESTError),
+            before=before_log(logger, logging.DEBUG),
+            reraise=True,
+        )
+        def _upload_dataframe(data_table: Any, file_data: FileData) -> None:
+            with self.get_table() as table:
+                self.upload_data_table(table, data_table, file_data)
+
+        try:
+            return _upload_dataframe(data_table, file_data)
+        except ProviderError:
+            raise
+        except Exception as e:
+            raise ProviderError(f"Failed to upload data to table: {e}")
 
     @requires_dependencies(["pandas"], extras="ibm-watsonx-s3")
     def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None: