Support pagination of get_users and add retries if rate limited

securitybot/chat/slack.py

@@ -7,12 +7,16 @@
 import logging
 from slackclient import SlackClient
 import json
+import time
 
 from securitybot.user import User
 from securitybot.chat.chat import Chat, ChatException
 
 from typing import Any, Dict, List
 
+RATE_LIMIT_SLEEP = 10 # sleep 10 seconds when rate limit
+RATE_LIMIT_TRIES = 6  # maximum 6 tries (60s) when rate limit reached
+
 class Slack(Chat):
     '''
     A wrapper around the Slack API designed for Securitybot.
@@ -37,7 +41,7 @@ def _validate(self):
             raise ChatException('Unable to connect to Slack API.')
         logging.info('Connection to Slack API successful!')
 
-    def _api_call(self, method, **kwargs):
+    def _api_call(self, method, rate_limit_retry=False, **kwargs):
         # type: (str, **Any) -> Dict[str, Any]
         '''
         Performs a _validated_ Slack API call. After performing a normal API
@@ -50,7 +54,19 @@ def _api_call(self, method, **kwargs):
         Returns:
             (dict): Parsed JSON from the response.
         '''
-        response = self._slack.api_call(method, **kwargs)
+        cur_try = 0
+        while True:
+            response = self._slack.api_call(method, **kwargs)
+            if cur_try > RATE_LIMIT_TRIES:
+                raise ChatException('Slack rate limit max tries reached.')
+
+            if response.get('error', '').lower() == 'ratelimited' and rate_limit_retry:
+                logging.debug('Rate limiting reached.  Sleeping {}.'.format(RATE_LIMIT_SLEEP))
+                cur_try += 1
+                time.sleep(RATE_LIMIT_SLEEP)
+            else:
+                break
+
         if not ('ok' in response and response['ok']):
             if kwargs:
                 logging.error('Bad Slack API request on {} with {}'.format(method, kwargs))
@@ -84,7 +100,20 @@ def get_users(self):
                     }
             }
         '''
-        return self._api_call('users.list')['members']
+        members = []
+        next_cursor = None
+        while True:
+            response = self._api_call('users.list', rate_limit_retry=True, cursor=next_cursor)
+            active_members = [m for m in response['members'] if not m.get('deleted')]
+            members.extend(active_members)
+            logging.debug('Fetched {} members'.format(len(members)))
+
+            metadata = response.get('response_metadata')
+            next_cursor = metadata.get('next_cursor') if metadata else None
+            if not metadata or not metadata.get('next_cursor'):
+                break
+
+        return members
 
     def get_messages(self):
         # type () -> List[Dict[str, Any]]