1#!/usr/bin/env python
2#
3# Copyright (C) 2019 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18try:
19    import apiclient.discovery
20    import apiclient.http
21    from oauth2client import client as oauth2_client
22except ImportError:
23    missingImportString = '''
24  Missing necessary libraries. Try doing the following:
25  $ sudo apt-get install python-pip
26  $ sudo pip install --upgrade google-api-python-client
27  $ sudo pip install --upgrade oauth2client
28'''
29    raise ImportError(missingImportString)
30
31import io
32import os
33
34ANDROID_BUILD_API_SCOPE = (
35    'https://www.googleapis.com/auth/androidbuild.internal')
36ANDROID_BUILD_API_NAME = 'androidbuildinternal'
37ANDROID_BUILD_API_VERSION = 'v2beta1'
38TRADEFED_KEY_FILE = '/google/data/ro/teams/tradefed/configs/tradefed.json'
39CHUNK_SIZE = 10 * 1024 * 1024  # 10M
40
41ANDROID_PGO_BUILD = 'pgo-coral-config1'
42
43
44class AndroidBuildClient(object):
45
46    def __init__(self):
47        creds = oauth2_client.GoogleCredentials.from_stream(TRADEFED_KEY_FILE)
48        self.creds = creds.create_scoped([ANDROID_BUILD_API_SCOPE])
49
50        self.client = apiclient.discovery.build(
51            ANDROID_BUILD_API_NAME,
52            ANDROID_BUILD_API_VERSION,
53            credentials=creds,
54            discoveryServiceUrl=apiclient.discovery.DISCOVERY_URI)
55
56    # Get the latest test invocation for a given test_tag for a given build.
57    def get_invocation_id(self, build, test_tag):
58        request = self.client.testresult().list(
59            buildId=build, target=ANDROID_PGO_BUILD, attemptId='latest')
60
61        response = request.execute()
62        testResultWithTag = [
63            r for r in response['testResults'] if r['testTag'] == test_tag
64        ]
65        if len(testResultWithTag) != 1:
66            raise RuntimeError(
67                'Expected one test with tag {} for build {}.  Found {}.  Full response is {}'
68                .format(test_tag, build, len(testResultWithTag), response))
69        return testResultWithTag[0]['id']
70
71    # Get the full artifact name for the zipped PGO profiles
72    # (_data_local_tmp_pgo_<hash>.zip) for a given <build, test_tag,
73    # invocation_id>.
74    def get_test_artifact_name(self, build, test_tag, invocation_id):
75        request = self.client.testartifact().list(
76            buildType='submitted',
77            buildId=build,
78            target=ANDROID_PGO_BUILD,
79            attemptId='latest',
80            testResultId=invocation_id,
81            maxResults=100)
82
83        response = request.execute()
84        profile_zip = [
85            f for f in response['test_artifacts']
86            if f['name'].endswith('zip') and '_data_local_tmp_pgo_' in f['name']
87        ]
88        if len(profile_zip) != 1:
89            raise RuntimeError(
90                'Expected one matching zipfile for invocation {} of {} for build {}.  Found {} ({})'
91                .format(invocation_id, test_tag, build, len(profile_zip),
92                        ', '.join(profile_zip)))
93        return profile_zip[0]['name']
94
95    # Download the zipped PGO profiles for a given <build, test_tag,
96    # invocation_id, artifact_name> into <output_zip>.
97    def download_test_artifact(self, build, invocation_id, artifact_name,
98                               output_zip):
99        request = self.client.testartifact().get_media(
100            buildType='submitted',
101            buildId=build,
102            target=ANDROID_PGO_BUILD,
103            attemptId='latest',
104            testResultId=invocation_id,
105            resourceId=artifact_name)
106
107        f = io.FileIO(output_zip, 'wb')
108        try:
109            downloader = apiclient.http.MediaIoBaseDownload(
110                f, request, chunksize=CHUNK_SIZE)
111            done = False
112            while not done:
113                status, done = downloader.next_chunk()
114        except apiclient.errors.HttpError as e:
115            if e.resp.status == 404:
116                raise RuntimeError(
117                    'Artifact {} does not exist for invocation {} for build {}.'
118                    .format(artifact_name, invocation_id, build))
119
120    # For a <build, test_tag>, find the invocation_id, artifact_name and
121    # download the artifact into <output_dir>/pgo_profiles.zip.
122    def download_pgo_zip(self, build, test_tag, output_dir):
123        output_zip = os.path.join(output_dir, 'pgo_profiles.zip')
124
125        invocation_id = self.get_invocation_id(build, test_tag)
126        artifact_name = self.get_test_artifact_name(build, test_tag,
127                                                    invocation_id)
128        self.download_test_artifact(build, invocation_id, artifact_name,
129                                    output_zip)
130        return output_zip
131