diff --git a/connectors/sources/outlook.py b/connectors/sources/outlook.py index eeb5867cb..ffb2c11b2 100644 --- a/connectors/sources/outlook.py +++ b/connectors/sources/outlook.py @@ -319,8 +319,7 @@ def __init__( self.ssl_enabled = ssl_enabled self.ssl_ca = ssl_ca - @cached_property - def _create_connection(self): + def _create_ldap_connection(self): return Connection( server=self.ad_server, user=self.user, @@ -329,47 +328,51 @@ def _create_connection(self): auto_bind=True, # pyright: ignore ) - async def close(self): - pass - - def _fetch_normal_users(self, search_query): + @retryable( + retries=RETRIES, + interval=RETRY_INTERVAL, + strategy=RetryStrategy.EXPONENTIAL_BACKOFF, + skipped_exceptions=[UsersFetchFailed], + ) + def _ldap_search(self, search_query, search_filter): + connection = self._create_ldap_connection() try: - has_value_for_normal_users, _, response, _ = self._create_connection.search( + has_value, _, response, _ = connection.search( search_query, - SEARCH_FILTER_FOR_NORMAL_USERS, + search_filter, attributes=["mail"], ) - if not has_value_for_normal_users: + if not has_value: msg = "Error while fetching users from Exchange Active Directory." raise UsersFetchFailed(msg) - for user in response: - yield user + return response + finally: + try: + connection.unbind() + except Exception as exc: + logger.debug("Failed to unbind LDAP connection: %s", exc) + async def close(self): + pass + + def _fetch_normal_users(self, search_query): + try: + for user in self._ldap_search(search_query, SEARCH_FILTER_FOR_NORMAL_USERS): + yield user + except UsersFetchFailed: + raise except Exception as e: msg = f"Something went wrong while fetching users. Error: {e}" raise UsersFetchFailed(msg) from e def _fetch_admin_users(self, search_query): try: - ( - has_value_for_admin_users, - _, - response_for_admin, - _, - ) = self._create_connection.search( - search_query, - SEARCH_FILTER_FOR_ADMIN, - attributes=["mail"], - ) - - if not has_value_for_admin_users: - msg = "Error while fetching users from Exchange Active Directory." - raise UsersFetchFailed(msg) - - for user in response_for_admin: + for user in self._ldap_search(search_query, SEARCH_FILTER_FOR_ADMIN): yield user + except UsersFetchFailed: + raise except Exception as e: msg = f"Something went wrong while fetching users. Error: {e}" raise UsersFetchFailed(msg) from e diff --git a/tests/sources/test_outlook.py b/tests/sources/test_outlook.py index 687a1f189..03d33a53f 100644 --- a/tests/sources/test_outlook.py +++ b/tests/sources/test_outlook.py @@ -728,6 +728,66 @@ async def test_fetch_admin_users(mock_connection): assert users == ["test.user@gmail.com", "dummy.user@gmail.com"] +@patch("connectors.utils.time_to_sleep_between_retries", return_value=0) +@patch("connectors.sources.outlook.Connection") +def test_ldap_search_retries_on_transient_error(mock_connection, _mock_sleep): + mock_connection_instance = mock_connection.return_value + mock_connection_instance.search.side_effect = [ + OSError("Connection reset by peer"), + (True, None, ["user@example.com"], None), + ] + + exchange_users = ExchangeUsers( + ad_server="ad.example.com", + domain="example.com", + exchange_server="exchange.example.com", + user="user", + password="password", + ssl_enabled=False, + ssl_ca=None, + ) + response = exchange_users._ldap_search("search_query", "filter") + assert list(response) == ["user@example.com"] + assert mock_connection.call_count == 2 + assert mock_connection_instance.search.call_count == 2 + assert mock_connection_instance.unbind.call_count == 2 + + +@patch("connectors.sources.outlook.Connection") +def test_exchange_ldap_search_uses_fresh_connection_per_search(mock_connection): + normal_user = {"mail": "normal@example.com"} + admin_user = {"mail": "admin@example.com"} + + first_connection = MagicMock() + first_connection.search.return_value = (True, None, [normal_user], None) + second_connection = MagicMock() + second_connection.search.return_value = (True, None, [admin_user], None) + mock_connection.side_effect = [first_connection, second_connection] + + exchange_users = ExchangeUsers( + ad_server="ad.example.com", + domain="example.com", + exchange_server="exchange.example.com", + user="user", + password="password", + ssl_enabled=False, + ssl_ca=None, + ) + + users = asyncio.run(_collect_exchange_users(exchange_users)) + assert users == [normal_user, admin_user] + assert mock_connection.call_count == 2 + first_connection.search.assert_called_once() + second_connection.search.assert_called_once() + + +async def _collect_exchange_users(exchange_users): + users = [] + async for user in exchange_users.get_users(): + users.append(user) + return users + + @pytest.mark.asyncio @pytest.mark.parametrize( "attachment, expected_content",