1# 2# Copyright (C) 2023 The Android Open Source Project 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# 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"""A Python interface to https://android.googlesource.com/tools/fetch_artifact/.""" 17import logging 18import urllib 19from collections.abc import AsyncIterable 20from logging import Logger 21from typing import cast 22 23from aiohttp import ClientSession 24 25_DEFAULT_QUERY_URL_BASE = "https://androidbuildinternal.googleapis.com" 26 27 28def _logger() -> Logger: 29 return logging.getLogger("fetchartifact") 30 31 32def _make_download_url( 33 target: str, 34 build_id: str, 35 artifact_name: str, 36 query_url_base: str, 37) -> str: 38 """Constructs the download URL. 39 40 Args: 41 target: Name of the build target from which to fetch the artifact. 42 build_id: ID of the build from which to fetch the artifact. 43 artifact_name: Name of the artifact to fetch. 44 45 Returns: 46 URL for the given artifact. 47 """ 48 # The Android build API does not handle / in artifact names, but urllib.parse.quote 49 # thinks those are safe by default. We need to escape them. 50 artifact_name = urllib.parse.quote(artifact_name, safe="") 51 return ( 52 f"{query_url_base}/android/internal/build/v3/builds/{build_id}/{target}/" 53 f"attempts/latest/artifacts/{artifact_name}/url" 54 ) 55 56 57async def fetch_artifact( 58 target: str, 59 build_id: str, 60 artifact_name: str, 61 session: ClientSession, 62 query_url_base: str = _DEFAULT_QUERY_URL_BASE, 63) -> bytes: 64 """Fetches an artifact from the build server. 65 66 Args: 67 target: Name of the build target from which to fetch the artifact. 68 build_id: ID of the build from which to fetch the artifact. 69 artifact_name: Name of the artifact to fetch. 70 session: The aiohttp ClientSession to use. If omitted, one will be created and 71 destroyed for every call. 72 query_url_base: The base of the endpoint used for querying download URLs. Uses 73 the android build service by default, but can be replaced for testing. 74 75 Returns: 76 The bytes of the downloaded artifact. 77 """ 78 download_url = _make_download_url(target, build_id, artifact_name, query_url_base) 79 _logger().debug("Beginning download from %s", download_url) 80 async with session.get(download_url) as response: 81 response.raise_for_status() 82 return await response.read() 83 84 85async def fetch_artifact_chunked( 86 target: str, 87 build_id: str, 88 artifact_name: str, 89 session: ClientSession, 90 chunk_size: int = 16 * 1024 * 1024, 91 query_url_base: str = _DEFAULT_QUERY_URL_BASE, 92) -> AsyncIterable[bytes]: 93 """Fetches an artifact from the build server. 94 95 Args: 96 target: Name of the build target from which to fetch the artifact. 97 build_id: ID of the build from which to fetch the artifact. 98 artifact_name: Name of the artifact to fetch. 99 session: The aiohttp ClientSession to use. If omitted, one will be created and 100 destroyed for every call. 101 query_url_base: The base of the endpoint used for querying download URLs. Uses 102 the android build service by default, but can be replaced for testing. 103 104 Returns: 105 Async iterable bytes of the artifact contents. 106 """ 107 download_url = _make_download_url(target, build_id, artifact_name, query_url_base) 108 _logger().debug("Beginning download from %s", download_url) 109 async with session.get(download_url) as response: 110 response.raise_for_status() 111 async for chunk in response.content.iter_chunked(chunk_size): 112 yield chunk 113