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" 26DEFAULT_CHUNK_SIZE = 16 * 1024 * 1024 27 28 29def _logger() -> Logger: 30 return logging.getLogger("fetchartifact") 31 32 33def _make_download_url( 34 target: str, 35 build_id: str, 36 artifact_name: str, 37 query_url_base: str, 38) -> str: 39 """Constructs the download URL. 40 41 Args: 42 target: Name of the build target from which to fetch the artifact. 43 build_id: ID of the build from which to fetch the artifact. 44 artifact_name: Name of the artifact to fetch. 45 46 Returns: 47 URL for the given artifact. 48 """ 49 # The Android build API does not handle / in artifact names, but urllib.parse.quote 50 # thinks those are safe by default. We need to escape them. 51 artifact_name = urllib.parse.quote(artifact_name, safe="") 52 return ( 53 f"{query_url_base}/android/internal/build/v3/builds/{build_id}/{target}/" 54 f"attempts/latest/artifacts/{artifact_name}/url" 55 ) 56 57 58async def fetch_artifact( 59 target: str, 60 build_id: str, 61 artifact_name: str, 62 session: ClientSession, 63 query_url_base: str = _DEFAULT_QUERY_URL_BASE, 64) -> bytes: 65 """Fetches an artifact from the build server. 66 67 Args: 68 target: Name of the build target from which to fetch the artifact. 69 build_id: ID of the build from which to fetch the artifact. 70 artifact_name: Name of the artifact to fetch. 71 session: The aiohttp ClientSession to use. If omitted, one will be created and 72 destroyed for every call. 73 query_url_base: The base of the endpoint used for querying download URLs. Uses 74 the android build service by default, but can be replaced for testing. 75 76 Returns: 77 The bytes of the downloaded artifact. 78 """ 79 download_url = _make_download_url(target, build_id, artifact_name, query_url_base) 80 _logger().debug("Beginning download from %s", download_url) 81 async with session.get(download_url) as response: 82 response.raise_for_status() 83 return await response.read() 84 85 86async def fetch_artifact_chunked( 87 target: str, 88 build_id: str, 89 artifact_name: str, 90 session: ClientSession, 91 chunk_size: int = DEFAULT_CHUNK_SIZE, 92 query_url_base: str = _DEFAULT_QUERY_URL_BASE, 93) -> AsyncIterable[bytes]: 94 """Fetches an artifact from the build server. 95 96 Args: 97 target: Name of the build target from which to fetch the artifact. 98 build_id: ID of the build from which to fetch the artifact. 99 artifact_name: Name of the artifact to fetch. 100 session: The aiohttp ClientSession to use. If omitted, one will be created and 101 destroyed for every call. 102 query_url_base: The base of the endpoint used for querying download URLs. Uses 103 the android build service by default, but can be replaced for testing. 104 105 Returns: 106 Async iterable bytes of the artifact contents. 107 """ 108 downloader = ArtifactDownloader(target, build_id, artifact_name, query_url_base) 109 async for chunk in downloader.download(session, chunk_size): 110 yield chunk 111 112 113class ArtifactDownloader: 114 """Similar to fetch_artifact_chunked but can be subclassed to report progress.""" 115 116 def __init__( 117 self, 118 target: str, 119 build_id: str, 120 artifact_name: str, 121 query_url_base: str = _DEFAULT_QUERY_URL_BASE, 122 ) -> None: 123 self.target = target 124 self.build_id = build_id 125 self.artifact_name = artifact_name 126 self.query_url_base = query_url_base 127 128 async def download( 129 self, session: ClientSession, chunk_size: int = DEFAULT_CHUNK_SIZE 130 ) -> AsyncIterable[bytes]: 131 """Fetches an artifact from the build server. 132 133 Once the Content-Length of the artifact is known, on_artifact_size will be 134 called. If no Content-Length header is present, on_artifact_size will not be 135 called. 136 137 After each chunk is yielded, after_chunk will be called. 138 139 Args: 140 session: The aiohttp ClientSession to use. If omitted, one will be created 141 and destroyed for every call. 142 chunk_size: The number of bytes to attempt to download before yielding and 143 reporting progress. 144 145 Returns: 146 Async iterable bytes of the artifact contents. 147 """ 148 download_url = _make_download_url( 149 self.target, self.build_id, self.artifact_name, self.query_url_base 150 ) 151 _logger().debug("Beginning download from %s", download_url) 152 async with session.get(download_url) as response: 153 response.raise_for_status() 154 try: 155 self.on_artifact_size(int(response.headers["Content-Length"])) 156 except KeyError: 157 pass 158 159 async for chunk in response.content.iter_chunked(chunk_size): 160 yield chunk 161 self.after_chunk(len(chunk)) 162 163 def on_artifact_size(self, size: int) -> None: 164 """Called once the artifact's Content-Length is known. 165 166 This can be overridden in a subclass to enable progress reporting. 167 168 If the artifact headers do not report the Content-Length, this function will not 169 be called. 170 171 Args: 172 size: The Content-Length of the artifact. 173 """ 174 175 def after_chunk(self, size: int) -> None: 176 """Called after each chunk is yielded. 177 178 This can be overridden in a subclass to enable progress reporting. 179 180 Args: 181 size: The size of the chunk. 182 """ 183