Fix slack pagination

This commit is contained in:
Weves 2023-08-15 17:17:55 -07:00 committed by Chris Weaver
parent 78b49f546c
commit 8fc74a4313
2 changed files with 16 additions and 12 deletions

View File

@ -36,7 +36,7 @@ ThreadType = list[MessageType]
def _make_paginated_slack_api_call( def _make_paginated_slack_api_call(
call: Callable[..., SlackResponse], **kwargs: Any call: Callable[..., SlackResponse], **kwargs: Any
) -> list[dict[str, Any]]: ) -> Generator[dict[str, Any], None, None]:
return make_slack_api_call_paginated( return make_slack_api_call_paginated(
make_slack_api_rate_limited(make_slack_api_call_logged(call)) make_slack_api_rate_limited(make_slack_api_call_logged(call))
)(**kwargs) )(**kwargs)
@ -50,9 +50,9 @@ def _make_slack_api_call(
def get_channel_info(client: WebClient, channel_id: str) -> ChannelType: def get_channel_info(client: WebClient, channel_id: str) -> ChannelType:
"""Get information about a channel. Needed to convert channel ID to channel name""" """Get information about a channel. Needed to convert channel ID to channel name"""
return _make_paginated_slack_api_call( return _make_slack_api_call(client.conversations_info, channel=channel_id)[0][
client.conversations_info, channel=channel_id "channel"
)[0]["channel"] ]
def get_channels( def get_channels(

View File

@ -1,6 +1,7 @@
import re import re
import time import time
from collections.abc import Callable from collections.abc import Callable
from collections.abc import Generator
from functools import wraps from functools import wraps
from typing import Any from typing import Any
from typing import cast from typing import cast
@ -45,20 +46,20 @@ def make_slack_api_call_logged(
def make_slack_api_call_paginated( def make_slack_api_call_paginated(
call: Callable[..., SlackResponse], call: Callable[..., SlackResponse],
) -> Callable[..., list[dict[str, Any]]]: ) -> Callable[..., Generator[dict[str, Any], None, None]]:
"""Wraps calls to slack API so that they automatically handle pagination""" """Wraps calls to slack API so that they automatically handle pagination"""
@wraps(call) @wraps(call)
def paginated_call(**kwargs: Any) -> list[dict[str, Any]]: def paginated_call(**kwargs: Any) -> Generator[dict[str, Any], None, None]:
results: list[dict[str, Any]] = []
cursor: str | None = None cursor: str | None = None
has_more = True has_more = True
while has_more: while has_more:
for result in call(cursor=cursor, limit=_SLACK_LIMIT, **kwargs): response = call(cursor=cursor, limit=_SLACK_LIMIT, **kwargs)
has_more = result.get("has_more", False) yield cast(dict[str, Any], response.validate())
cursor = result.get("response_metadata", {}).get("next_cursor", "") cursor = cast(dict[str, Any], response.get("response_metadata", {})).get(
results.append(cast(dict[str, Any], result)) "next_cursor", ""
return results )
has_more = bool(cursor)
return paginated_call return paginated_call
@ -84,6 +85,9 @@ def make_slack_api_rate_limited(
if e.response["error"] == "ratelimited": if e.response["error"] == "ratelimited":
# Handle rate limiting: get the 'Retry-After' header value and sleep for that duration # Handle rate limiting: get the 'Retry-After' header value and sleep for that duration
retry_after = int(e.response.headers.get("Retry-After", 1)) retry_after = int(e.response.headers.get("Retry-After", 1))
logger.info(
f"Slack call rate limited, retrying after {retry_after} seconds. Exception: {e}"
)
time.sleep(retry_after) time.sleep(retry_after)
else: else:
# Raise the error for non-transient errors # Raise the error for non-transient errors