retry with wait when throttling error happen in sharepoint (#588)

Author: bryan-unstructured

CHANGELOG.md

@@ -1,9 +1,7 @@
-## 1.2.12-dev1
+## 1.2.12
 
+* **Fix: retry with wait when throttling error happens in Sharepoint connector**
 * **Fix: fix Milvus stager to use correct exception**
-
-## 1.2.12-dev0
-
 * **Fix: confluence integration test to use new link and credential**
 
 ## 1.2.11

pyproject.toml

@@ -143,6 +143,7 @@ test = [
     "pyarrow",
     "networkx",
     "htmlbuilder",
+    "office365-rest-python-client",
 ]
 # Add constraints needed for CI
 ci = [

test/unit/connectors/test_sharepoint.py

@@ -0,0 +1,135 @@
+from unittest.mock import Mock
+
+import pytest
+
+from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.error import SourceConnectionError
+from unstructured_ingest.processes.connectors.sharepoint import (
+    SharepointConnectionConfig,
+    SharepointDownloader,
+    SharepointDownloaderConfig,
+)
+
+
+@pytest.fixture
+def mock_client():
+    return Mock()
+
+
+@pytest.fixture
+def mock_site():
+    return Mock()
+
+
+@pytest.fixture
+def mock_drive_item():
+    return Mock()
+
+
+@pytest.fixture
+def mock_file():
+    return Mock()
+
+
+@pytest.fixture
+def mock_connection_config(mock_client, mock_drive_item):
+    config = Mock(spec=SharepointConnectionConfig)
+    config.site = "https://test.sharepoint.com/sites/test"
+    config.get_client.return_value = mock_client
+    config._get_drive_item.return_value = mock_drive_item
+    return config
+
+
+@pytest.fixture
+def mock_download_config():
+    config = Mock(spec=SharepointDownloaderConfig)
+    config.max_retries = 3
+    return config
+
+
+@pytest.fixture
+def sharepoint_downloader(mock_connection_config, mock_download_config):
+    downloader = SharepointDownloader(
+        connection_config=mock_connection_config, download_config=mock_download_config
+    )
+    return downloader
+
+
+@pytest.fixture
+def file_data():
+    return FileData(
+        source_identifiers=SourceIdentifiers(
+            filename="test.docx", fullpath="/sites/test/Shared Documents/test.docx"
+        ),
+        connector_type="sharepoint",
+        identifier="test-id",
+    )
+
+
+def test_fetch_file(
+    mock_client, mock_drive_item, mock_site, mock_file, sharepoint_downloader, file_data
+):
+    """Test successful file fetch without any errors"""
+    mock_client.sites.get_by_url.return_value.get.return_value.execute_query.return_value = (
+        mock_site
+    )
+    mock_drive_item.get_by_path.return_value.get.return_value.execute_query.return_value = mock_file
+    result = sharepoint_downloader._fetch_file(file_data)
+
+    assert result == mock_file
+    assert mock_client.sites.get_by_url.return_value.get.return_value.execute_query.call_count == 1
+    assert mock_drive_item.get_by_path.return_value.get.return_value.execute_query.call_count == 1
+    mock_drive_item.get_by_path.assert_called_with("/sites/test/Shared Documents/test.docx")
+
+
+def test_fetch_file_retries_on_429_error(
+    mock_client, mock_drive_item, mock_site, sharepoint_downloader, file_data
+):
+    """Test that _fetch_file retries when encountering 429 errors"""
+    mock_client.sites.get_by_url.return_value.get.return_value.execute_query.return_value = (
+        mock_site
+    )
+    mock_drive_item.get_by_path.return_value.get.return_value.execute_query.side_effect = [
+        Exception("429 Client Error"),
+        Exception("Request has been throttled"),
+        mock_file,
+    ]
+
+    result = sharepoint_downloader._fetch_file(file_data)
+    assert result == mock_file
+    assert mock_drive_item.get_by_path.return_value.get.return_value.execute_query.call_count == 3
+
+
+def test_fetch_file_fails_after_max_retries(
+    mock_client, mock_drive_item, mock_site, sharepoint_downloader, file_data
+):
+    """Test that _fetch_file fails after exhausting max retries"""
+    mock_client.sites.get_by_url.return_value.get.return_value.execute_query.return_value = (
+        mock_site
+    )
+    mock_drive_item.get_by_path.return_value.get.return_value.execute_query.side_effect = Exception(
+        "429 Client Error"
+    )
+
+    with pytest.raises(Exception, match="429"):
+        sharepoint_downloader._fetch_file(file_data)
+
+    expected_calls = sharepoint_downloader.download_config.max_retries
+    assert (
+        mock_drive_item.get_by_path.return_value.get.return_value.execute_query.call_count
+        == expected_calls
+    )
+
+
+def test_fetch_file_handles_site_not_found_immediately(
+    mock_client, sharepoint_downloader, file_data
+):
+    """Test that site not found errors are not retried"""
+    mock_client.sites.get_by_url.return_value.get.return_value.execute_query.side_effect = (
+        Exception("Site not found")
+    )
+
+    with pytest.raises(SourceConnectionError, match="Site not found"):
+        sharepoint_downloader._fetch_file(file_data)
+
+    assert mock_client.sites.get_by_url.return_value.get.return_value.execute_query.call_count == 1

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.2.12-dev1"  # pragma: no cover
+__version__ = "1.2.12"  # pragma: no cover

unstructured_ingest/processes/connectors/sharepoint.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, AsyncIterator, Optional
 
@@ -210,7 +211,7 @@ async def run_async(self, **kwargs: Any) -> AsyncIterator[FileData]:
 
 
 class SharepointDownloaderConfig(OnedriveDownloaderConfig):
-    pass
+    max_retries: int = 10
 
 
 @dataclass
@@ -219,10 +220,22 @@ class SharepointDownloader(OnedriveDownloader):
     download_config: SharepointDownloaderConfig
     connector_type: str = CONNECTOR_TYPE
 
+    @staticmethod
+    def retry_on_status_code(exc):
+        error_msg = str(exc).lower()
+        return "429" in error_msg or "activitylimitreached" in error_msg or "throttled" in error_msg
+
     @SourceConnectionNetworkError.wrap
     @requires_dependencies(["office365"], extras="sharepoint")
     def _fetch_file(self, file_data: FileData) -> DriveItem:
         from office365.runtime.client_request_exception import ClientRequestException
+        from tenacity import (
+            before_log,
+            retry,
+            retry_if_exception,
+            stop_after_attempt,
+            wait_exponential,
+        )
 
         if file_data.source_identifiers is None or not file_data.source_identifiers.fullpath:
             raise ValueError(
@@ -233,13 +246,27 @@ def _fetch_file(self, file_data: FileData) -> DriveItem:
         server_relative_path = file_data.source_identifiers.fullpath
         client = self.connection_config.get_client()
 
-        try:
-            client_site = client.sites.get_by_url(self.connection_config.site).get().execute_query()
-            site_drive_item = self.connection_config._get_drive_item(client_site)
-        except ClientRequestException:
-            logger.info("Site not found")
-            raise SourceConnectionError(f"Site not found: {self.connection_config.site}")
-        file = site_drive_item.get_by_path(server_relative_path).get().execute_query()
+        @retry(
+            stop=stop_after_attempt(self.download_config.max_retries),
+            wait=wait_exponential(exp_base=2, multiplier=1, min=2, max=10),
+            retry=retry_if_exception(self.retry_on_status_code),
+            before=before_log(logger, logging.DEBUG),
+            reraise=True,
+        )
+        def _get_item_by_path() -> DriveItem:
+            try:
+                client_site = (
+                    client.sites.get_by_url(self.connection_config.site).get().execute_query()
+                )
+                site_drive_item = self.connection_config._get_drive_item(client_site)
+            except ClientRequestException:
+                logger.info(f"Site not found: {self.connection_config.site}")
+                raise SourceConnectionError(f"Site not found: {self.connection_config.site}")
+            file = site_drive_item.get_by_path(server_relative_path).get().execute_query()
+            return file
+
+        # Call the retry-wrapped function
+        file = _get_item_by_path()
 
         if not file:
             raise NotFoundError(f"file not found: {server_relative_path}")