1#!/usr/bin/env python3
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 python3-pip
26  $ pip install --user --upgrade google-api-python-client
27  $ pip install --user --upgrade oauth2client
28"""
29    raise ImportError(missingImportString)
30
31import io
32import getpass
33import os
34
35import utils
36
37ANDROID_BUILD_API_SCOPE = (
38    'https://www.googleapis.com/auth/androidbuild.internal')
39ANDROID_BUILD_API_NAME = 'androidbuildinternal'
40ANDROID_BUILD_API_VERSION = 'v2beta1'
41CHUNK_SIZE = 10 * 1024 * 1024  # 10M
42
43ANDROID_PGO_BUILD = 'pgo-coral-config1'
44
45STUBBY_COMMAND_PATH = '/google/data/ro/teams/android-llvm/tests/sso_stubby_cmd.sh'
46STUBBY_REQUEST = """
47target: {{
48  scope: GAIA_USER
49  name: "{user}@google.com"
50}}
51target_credential: {{
52  type: OAUTH2_TOKEN
53  oauth2_attributes: {{
54    scope: '{scope}'
55  }}
56}}
57"""
58
59
60def _get_oauth2_token():
61    request = STUBBY_REQUEST.format(
62        user=getpass.getuser(), scope=ANDROID_BUILD_API_SCOPE)
63    with open(STUBBY_COMMAND_PATH) as stubby_command_file:
64        stubby_command = stubby_command_file.read().strip().split()
65    output = utils.check_output(stubby_command, input=request)
66    # output is of the format:
67    # oauth2_token: "<TOKEN>"
68    return output.split('"')[1]
69
70
71class AndroidBuildClient(object):
72
73    def __init__(self):
74        creds = oauth2_client.AccessTokenCredentials(
75            access_token=_get_oauth2_token(), user_agent='unused/1.0')
76
77        self.client = apiclient.discovery.build(
78            ANDROID_BUILD_API_NAME,
79            ANDROID_BUILD_API_VERSION,
80            credentials=creds,
81            discoveryServiceUrl=apiclient.discovery.DISCOVERY_URI)
82
83    # Get the latest test invocation for a given test_tag for a given build.
84    def get_invocation_id(self, build, test_tag):
85        request = self.client.testresult().list(
86            buildId=build, target=ANDROID_PGO_BUILD, attemptId='latest')
87
88        response = request.execute()
89        testResultWithTag = [
90            r for r in response['testResults'] if r['testTag'] == test_tag
91        ]
92        if len(testResultWithTag) != 1:
93            raise RuntimeError(
94                'Expected one test with tag {} for build {}.  Found {}.  Full response is {}'
95                .format(test_tag, build, len(testResultWithTag), response))
96        return testResultWithTag[0]['id']
97
98    # Get the full artifact name for the zipped PGO profiles
99    # (_data_local_tmp_pgo_<hash>.zip) for a given <build, test_tag,
100    # invocation_id>.
101    def get_test_artifact_name(self, build, test_tag, invocation_id):
102        request = self.client.testartifact().list(
103            buildType='submitted',
104            buildId=build,
105            target=ANDROID_PGO_BUILD,
106            attemptId='latest',
107            testResultId=invocation_id,
108            maxResults=100)
109
110        response = request.execute()
111        profile_zip = [
112            f for f in response['test_artifacts']
113            if f['name'].endswith('zip') and '_data_local_tmp_pgo_' in f['name']
114        ]
115        if len(profile_zip) != 1:
116            raise RuntimeError(
117                'Expected one matching zipfile for invocation {} of {} for build {}.  Found {} ({})'
118                .format(invocation_id, test_tag, build, len(profile_zip),
119                        ', '.join(profile_zip)))
120        return profile_zip[0]['name']
121
122    # Download the zipped PGO profiles for a given <build, test_tag,
123    # invocation_id, artifact_name> into <output_zip>.
124    def download_test_artifact(self, build, invocation_id, artifact_name,
125                               output_zip):
126        request = self.client.testartifact().get_media(
127            buildType='submitted',
128            buildId=build,
129            target=ANDROID_PGO_BUILD,
130            attemptId='latest',
131            testResultId=invocation_id,
132            resourceId=artifact_name)
133
134        f = io.FileIO(output_zip, 'wb')
135        try:
136            downloader = apiclient.http.MediaIoBaseDownload(
137                f, request, chunksize=CHUNK_SIZE)
138            done = False
139            while not done:
140                status, done = downloader.next_chunk()
141        except apiclient.errors.HttpError as e:
142            if e.resp.status == 404:
143                raise RuntimeError(
144                    'Artifact {} does not exist for invocation {} for build {}.'
145                    .format(artifact_name, invocation_id, build))
146
147    # For a <build, test_tag>, find the invocation_id, artifact_name and
148    # download the artifact into <output_dir>/pgo_profiles.zip.
149    def download_pgo_zip(self, build, test_tag, output_dir):
150        output_zip = os.path.join(output_dir, 'pgo_profiles.zip')
151
152        invocation_id = self.get_invocation_id(build, test_tag)
153        artifact_name = self.get_test_artifact_name(build, test_tag,
154                                                    invocation_id)
155        self.download_test_artifact(build, invocation_id, artifact_name,
156                                    output_zip)
157        return output_zip
158