1#!/usr/bin/env -S python3 -B
2#
3# Copyright (C) 2021 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"""Downloads ART Module prebuilts and creates CLs to update them in git."""
18
19import argparse
20import collections
21import os
22import re
23import subprocess
24import sys
25import tempfile
26
27
28# Prebuilt description used in commit message
29PREBUILT_DESCR = "ART Module"
30
31# fetch_artifact branch and targets
32BRANCH = "aosp-master-art"
33MODULE_TARGET = "DOES_NOT_EXIST" # There is currently no CI build in AOSP.
34SDK_TARGET = "mainline_modules_sdks"
35
36# Where to install the APEX modules
37MODULE_PATH = "packages/modules/ArtPrebuilt"
38
39# Where to install the SDKs and module exports
40SDK_PATH = "prebuilts/module_sdk/art"
41
42SDK_VERSION = "current"
43
44# Paths to git projects to prepare CLs in
45GIT_PROJECT_ROOTS = [MODULE_PATH, SDK_PATH]
46
47SCRIPT_PATH = MODULE_PATH + "/update-art-module-prebuilts.py"
48
49
50InstallEntry = collections.namedtuple("InstallEntry", [
51    # Artifact path in the build, passed to fetch_target
52    "source_path",
53    # Local install path
54    "install_path",
55    # True if this is a module SDK, to be skipped by --skip-module-sdk.
56    "module_sdk",
57    # True if the entry is a zip file that should be unzipped to install_path
58    "install_unzipped",
59])
60
61
62def install_apks_entry(apex_name):
63  return [InstallEntry(
64      os.path.join(apex_name + ".apks"),
65      os.path.join(MODULE_PATH, apex_name + ".apks"),
66      module_sdk=False,
67      install_unzipped=False)]
68
69
70def install_sdk_entries(apex_name, mainline_sdk_name, sdk_dir):
71  return [InstallEntry(
72      os.path.join("mainline-sdks",
73                   "for-latest-build",
74                   SDK_VERSION,
75                   apex_name,
76                   sdk_dir,
77                   mainline_sdk_name + "-" + SDK_VERSION + ".zip"),
78      os.path.join(SDK_PATH, SDK_VERSION, sdk_dir),
79      module_sdk=True,
80      install_unzipped=True)]
81
82
83install_entries = (
84    install_apks_entry("com.android.art") +
85    install_sdk_entries("com.android.art",
86                        "art-module-sdk", "sdk") +
87    install_sdk_entries("com.android.art",
88                        "art-module-host-exports", "host-exports") +
89    install_sdk_entries("com.android.art",
90                        "art-module-test-exports", "test-exports")
91)
92
93
94def check_call(cmd, **kwargs):
95  """Proxy for subprocess.check_call with logging."""
96  msg = " ".join(cmd) if isinstance(cmd, list) else cmd
97  if "cwd" in kwargs:
98    msg = "In " + kwargs["cwd"] + ": " + msg
99  print(msg)
100  subprocess.check_call(cmd, **kwargs)
101
102
103def fetch_artifact(branch, target, build, fetch_pattern, local_dir):
104  """Fetches artifact from the build server."""
105  fetch_artifact_path = "/google/data/ro/projects/android/fetch_artifact"
106  cmd = [fetch_artifact_path, "--branch", branch, "--target", target,
107         "--bid", build, fetch_pattern]
108  check_call(cmd, cwd=local_dir)
109
110
111def start_branch(git_branch_name, git_dirs):
112  """Creates a new repo branch in the given projects."""
113  check_call(["repo", "start", git_branch_name] + git_dirs)
114  # In case the branch already exists we reset it to upstream, to get a clean
115  # update CL.
116  for git_dir in git_dirs:
117    check_call(["git", "reset", "--hard", "@{upstream}"], cwd=git_dir)
118
119
120def upload_branch(git_root, git_branch_name):
121  """Uploads the CLs in the given branch in the given project."""
122  # Set the branch as topic to bundle with the CLs in other git projects (if
123  # any).
124  check_call(["repo", "upload", "-t", "--br=" + git_branch_name, git_root])
125
126
127def remove_files(git_root, subpaths, stage_removals):
128  """Removes files in the work tree, optionally staging them in git."""
129  if stage_removals:
130    check_call(["git", "rm", "-qrf", "--ignore-unmatch"] + subpaths, cwd=git_root)
131  # Need a plain rm afterwards even if git rm was executed, because git won't
132  # remove directories if they have non-git files in them.
133  check_call(["rm", "-rf"] + subpaths, cwd=git_root)
134
135
136def commit(git_root, prebuilt_descr, branch, target, build, add_paths, bug_number):
137  """Commits the new prebuilts."""
138  check_call(["git", "add"] + add_paths, cwd=git_root)
139
140  if build:
141    message = (
142        "Update {prebuilt_descr} prebuilts to build {build}.\n\n"
143        "Taken from branch {branch}, target {target}."
144        .format(prebuilt_descr=prebuilt_descr, branch=branch, target=target,
145                build=build))
146  else:
147    message = (
148        "DO NOT SUBMIT: Update {prebuilt_descr} prebuilts from local build."
149        .format(prebuilt_descr=prebuilt_descr))
150  message += ("\n\nCL prepared by {}."
151              "\n\nTest: Presubmits".format(SCRIPT_PATH))
152  if bug_number:
153    message += ("\nBug: {}".format(bug_number))
154  msg_fd, msg_path = tempfile.mkstemp()
155  with os.fdopen(msg_fd, "w") as f:
156    f.write(message)
157
158  # Do a diff first to skip the commit without error if there are no changes to
159  # commit.
160  check_call("git diff-index --quiet --cached HEAD -- || "
161             "git commit -F " + msg_path, shell=True, cwd=git_root)
162  os.unlink(msg_path)
163
164
165def install_entry(branch, target, build, local_dist, entry):
166  """Installs one file specified by entry."""
167
168  install_dir, install_file = os.path.split(entry.install_path)
169  if install_dir and not os.path.exists(install_dir):
170    os.makedirs(install_dir)
171
172  if build:
173    fetch_artifact(branch, target, build, entry.source_path, install_dir)
174  else:
175    check_call(["cp", os.path.join(local_dist, entry.source_path), install_dir])
176  source_file = os.path.basename(entry.source_path)
177
178  if entry.install_unzipped:
179    check_call(["mkdir", install_file], cwd=install_dir)
180    # Add -DD to not extract timestamps that may confuse the build system.
181    check_call(["unzip", "-DD", source_file, "-d", install_file],
182               cwd=install_dir)
183    check_call(["rm", source_file], cwd=install_dir)
184
185  elif source_file != install_file:
186    check_call(["mv", source_file, install_file], cwd=install_dir)
187
188
189def install_paths_per_git_root(roots, paths):
190  """Partitions the given paths into subpaths within the given roots.
191
192  Args:
193    roots: List of root paths.
194    paths: List of paths relative to the same directory as the root paths.
195
196  Returns:
197    A dict mapping each root to the subpaths under it. It's an error if some
198    path doesn't go into any root.
199  """
200  res = collections.defaultdict(list)
201  for path in paths:
202    found = False
203    for root in roots:
204      if path.startswith(root + "/"):
205        res[root].append(path[len(root) + 1:])
206        found = True
207        break
208    if not found:
209      sys.exit("Install path {} is not in any of the git roots: {}"
210               .format(path, " ".join(roots)))
211  return res
212
213
214def get_args():
215  """Parses and returns command line arguments."""
216  parser = argparse.ArgumentParser(
217      epilog="Either --build or --local-dist is required.")
218
219  parser.add_argument("--branch", default=BRANCH,
220                      help="Branch to fetch, defaults to " + BRANCH)
221  parser.add_argument("--module-target", default=MODULE_TARGET,
222                      help="Target to fetch modules from, defaults to " +
223                      MODULE_TARGET)
224  parser.add_argument("--sdk-target", default=SDK_TARGET,
225                      help="Target to fetch SDKs from, defaults to " +
226                      SDK_TARGET)
227  parser.add_argument("--build", metavar="NUMBER",
228                      help="Build number to fetch")
229  parser.add_argument("--local-dist", metavar="PATH",
230                      help="Take prebuilts from this local dist dir instead of "
231                      "using fetch_artifact")
232  parser.add_argument("--skip-apex", default=True, action="store_true",
233                      help="Do not fetch .apex files. Defaults to true.")
234  parser.add_argument("--skip-module-sdk", action="store_true",
235                      help="Do not fetch and unpack sdk and module_export zips.")
236  parser.add_argument("--skip-cls", action="store_true",
237                      help="Do not create branches or git commits")
238  parser.add_argument("--bug", metavar="NUMBER",
239                      help="Add a 'Bug' line with this number to commit "
240                      "messages.")
241  parser.add_argument("--upload", action="store_true",
242                      help="Upload the CLs to Gerrit")
243
244  args = parser.parse_args()
245  if ((not args.build and not args.local_dist) or
246      (args.build and args.local_dist)):
247    sys.exit(parser.format_help())
248  return args
249
250
251def main():
252  """Program entry point."""
253  args = get_args()
254
255  if any(path for path in GIT_PROJECT_ROOTS if not os.path.exists(path)):
256    sys.exit("This script must be run in the root of the Android build tree.")
257
258  entries = install_entries
259  if args.skip_apex:
260    entries = [entry for entry in entries if entry.module_sdk]
261  if args.skip_module_sdk:
262    entries = [entry for entry in entries if not entry.module_sdk]
263  if not entries:
264    sys.exit("Both APEXes and SDKs skipped - nothing to do.")
265
266  install_paths = [entry.install_path for entry in entries]
267  install_paths_per_root = install_paths_per_git_root(
268      GIT_PROJECT_ROOTS, install_paths)
269
270  git_branch_name = PREBUILT_DESCR.lower().replace(" ", "-") + "-update"
271  if args.build:
272    git_branch_name += "-" + args.build
273
274  if not args.skip_cls:
275    git_paths = list(install_paths_per_root.keys())
276    start_branch(git_branch_name, git_paths)
277
278  for git_root, subpaths in install_paths_per_root.items():
279    remove_files(git_root, subpaths, not args.skip_cls)
280  for entry in entries:
281    target = args.sdk_target if entry.module_sdk else args.module_target
282    install_entry(args.branch, target, args.build, args.local_dist, entry)
283
284  if not args.skip_cls:
285    for git_root, subpaths in install_paths_per_root.items():
286      target = args.sdk_target if git_root == SDK_PATH else args.module_target
287      commit(git_root, PREBUILT_DESCR, args.branch, target, args.build, subpaths,
288             args.bug)
289
290    if args.upload:
291      # Don't upload all projects in a single repo upload call, because that
292      # makes it pop up an interactive editor.
293      for git_root in install_paths_per_root:
294        upload_branch(git_root, git_branch_name)
295
296
297if __name__ == "__main__":
298  main()
299