Skip to content

Commit 0d1f25d

Browse files
Amar3ttoAmar3tto
authored
Merge pull request #39196 from raman118/fix/issue-39188-github-api-retries
fix: add retries and URL encoding to Unmanaged Keys Audit workflow (#39188)
2 parents d322070 + 7a16876 commit 0d1f25d

3 files changed

Lines changed: 163 additions & 13 deletions

File tree

‎infra/enforcement/account_keys.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -164,14 +164,15 @@ def _get_all_live_service_accounts(self) -> List[str]:
164164
request.name = f"projects/{self.project_id}"
165165

166166
try:
167-
accounts = self.service_account_client.list_service_accounts(request=request)
168-
self.logger.debug(f"Retrieved {len(accounts.accounts)} service accounts for project {self.project_id}")
167+
accounts_pager = self.service_account_client.list_service_accounts(request=request)
168+
accounts = list(accounts_pager)
169+
self.logger.debug(f"Retrieved {len(accounts)} service accounts for project {self.project_id}")
169170

170171
if not accounts:
171172
self.logger.warning(f"No service accounts found in project {self.project_id}.")
172173
return []
173174

174-
return [self._normalize_account_email(account.email) for account in accounts.accounts if not account.disabled]
175+
return [self._normalize_account_email(account.email) for account in accounts if not account.disabled]
175176
except Exception as e:
176177
self.logger.error(f"Failed to retrieve service accounts for project {self.project_id}: {e}")
177178
raise

‎infra/enforcement/sending.py‎

Lines changed: 56 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -59,25 +59,71 @@ def __init__(self, logger: logging.Logger, github_token: str, github_repo: str,
5959
self.logger = logger
6060
self.github_api_url = "https://api.github.com"
6161

62-
def _make_github_request(self, method: str, endpoint: str, json: Optional[dict] = None) -> requests.Response:
62+
def _make_github_request(self, method: str, endpoint: str, json: Optional[dict] = None, params: Optional[dict] = None) -> requests.Response:
6363
"""
64-
Makes a request to the GitHub API.
64+
Makes a request to the GitHub API with retry logic for transient errors and rate limiting.
6565
6666
Args:
6767
method (str): The HTTP method to use (e.g., "GET", "POST", "PATCH").
6868
endpoint (str): The API endpoint to call.
6969
json (Optional[dict]): The JSON payload to send with the request.
70+
params (Optional[dict]): The URL parameters to send with the request.
7071
7172
Returns:
7273
requests.Response: The response from the API.
7374
"""
75+
import time
7476
url = f"{self.github_api_url}/{endpoint}"
75-
response = requests.request(method, url, headers=self.headers, json=json)
77+
max_retries = 5
78+
backoff = 2
7679

77-
if not response.ok:
78-
self.logger.error(f"Failed GitHub API request to {endpoint}: {response.status_code} - {response.text}")
79-
response.raise_for_status()
80-
80+
for attempt in range(max_retries):
81+
try:
82+
response = requests.request(method, url, headers=self.headers, json=json, params=params)
83+
84+
# Check for rate limiting / secondary rate limit (403, 429)
85+
if response.status_code in [403, 429]:
86+
retry_after = response.headers.get("Retry-After")
87+
if retry_after:
88+
sleep_seconds = int(retry_after)
89+
else:
90+
sleep_seconds = backoff
91+
92+
self.logger.warning(
93+
f"GitHub API rate limit hit ({response.status_code}) on {endpoint}. "
94+
f"Retrying in {sleep_seconds} seconds... (Attempt {attempt + 1}/{max_retries})"
95+
)
96+
time.sleep(sleep_seconds)
97+
backoff *= 2
98+
continue
99+
100+
# Check for transient server errors (500, 502, 503, 504)
101+
if response.status_code in [500, 502, 503, 504]:
102+
self.logger.warning(
103+
f"GitHub API server error ({response.status_code}) on {endpoint}. "
104+
f"Retrying in {backoff} seconds... (Attempt {attempt + 1}/{max_retries})"
105+
)
106+
time.sleep(backoff)
107+
backoff *= 2
108+
continue
109+
110+
if not response.ok:
111+
self.logger.error(f"Failed GitHub API request to {endpoint}: {response.status_code} - {response.text}")
112+
response.raise_for_status()
113+
114+
return response
115+
except requests.exceptions.RequestException as e:
116+
if attempt == max_retries - 1:
117+
raise
118+
self.logger.warning(
119+
f"GitHub API request exception on {endpoint}: {e}. "
120+
f"Retrying in {backoff} seconds... (Attempt {attempt + 1}/{max_retries})"
121+
)
122+
time.sleep(backoff)
123+
backoff *= 2
124+
125+
self.logger.error(f"Failed GitHub API request to {endpoint} after {max_retries} attempts.")
126+
response.raise_for_status()
81127
return response
82128

83129
def _send_email(self, title: str, body: str, recipient: str) -> None:
@@ -97,13 +143,13 @@ def _send_email(self, title: str, body: str, recipient: str) -> None:
97143

98144
def _get_open_issues(self, title: str) -> List[GitHubIssue]:
99145
"""
100-
Retrieves the number of open GitHub issues with a given title.
146+
Retrieves the open GitHub issues with a given title.
101147
102148
Args:
103149
title (str): The title of the GitHub issue.
104150
"""
105-
endpoint = f"search/issues?q=is:issue+repo:{self.github_repo}+in:title+{title}+is:open"
106-
response = self._make_github_request("GET", endpoint)
151+
q = f'is:issue repo:{self.github_repo} in:title "{title}" is:open'
152+
response = self._make_github_request("GET", "search/issues", params={"q": q})
107153
issues = response.json().get('items', [])
108154
parsed_issues = []
109155
for issue in issues:

‎infra/enforcement/test_sending.py‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one or more
2+
# contributor license agreements. See the NOTICE file distributed with
3+
# this work for additional information regarding copyright ownership.
4+
# The ASF licenses this file to You under the Apache License, Version 2.0
5+
# (the "License"); you may not use this file except in compliance with
6+
# the License. You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
import unittest
17+
from unittest.mock import patch, MagicMock
18+
import logging
19+
import requests
20+
from sending import SendingClient
21+
22+
class TestSendingClient(unittest.TestCase):
23+
def setUp(self):
24+
self.logger = logging.getLogger("TestLogger")
25+
self.logger.setLevel(logging.DEBUG)
26+
# Create a SendingClient with dummy values
27+
self.client = SendingClient(
28+
logger=self.logger,
29+
github_token="dummy-token",
30+
github_repo="apache/beam",
31+
smtp_server="smtp.example.com",
32+
smtp_port=587,
33+
email="sender@example.com",
34+
password="password"
35+
)
36+
37+
@patch("requests.request")
38+
def test_get_open_issues_flaky_retry(self, mock_request):
39+
# We simulate a 403 secondary rate limit error on the first request,
40+
# and a successful 200 response on the second request.
41+
mock_response_fail = MagicMock()
42+
mock_response_fail.status_code = 403
43+
mock_response_fail.ok = False
44+
mock_response_fail.headers = {"Retry-After": "1"}
45+
mock_response_fail.raise_for_status.side_effect = requests.exceptions.HTTPError("403 Client Error")
46+
47+
mock_response_success = MagicMock()
48+
mock_response_success.status_code = 200
49+
mock_response_success.ok = True
50+
mock_response_success.json.return_value = {
51+
"items": [
52+
{
53+
"number": 1234,
54+
"title": "[SECURITY] Action Required: Unmanaged Service Account Keys Detected",
55+
"body": "Test body",
56+
"state": "open",
57+
"html_url": "https://github.com/apache/beam/issues/1234",
58+
"created_at": "2026-07-02T00:00:00Z",
59+
"updated_at": "2026-07-02T00:00:00Z"
60+
}
61+
]
62+
}
63+
64+
mock_request.side_effect = [mock_response_fail, mock_response_success]
65+
66+
# Call get_open_issues
67+
issues = self.client._get_open_issues("[SECURITY] Action Required: Unmanaged Service Account Keys Detected")
68+
69+
# Verify that two requests were made (one retry)
70+
self.assertEqual(mock_request.call_count, 2)
71+
self.assertEqual(len(issues), 1)
72+
self.assertEqual(issues[0].number, 1234)
73+
74+
@patch("requests.request")
75+
def test_get_open_issues_query_format(self, mock_request):
76+
mock_response = MagicMock()
77+
mock_response.status_code = 200
78+
mock_response.ok = True
79+
mock_response.json.return_value = {"items": []}
80+
mock_request.return_value = mock_response
81+
82+
title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected"
83+
self.client._get_open_issues(title)
84+
85+
# Verify that the query parameter was passed correctly to requests
86+
mock_request.assert_called_once()
87+
args, kwargs = mock_request.call_args
88+
89+
# Check that we passed params dict containing 'q'
90+
self.assertIn("params", kwargs)
91+
self.assertIn("q", kwargs["params"])
92+
93+
q = kwargs["params"]["q"]
94+
# The query should specify the repo, title (quoted), state and type
95+
self.assertIn('repo:apache/beam', q)
96+
self.assertIn('is:issue', q)
97+
self.assertIn('is:open', q)
98+
self.assertIn('in:title', q)
99+
self.assertIn(f'"{title}"', q)
100+
101+
102+
if __name__ == "__main__":
103+
unittest.main()

0 commit comments

Comments
 (0)