fix: add retries and URL encoding to Unmanaged Keys Audit workflow (#39188)

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

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:

infra/enforcement/test_sending.py

@@ -0,0 +1,103 @@
+# 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
+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()