Add command to get github api key statuses

augur/application/cli/github.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: MIT
+from os import walk, chdir, environ, chmod, path
+import os
+import logging
+from sys import exit
+import stat
+from collections import OrderedDict
+from subprocess import call
+import random
+import string
+import csv
+import click
+import sqlalchemy as s
+import pandas as pd
+import requests
+import json
+import sqlalchemy as s
+import re
+
+from augur.application.cli import test_connection, test_db_connection 
+
+from augur.application.db.session import DatabaseSession
+from augur.application.logs import AugurLogger
+from augur.application.db.engine import DatabaseEngine
+from sqlalchemy import update
+from datetime import datetime
+from augur.application.db.models import Repo
+
+
+import httpx
+from collections import Counter
+
+
+from augur.tasks.github.util.github_api_key_handler import GithubApiKeyHandler
+
+
+logger = logging.getLogger(__name__)
+
+@click.group("github", short_help="Github utilities")
+def cli():
+    pass
+
+@cli.command("api-keys")
+@test_connection
+@test_db_connection
+def update_api_key():
+    """
+    Get the ratelimit of Github API keys
+    """
+
+    with DatabaseEngine() as engine, engine.connect() as connection:
+
+        get_api_keys_sql = s.sql.text(
+            """
+            SELECT value as github_key from config Where section_name='Keys' AND setting_name='github_api_key'
+            UNION All
+            SELECT access_token as github_key from worker_oauth ORDER BY github_key DESC;
+            """
+        )
+
+        result = connection.execute(get_api_keys_sql).fetchall()
+        keys = [x[0] for x in result]
+
+        with httpx.Client() as client:
+
+            invalid_keys = []
+            valid_key_data = []
+            for key in keys:
+                core_key_data, graphql_key_data = GithubApiKeyHandler.get_key_rate_limit(client, key)
+                if core_key_data is None or graphql_key_data is None:
+                    invalid_keys.append(key)
+                else:
+                    valid_key_data.append((key, core_key_data, graphql_key_data))
+
+            valid_key_data = sorted(valid_key_data, key=lambda x: x[1]["requests_remaining"])
+
+            core_request_header = "Core Requests Left"
+            core_reset_header = "Core Reset Time"
+            graphql_request_header = "Graphql Requests Left"
+            graphql_reset_header = "Graphql Reset Time"
+            print(f"{'Key'.center(40)}   {core_request_header}   {core_reset_header}   {graphql_request_header}   {graphql_reset_header}")
+            for key, core_key_data, graphql_key_data in valid_key_data:
+                core_requests = str(core_key_data['requests_remaining']).center(len(core_request_header))
+                core_reset_time = str(epoch_to_local_time_with_am_pm(core_key_data["reset_epoch"])).center(len(core_reset_header))
+
+                graphql_requests = str(graphql_key_data['requests_remaining']).center(len(graphql_request_header))
+                graphql_reset_time = str(epoch_to_local_time_with_am_pm(graphql_key_data["reset_epoch"])).center(len(graphql_reset_header))
+
+                print(f"{key} | {core_requests} | {core_reset_time} | {graphql_requests} | {graphql_reset_time} |")
+
+            valid_key_list = [x[0] for x in valid_key_data]
+            duplicate_keys = find_duplicates(valid_key_list)
+            if len(duplicate_keys) > 0:
+                print("\n\nWARNING: There are duplicate keys this will slow down collection")
+                print("Duplicate keys".center(40))
+                for key in duplicate_keys:
+                    print(key)
+
+
+            if len(invalid_keys) > 0:
+                invalid_key_header = "Invalid Keys".center(40)
+                print("\n")
+                print(invalid_key_header)
+                for key in invalid_keys:
+                    print(key)
+                print("")
+
+
+
+    engine.dispose()
+
+
+def epoch_to_local_time_with_am_pm(epoch):
+    local_time = datetime.fromtimestamp(epoch)
+    formatted_time = local_time.strftime('%I:%M %p')  # This format includes the date as well
+    return formatted_time
+
+
+def find_duplicates(lst):
+    counter = Counter(lst)
+    return [item for item, count in counter.items() if count > 1]
+

augur/tasks/github/util/github_api_key_handler.py

@@ -3,11 +3,13 @@
 import random
 
 from typing import Optional, List
 
 from augur.tasks.util.redis_list import RedisList
 from augur.application.db.session import DatabaseSession
 from augur.application.config import AugurConfig
 
+RATE_LIMIT_URL = "https://api.github.com/rate_limit"
+
 
 class NoValidKeysError(Exception):
     pass
@@ -96,7 +98,7 @@
         attempts = 0
         while attempts < 3:
 
             try:
                 keys = self.get_api_keys_from_database()
                 break
             except:
@@ -143,17 +145,35 @@
             True if key is bad. False if the key is good
         """
 
-        # this endpoint allows us to check the rate limit, but it does not use one of our 5000 requests
-        url = "https://api.github.com/rate_limit"
-
         headers = {'Authorization': f'token {oauth_key}'}
 
-        data = client.request(method="GET", url=url, headers=headers, timeout=180).json()
+        data = client.request(method="GET", url=RATE_LIMIT_URL, headers=headers, timeout=180).json()
 
         try:
             if data["message"] == "Bad credentials":
                 return True
         except KeyError:
             pass
 
-        return False
+        return False
+
+    @staticmethod
+    def get_key_rate_limit(client, github_key):
+
+        headers = {'Authorization': f'token {github_key}'}
+
+        data = client.request(method="GET", url=RATE_LIMIT_URL, headers=headers, timeout=180).json()
+
+        if "message" in data:
+            return None, None
+
+        def convert_rate_limit_request(data):
+            return {
+            "requests_remaining": data["remaining"],
+            "reset_epoch": data["reset"]
+            }
+
+        core_data = convert_rate_limit_request(data["resources"]["core"])
+        graphql_data = convert_rate_limit_request(data["resources"]["graphql"])
+
+        return core_data, graphql_data