diff --git a/main.py b/main.py index a22a472..5d6e59b 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,11 @@ from __future__ import annotations import asyncio +import re +import time import os from asyncio import Semaphore -from datetime import datetime +from datetime import datetime, timedelta from enum import Enum from fnmatch import fnmatch from sys import argv @@ -67,6 +69,21 @@ class PackageResponse(BaseModel): updated_at: datetime | None +async def wait_for_ratelimit(*, response: Response, eligible_for_secondary_limit: bool = False) -> None: + ratelimit_remaining = int(response.headers['x-ratelimit-remaining']) + if ratelimit_remaining == 0: + print('ratelimit exceeded') + ratelimit_reset = datetime.fromtimestamp(int(response.headers['x-ratelimit-reset'])) + delta = ratelimit_reset - datetime.now() + if delta > timedelta(0): + print(f'sleeping for {delta}s') + time.sleep(delta.total_seconds()) + print('done sleeping') + elif eligible_for_secondary_limit: + # https://docs.github.com/en/rest/guides/best-practices-for-integrators#dealing-with-secondary-rate-limits + time.sleep(1) + + async def list_org_packages(*, org_name: str, http_client: AsyncClient) -> list[PackageResponse]: """ List all packages, for an organization. @@ -120,11 +137,11 @@ async def list_org_package_versions( :param http_client: HTTP client. :return: List of image objects. """ - response = await http_client.get( - f'{BASE_URL}/orgs/{org_name}/packages/container/{image_name.encoded}/versions?per_page=100' + packages = await get_all_pages( + url=f'{BASE_URL}/orgs/{org_name}/packages/container/{image_name.encoded}/versions?per_page=100', + http_client=http_client, ) - response.raise_for_status() - return [PackageVersionResponse(**i) for i in response.json()] + return [PackageVersionResponse(**i) for i in packages] async def list_package_versions(*, image_name: ImageName, http_client: AsyncClient) -> list[PackageVersionResponse]: @@ -135,9 +152,34 @@ async def list_package_versions(*, image_name: ImageName, http_client: AsyncClie :param http_client: HTTP client. :return: List of image objects. """ - response = await http_client.get(f'{BASE_URL}/user/packages/container/{image_name.encoded}/versions?per_page=100') - response.raise_for_status() - return [PackageVersionResponse(**i) for i in response.json()] + packages = await get_all_pages( + url=f'{BASE_URL}/user/packages/container/{image_name.encoded}/versions?per_page=100', http_client=http_client + ) + return [PackageVersionResponse(**i) for i in packages] + + +async def get_all_pages(*, url: str, http_client: AsyncClient) -> list[dict]: + """ + Accumulate all pages of a paginated API endpoint. + + :param url: The full API URL + :param http_client: HTTP client. + :return: List of objects. + """ + result = [] + rel_regex = re.compile(r'<([^<>]*)>; rel="(\w+)"') + rels = {'next': url} + + while 'next' in rels: + response = await http_client.get(rels['next']) + response.raise_for_status() + result.extend(response.json()) + + rels = {rel: url for url, rel in rel_regex.findall(response.headers['link'])} + + await wait_for_ratelimit(response=response) + + return result def post_deletion_output(*, response: Response, image_name: ImageName, version_id: int) -> None: @@ -176,6 +218,7 @@ async def delete_org_package_versions( await semaphore.acquire() try: response = await http_client.delete(url) + await wait_for_ratelimit(response=response, eligible_for_secondary_limit=True) post_deletion_output(response=response, image_name=image_name, version_id=version_id) except TimeoutException as e: print(f'Request to delete {image_name.value} timed out with error `{e}`') @@ -198,6 +241,7 @@ async def delete_package_versions( await semaphore.acquire() try: response = await http_client.delete(url) + await wait_for_ratelimit(response=response, eligible_for_secondary_limit=True) post_deletion_output(response=response, image_name=image_name, version_id=version_id) except TimeoutException as e: print(f'Request to delete {image_name.value} timed out with error `{e}`') diff --git a/main_tests.py b/main_tests.py index 27e0258..355431d 100644 --- a/main_tests.py +++ b/main_tests.py @@ -30,9 +30,11 @@ from main import post_deletion_output mock_response = Mock() +mock_response.headers = {'x-ratelimit-remaining': '1', 'link': ''} mock_response.json.return_value = [] mock_response.is_error = False mock_bad_response = Mock() +mock_bad_response.headers = {'x-ratelimit-remaining': '1', 'link': ''} mock_bad_response.is_error = True mock_http_client = AsyncMock() mock_http_client.get.return_value = mock_response @@ -450,11 +452,13 @@ async def test_public_images_with_more_than_5000_downloads(mocker, capsys): them once at the end, with the necessary context to act on them if wanted. """ mock_delete_response = Mock() + mock_delete_response.headers = {'x-ratelimit-remaining': '1', 'link': ''} mock_delete_response.is_error = True mock_delete_response.status_code = 400 mock_delete_response.json = lambda: {'message': main.GITHUB_ASSISTANCE_MSG} mock_list_response = Mock() + mock_list_response.headers = {'x-ratelimit-remaining': '1', 'link': ''} mock_list_response.is_error = True mock_list_response.status_code = 400 @@ -554,6 +558,7 @@ def json(self): @pytest.mark.asyncio async def test_outputs_are_set(mocker): mock_list_response = Mock() + mock_list_response.headers = {'x-ratelimit-remaining': '1', 'link': ''} mock_list_response.is_error = True mock_list_response.status_code = 200 mock_list_response.json = lambda: [