Skip to content

Commit

Permalink
feat: add bypass to rate limiter
Browse files Browse the repository at this point in the history
Signed-off-by: Zack Koppert <[email protected]>
  • Loading branch information
zkoppert committed Oct 5, 2024
1 parent d77140c commit e64e4e7
Show file tree
Hide file tree
Showing 5 changed files with 51 additions and 5 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ This action can be configured to authenticate with GitHub App Installation or Pe
| field | required | default | description |
| ----------------------------- | -------- | ------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `GH_ENTERPRISE_URL` | False | `""` | URL of GitHub Enterprise instance to use for auth instead of github.com |
| `RATE_LIMIT_BYPASS` | False | `false` | If set to `true`, the rate limit will be bypassed. This is useful if being run on an local github server with rate limiting disabled. |
| `HIDE_AUTHOR` | False | False | If set to `true`, the author will not be displayed in the generated Markdown file. |
| `HIDE_ITEMS_CLOSED_COUNT` | False | False | If set to `true`, the number of items closed metric will not be displayed in the generated Markdown file. |
| `HIDE_LABEL_METRICS` | False | False | If set to `true`, the time in label metrics will not be displayed in the generated Markdown file. |
Expand Down
6 changes: 6 additions & 0 deletions config.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ class EnvVars:
non_mentioning_links (bool): If set to TRUE, links do not cause a notification in the desitnation repository
report_title (str): The title of the report
output_file (str): The name of the file to write the report to
rate_limit_bypass (bool): If set to TRUE, bypass the rate limit for the GitHub API
"""

def __init__(
Expand All @@ -68,6 +69,7 @@ def __init__(
non_mentioning_links: bool,
report_title: str,
output_file: str,
rate_limit_bypass: bool = False,
):
self.gh_app_id = gh_app_id
self.gh_app_installation_id = gh_app_installation_id
Expand All @@ -90,6 +92,7 @@ def __init__(
self.non_mentioning_links = non_mentioning_links
self.report_title = report_title
self.output_file = output_file
self.rate_limit_bypass = rate_limit_bypass

def __repr__(self):
return (
Expand All @@ -115,6 +118,7 @@ def __repr__(self):
f"{self.non_mentioning_links}"
f"{self.report_title}"
f"{self.output_file}"
f"{self.rate_limit_bypass}"
)


Expand Down Expand Up @@ -198,6 +202,7 @@ def get_env_vars(test: bool = False) -> EnvVars:

report_title = os.getenv("REPORT_TITLE", "Issue Metrics")
output_file = os.getenv("OUTPUT_FILE", "")
rate_limit_bypass = get_bool_env_var("RATE_LIMIT_BYPASS", False)

# Hidden columns
hide_author = get_bool_env_var("HIDE_AUTHOR", False)
Expand Down Expand Up @@ -234,4 +239,5 @@ def get_env_vars(test: bool = False) -> EnvVars:
non_mentioning_links,
report_title,
output_file,
rate_limit_bypass,
)
20 changes: 15 additions & 5 deletions issue_metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ def search_issues(
search_query: str,
github_connection: github3.GitHub,
owners_and_repositories: List[dict],
rate_limit_bypass: bool = False,
) -> List[github3.search.IssueSearchResult]: # type: ignore
"""
Searches for issues/prs/discussions in a GitHub repository that match
Expand All @@ -66,7 +67,13 @@ def search_issues(
"""

# Rate Limit Handling: API only allows 30 requests per minute
def wait_for_api_refresh(iterator: github3.structs.SearchIterator):
def wait_for_api_refresh(
iterator: github3.structs.SearchIterator, rate_limit_bypass: bool = False
):
# If the rate limit bypass is enabled, don't wait for the API to refresh
if rate_limit_bypass:
return

max_retries = 5
retry_count = 0
sleep_time = 70
Expand All @@ -90,7 +97,7 @@ def wait_for_api_refresh(iterator: github3.structs.SearchIterator):
issues_iterator = github_connection.search_issues(
search_query, per_page=issues_per_page
)
wait_for_api_refresh(issues_iterator)
wait_for_api_refresh(issues_iterator, rate_limit_bypass)

issues = []
repos_and_owners_string = ""
Expand All @@ -107,7 +114,7 @@ def wait_for_api_refresh(iterator: github3.structs.SearchIterator):

# requests are sent once per page of issues
if idx % issues_per_page == 0:
wait_for_api_refresh(issues_iterator)
wait_for_api_refresh(issues_iterator, rate_limit_bypass)

except github3.exceptions.ForbiddenError:
print(
Expand Down Expand Up @@ -288,7 +295,7 @@ def get_owners_and_repositories(
return results_list


def main():
def main(): # pragma: no cover
"""Run the issue-metrics script.
This function authenticates to GitHub, searches for issues/prs/discussions
Expand All @@ -312,6 +319,7 @@ def main():
non_mentioning_links = env_vars.non_mentioning_links
report_title = env_vars.report_title
output_file = env_vars.output_file
rate_limit_bypass = env_vars.rate_limit_bypass

gh_app_id = env_vars.gh_app_id
gh_app_installation_id = env_vars.gh_app_installation_id
Expand Down Expand Up @@ -363,7 +371,9 @@ def main():
write_to_markdown(None, None, None, None, None, None, None, None)
return
else:
issues = search_issues(search_query, github_connection, owners_and_repositories)
issues = search_issues(
search_query, github_connection, owners_and_repositories, rate_limit_bypass
)
if len(issues) <= 0:
print("No issues found")
write_to_markdown(None, None, None, None, None, None, None, None)
Expand Down
3 changes: 3 additions & 0 deletions test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ def setUp(self):
"OUTPUT_FILE",
"REPORT_TITLE",
"SEARCH_QUERY",
"RATE_LIMIT_BYPASS",
]
for key in env_keys:
if key in os.environ:
Expand All @@ -108,6 +109,7 @@ def setUp(self):
"OUTPUT_FILE": "",
"REPORT_TITLE": "",
"SEARCH_QUERY": SEARCH_QUERY,
"RATE_LIMIT_BYPASS": "false",
},
clear=True,
)
Expand Down Expand Up @@ -135,6 +137,7 @@ def test_get_env_vars_with_github_app(self):
False,
"",
"",
False,
)
result = get_env_vars(True)
self.assertEqual(str(result), str(expected_result))
Expand Down
26 changes: 26 additions & 0 deletions test_issue_metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,32 @@ def test_search_issues_with_just_owner_or_org(self):
issues = search_issues("is:open", mock_connection, owners)
self.assertEqual(issues, mock_issues)

def test_search_issues_with_just_owner_or_org_with_bypass(self):
"""Test that search_issues with just an owner/org returns the correct issues."""

# Set up the mock GitHub connection object
mock_issues = [
MagicMock(title="Issue 1"),
MagicMock(title="Issue 2"),
MagicMock(title="Issue 3"),
]

# simulating github3.structs.SearchIterator return value
mock_search_result = MagicMock()
mock_search_result.__iter__.return_value = iter(mock_issues)
mock_search_result.ratelimit_remaining = 30

mock_connection = MagicMock()
mock_connection.search_issues.return_value = mock_search_result

# Call search_issues and check that it returns the correct issues
org = {"owner": "org1"}
owners = [org]
issues = search_issues(
"is:open", mock_connection, owners, rate_limit_bypass=True
)
self.assertEqual(issues, mock_issues)


class TestGetOwnerAndRepository(unittest.TestCase):
"""Unit tests for the get_owners_and_repositories function.
Expand Down

0 comments on commit e64e4e7

Please sign in to comment.