feat: Refresh app tokens

tap_github/authenticator.py

@@ -113,25 +113,15 @@ def generate_app_access_token(
 class GitHubTokenAuthenticator(APIAuthenticatorBase):
     """Base class for offloading API auth."""
 
-    def prepare_tokens(self) -> Dict[str, TokenRateLimit]:
-        # Save GitHub tokens
-        available_tokens: List[str] = []
-        if "auth_token" in self._config:
-            available_tokens = available_tokens + [self._config["auth_token"]]
-        if "additional_auth_tokens" in self._config:
-            available_tokens = available_tokens + self._config["additional_auth_tokens"]
-        else:
-            # Accept multiple tokens using environment variables GITHUB_TOKEN*
-            env_tokens = [
-                value
-                for key, value in environ.items()
-                if key.startswith("GITHUB_TOKEN")
-            ]
-            if len(env_tokens) > 0:
-                self.logger.info(
-                    f"Found {len(env_tokens)} 'GITHUB_TOKEN' environment variables for authentication."
-                )
-                available_tokens = env_tokens
+    last_private_key_token_refresh: Optional[datetime]
+
+    def refresh_app_token(self):
+        if self.last_private_key_token_refresh:
+            # Do not refresh token if less than 10 minutes have passed, if this is requested the app probably hit its rate limit.
+            if (
+                datetime.now() - self.last_private_key_token_refresh
+            ).total_seconds() < 600:
+                return
 
         # Parse App level private key and generate a token
         if "GITHUB_APP_PRIVATE_KEY" in environ.keys():
@@ -152,7 +142,32 @@ def prepare_tokens(self) -> Dict[str, TokenRateLimit]:
                 app_token = generate_app_access_token(
                     github_app_id, github_private_key, github_installation_id or None
                 )
-                available_tokens = available_tokens + [app_token]
+                # Get rate_limit_buffer
+                rate_limit_buffer = self._config.get("rate_limit_buffer", None)
+                # Update last refresh timestamp
+                self.last_private_key_token_refresh = datetime.now()
+
+                return TokenRateLimit(app_token, rate_limit_buffer)
+
+    def prepare_tokens(self) -> Dict[str, TokenRateLimit]:
+        # Save GitHub tokens
+        available_tokens: List[str] = []
+        if "auth_token" in self._config:
+            available_tokens = available_tokens + [self._config["auth_token"]]
+        if "additional_auth_tokens" in self._config:
+            available_tokens = available_tokens + self._config["additional_auth_tokens"]
+        else:
+            # Accept multiple tokens using environment variables GITHUB_TOKEN*
+            env_tokens = [
+                value
+                for key, value in environ.items()
+                if key.startswith("GITHUB_TOKEN")
+            ]
+            if len(env_tokens) > 0:
+                self.logger.info(
+                    f"Found {len(env_tokens)} 'GITHUB_TOKEN' environment variables for authentication."
+                )
+                available_tokens = env_tokens
 
         # Get rate_limit_buffer
         rate_limit_buffer = self._config.get("rate_limit_buffer", None)
@@ -180,8 +195,6 @@ def prepare_tokens(self) -> Dict[str, TokenRateLimit]:
         self.logger.info(f"Tap will run with {len(filtered_tokens)} auth tokens")
 
         # Create a dict of TokenRateLimit
-        # TODO - separate app_token and add logic to refresh the token
-        # using generate_app_access_token.
         return {
             token: TokenRateLimit(token, rate_limit_buffer) for token in filtered_tokens
         }
@@ -200,6 +213,10 @@ def __init__(self, stream: RESTStream) -> None:
         self.active_token: Optional[TokenRateLimit] = (
             choice(list(self.tokens_map.values())) if len(self.tokens_map) else None
         )
+        self.last_private_key_token_refresh = None
+        # Refresh tokens from private key if it was supplied
+        if self.active_token is None:
+            self.active_token = self.refresh_app_token()
 
     def get_next_auth_token(self) -> None:
         tokens_list = list(self.tokens_map.items())
@@ -234,6 +251,16 @@ def auth_headers(self) -> Dict[str, str]:
             HTTP headers for authentication.
         """
         result = super().auth_headers
+
+        if self.last_private_key_token_refresh is not None:
+            # Refresh token once every 30 minutes if we have a private key
+            if (
+                datetime.now() - self.last_private_key_token_refresh
+            ).total_seconds() > 1800:
+                new_token = self.refresh_app_token()
+                if new_token:
+                    self.active_token = new_token
+
         if self.active_token:
             # Make sure that our token is still valid or update it.
             if not self.active_token.is_valid():