From e65186e1bb07069e0a61ed8254a0bbc680fd40ae Mon Sep 17 00:00:00 2001 From: raman118 Date: Thu, 2 Jul 2026 04:52:25 +0530 Subject: [PATCH 1/2] fix: add retries and query parameter encoding for GitHub API requests (closes #39188) --- infra/enforcement/account_keys.py | 7 +-- infra/enforcement/sending.py | 66 +++++++++++++++++++---- infra/enforcement/test_sending.py | 88 +++++++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 13 deletions(-) create mode 100644 infra/enforcement/test_sending.py diff --git a/infra/enforcement/account_keys.py b/infra/enforcement/account_keys.py index 31c1354319d6..56ccbf654b03 100644 --- a/infra/enforcement/account_keys.py +++ b/infra/enforcement/account_keys.py @@ -164,14 +164,15 @@ def _get_all_live_service_accounts(self) -> List[str]: request.name = f"projects/{self.project_id}" try: - accounts = self.service_account_client.list_service_accounts(request=request) - self.logger.debug(f"Retrieved {len(accounts.accounts)} service accounts for project {self.project_id}") + accounts_pager = self.service_account_client.list_service_accounts(request=request) + accounts = list(accounts_pager) + self.logger.debug(f"Retrieved {len(accounts)} service accounts for project {self.project_id}") if not accounts: self.logger.warning(f"No service accounts found in project {self.project_id}.") return [] - return [self._normalize_account_email(account.email) for account in accounts.accounts if not account.disabled] + return [self._normalize_account_email(account.email) for account in accounts if not account.disabled] except Exception as e: self.logger.error(f"Failed to retrieve service accounts for project {self.project_id}: {e}") raise diff --git a/infra/enforcement/sending.py b/infra/enforcement/sending.py index 37de025a207f..bd9787b6ce87 100644 --- a/infra/enforcement/sending.py +++ b/infra/enforcement/sending.py @@ -59,25 +59,71 @@ def __init__(self, logger: logging.Logger, github_token: str, github_repo: str, self.logger = logger self.github_api_url = "https://api.github.com" - def _make_github_request(self, method: str, endpoint: str, json: Optional[dict] = None) -> requests.Response: + def _make_github_request(self, method: str, endpoint: str, json: Optional[dict] = None, params: Optional[dict] = None) -> requests.Response: """ - Makes a request to the GitHub API. + Makes a request to the GitHub API with retry logic for transient errors and rate limiting. Args: method (str): The HTTP method to use (e.g., "GET", "POST", "PATCH"). endpoint (str): The API endpoint to call. json (Optional[dict]): The JSON payload to send with the request. + params (Optional[dict]): The URL parameters to send with the request. Returns: requests.Response: The response from the API. """ + import time url = f"{self.github_api_url}/{endpoint}" - response = requests.request(method, url, headers=self.headers, json=json) + max_retries = 5 + backoff = 2 - if not response.ok: - self.logger.error(f"Failed GitHub API request to {endpoint}: {response.status_code} - {response.text}") - response.raise_for_status() - + for attempt in range(max_retries): + try: + response = requests.request(method, url, headers=self.headers, json=json, params=params) + + # Check for rate limiting / secondary rate limit (403, 429) + if response.status_code in [403, 429]: + retry_after = response.headers.get("Retry-After") + if retry_after: + sleep_seconds = int(retry_after) + else: + sleep_seconds = backoff + + self.logger.warning( + f"GitHub API rate limit hit ({response.status_code}) on {endpoint}. " + f"Retrying in {sleep_seconds} seconds... (Attempt {attempt + 1}/{max_retries})" + ) + time.sleep(sleep_seconds) + backoff *= 2 + continue + + # Check for transient server errors (500, 502, 503, 504) + if response.status_code in [500, 502, 503, 504]: + self.logger.warning( + f"GitHub API server error ({response.status_code}) on {endpoint}. " + f"Retrying in {backoff} seconds... (Attempt {attempt + 1}/{max_retries})" + ) + time.sleep(backoff) + backoff *= 2 + continue + + if not response.ok: + self.logger.error(f"Failed GitHub API request to {endpoint}: {response.status_code} - {response.text}") + response.raise_for_status() + + return response + except requests.exceptions.RequestException as e: + if attempt == max_retries - 1: + raise + self.logger.warning( + f"GitHub API request exception on {endpoint}: {e}. " + f"Retrying in {backoff} seconds... (Attempt {attempt + 1}/{max_retries})" + ) + time.sleep(backoff) + backoff *= 2 + + self.logger.error(f"Failed GitHub API request to {endpoint} after {max_retries} attempts.") + response.raise_for_status() return response 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: def _get_open_issues(self, title: str) -> List[GitHubIssue]: """ - Retrieves the number of open GitHub issues with a given title. + Retrieves the open GitHub issues with a given title. Args: title (str): The title of the GitHub issue. """ - endpoint = f"search/issues?q=is:issue+repo:{self.github_repo}+in:title+{title}+is:open" - response = self._make_github_request("GET", endpoint) + q = f'is:issue repo:{self.github_repo} in:title "{title}" is:open' + response = self._make_github_request("GET", "search/issues", params={"q": q}) issues = response.json().get('items', []) parsed_issues = [] for issue in issues: diff --git a/infra/enforcement/test_sending.py b/infra/enforcement/test_sending.py new file mode 100644 index 000000000000..70103b0fbcb7 --- /dev/null +++ b/infra/enforcement/test_sending.py @@ -0,0 +1,88 @@ +import unittest +from unittest.mock import patch, MagicMock +import logging +import requests +from sending import SendingClient + +class TestSendingClient(unittest.TestCase): + def setUp(self): + self.logger = logging.getLogger("TestLogger") + self.logger.setLevel(logging.DEBUG) + # Create a SendingClient with dummy values + self.client = SendingClient( + logger=self.logger, + github_token="dummy-token", + github_repo="apache/beam", + smtp_server="smtp.example.com", + smtp_port=587, + email="sender@example.com", + password="password" + ) + + @patch("requests.request") + def test_get_open_issues_flaky_retry(self, mock_request): + # We simulate a 403 secondary rate limit error on the first request, + # and a successful 200 response on the second request. + mock_response_fail = MagicMock() + mock_response_fail.status_code = 403 + mock_response_fail.ok = False + mock_response_fail.headers = {"Retry-After": "1"} + mock_response_fail.raise_for_status.side_effect = requests.exceptions.HTTPError("403 Client Error") + + mock_response_success = MagicMock() + mock_response_success.status_code = 200 + mock_response_success.ok = True + mock_response_success.json.return_value = { + "items": [ + { + "number": 1234, + "title": "[SECURITY] Action Required: Unmanaged Service Account Keys Detected", + "body": "Test body", + "state": "open", + "html_url": "https://github.com/apache/beam/issues/1234", + "created_at": "2026-07-02T00:00:00Z", + "updated_at": "2026-07-02T00:00:00Z" + } + ] + } + + mock_request.side_effect = [mock_response_fail, mock_response_success] + + # Call get_open_issues + issues = self.client._get_open_issues("[SECURITY] Action Required: Unmanaged Service Account Keys Detected") + + # Verify that two requests were made (one retry) + self.assertEqual(mock_request.call_count, 2) + self.assertEqual(len(issues), 1) + self.assertEqual(issues[0].number, 1234) + + @patch("requests.request") + def test_get_open_issues_query_format(self, mock_request): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.ok = True + mock_response.json.return_value = {"items": []} + mock_request.return_value = mock_response + + title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected" + self.client._get_open_issues(title) + + # Verify that the query parameter was passed correctly to requests + mock_request.assert_called_once() + args, kwargs = mock_request.call_args + + # Check that we passed params dict containing 'q' + self.assertIn("params", kwargs) + self.assertIn("q", kwargs["params"]) + + q = kwargs["params"]["q"] + # The query should specify the repo, title (quoted), state and type + self.assertIn('repo:apache/beam', q) + self.assertIn('is:issue', q) + self.assertIn('is:open', q) + self.assertIn('in:title', q) + self.assertIn(f'"{title}"', q) + + +if __name__ == "__main__": + unittest.main() From 7a16876ae37f57126bda32682c749a3e8eeae9e5 Mon Sep 17 00:00:00 2001 From: raman118 Date: Thu, 2 Jul 2026 05:03:20 +0530 Subject: [PATCH 2/2] fix: add Apache license header to test_sending.py --- infra/enforcement/test_sending.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/infra/enforcement/test_sending.py b/infra/enforcement/test_sending.py index 70103b0fbcb7..26d4080adec5 100644 --- a/infra/enforcement/test_sending.py +++ b/infra/enforcement/test_sending.py @@ -1,3 +1,18 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from unittest.mock import patch, MagicMock import logging