diff --git a/augur/application/cli/github.py b/augur/application/cli/github.py new file mode 100644 index 000000000..cad13be79 --- /dev/null +++ b/augur/application/cli/github.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: MIT +import logging +import click +import sqlalchemy as s +from datetime import datetime +import httpx +from collections import Counter + +from augur.application.cli import test_connection, test_db_connection + +from augur.application.db.engine import DatabaseEngine +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] + diff --git a/augur/tasks/github/util/github_api_key_handler.py b/augur/tasks/github/util/github_api_key_handler.py index 4f8178e7c..9ca777a1b 100644 --- a/augur/tasks/github/util/github_api_key_handler.py +++ b/augur/tasks/github/util/github_api_key_handler.py @@ -9,6 +9,8 @@ from augur.application.db.lib import get_value, get_worker_oauth_keys from sqlalchemy import func +RATE_LIMIT_URL = "https://api.github.com/rate_limit" + class NoValidKeysError(Exception): pass @@ -152,12 +154,9 @@ def is_bad_api_key(self, client: httpx.Client, oauth_key: str) -> bool: 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": @@ -165,4 +164,25 @@ def is_bad_api_key(self, client: httpx.Client, oauth_key: str) -> bool: except KeyError: pass - return False \ No newline at end of file + 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 \ No newline at end of file