1#!/usr/bin/python
2
3"""Automatically update the afe_stable_versions table.
4
5This command updates the stable repair version for selected boards
6in the lab.  For each board, if the version that Omaha is serving
7on the Beta channel for the board is more recent than the current
8stable version in the AFE database, then the AFE is updated to use
9the version on Omaha.
10
11The upgrade process is applied to every "managed board" in the test
12lab.  Generally, a managed board is a board with both spare and
13critical scheduling pools.
14
15See `autotest_lib.site_utils.lab_inventory` for the full definition
16of "managed board".
17
18The command accepts two mutually exclusive options determining
19how changes will be handled:
20  * With no options, the command will make RPC calls to the AFE to
21    update the state according to the rules.
22  * With the `--shell-mode` option, the command will print a series
23    of `atest` commands that will accomplish the changes.
24  * With the `--dry-run` option, the command will perform all normal
25    printing, but will skip actual RPC calls to change the database.
26
27The `--shell-mode` and `--dry-run` options are mutually exclusive.
28"""
29
30import argparse
31import json
32import subprocess
33import sys
34
35import common
36from autotest_lib.client.common_lib import utils
37from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
38from autotest_lib.site_utils import lab_inventory
39
40
41# _OMAHA_STATUS - URI of a file in GoogleStorage with a JSON object
42# summarizing all versions currently being served by Omaha.
43#
44# The principle data is in an array named 'omaha_data'.  Each entry
45# in the array contains information relevant to one image being
46# served by Omaha, including the following information:
47#   * The board name of the product, as known to Omaha.
48#   * The channel associated with the image.
49#   * The Chrome and Chrome OS version strings for the image
50#     being served.
51#
52_OMAHA_STATUS = 'gs://chromeos-build-release-console/omaha_status.json'
53
54# _SET_VERSION - `atest` command that will assign a specific board a
55# specific stable version in the AFE.
56#
57# _DELETE_VERSION - `atest` command that will delete a stable version
58# mapping from the AFE.
59#
60# _DEFAULT_BOARD - The distinguished board name used to identify a
61# stable version mapping that is used for any board without an explicit
62# mapping of its own.
63#
64_SET_VERSION = 'atest stable_version modify --board %s --version %s'
65_DELETE_VERSION = ('atest stable_version delete --no-confirmation '
66                   '--board %s')
67_DEFAULT_BOARD = 'DEFAULT'
68
69
70# Execution modes:
71#
72# _NORMAL_MODE:  no command line options.
73# _DRY_RUN: --dry-run on the command line.
74# _SHELL_MODE: --shell-mode on the command line.
75#
76_NORMAL_MODE = 0
77_DRY_RUN = 1
78_SHELL_MODE = 2
79
80
81def _get_omaha_board(json_entry):
82    """Get the board name from an 'omaha_data' entry.
83
84    @param json_entry   Deserialized JSON object for one entry of the
85                        'omaha_data' array.
86    @return Returns a version number in the form R##-####.#.#.
87    """
88    return json_entry['board']['public_codename']
89
90
91def _get_omaha_version(json_entry):
92    """Get a Chrome OS version string from an 'omaha_data' entry.
93
94    @param json_entry   Deserialized JSON object for one entry of the
95                        'omaha_data' array.
96    @return Returns a version number in the form R##-####.#.#.
97    """
98    milestone = json_entry['chrome_version'].split('.')[0]
99    build = json_entry['chrome_os_version']
100    return 'R%s-%s' % (milestone, build)
101
102
103def _get_omaha_versions():
104    """Get the current Beta versions serving on Omaha.
105
106    Returns a dictionary mapping board names to the currently preferred
107    version for the Beta channel as served by Omaha.  The board names
108    are the names as known to Omaha:  If the board name in the AFE has a
109    '_', the corresponding Omaha name uses a '-' instead.  The boards
110    mapped may include boards not in the list of managed boards in the
111    lab.
112
113    The beta channel versions are found by searching the `_OMAHA_STATUS`
114    file.  That file is calculated by GoldenEye from Omaha.  It's
115    accurate, but could be out-of-date for a small time window.
116
117    @return A dictionary mapping Omaha boards to Beta versions.
118    """
119    sp = subprocess.Popen(['gsutil', 'cat', _OMAHA_STATUS],
120                          stdout=subprocess.PIPE)
121    omaha_status = json.load(sp.stdout)
122    return {_get_omaha_board(e): _get_omaha_version(e)
123                for e in omaha_status['omaha_data']
124                if e['channel'] == 'beta'}
125
126
127def _get_upgrade_versions(afe_versions, omaha_versions, boards):
128    """Get the new stable versions to which we should update.
129
130    The new versions are returned as a tuple of a dictionary mapping
131    board names to versions, plus a new default board setting.  The
132    new default is determined as the most commonly used version
133    across the given boards.
134
135    The new dictionary will have a mapping for every board in `boards`.
136    That mapping will be taken from `afe_versions`, unless the board has
137    a mapping in `omaha_versions` _and_ the omaha version is more recent
138    than the AFE version.
139
140    @param afe_versions     The current board->version mappings in the
141                            AFE.
142    @param omaha_versions   The current board->version mappings from
143                            Omaha for the Beta channel.
144    @param boards           Set of boards to be upgraded.
145    @return Tuple of (mapping, default) where mapping is a dictionary
146            mapping boards to versions, and default is a version string.
147    """
148    upgrade_versions = {}
149    version_counts = {}
150    afe_default = afe_versions[_DEFAULT_BOARD]
151    for board in boards:
152        version = afe_versions.get(board, afe_default)
153        omaha_version = omaha_versions.get(board.replace('_', '-'))
154        if (omaha_version is not None and
155                utils.compare_versions(version, omaha_version) < 0):
156            version = omaha_version
157        upgrade_versions[board] = version
158        version_counts.setdefault(version, 0)
159        version_counts[version] += 1
160    return (upgrade_versions,
161            max(version_counts.items(), key=lambda x: x[1])[0])
162
163
164def _set_stable_version(afe, mode, board, version):
165    """Call the AFE to change a stable version mapping.
166
167    Setting the mapping for the distinguished board name
168    `_DEFAULT_BOARD` will change the default mapping for any board
169    that doesn't have its own mapping.
170
171    @param afe          AFE object for RPC calls.
172    @param mode         Mode indicating whether to print a shell
173                        command, call an RPC, or do nothing.
174    @param board        Update the mapping for this board.
175    @param version      Update the board to this version.
176    """
177    if mode == _SHELL_MODE:
178        print _SET_VERSION % (board, version)
179    elif mode == _NORMAL_MODE:
180        afe.run('set_stable_version', board=board, version=version)
181
182
183def _delete_stable_version(afe, mode, board):
184    """Call the AFE to delete a stable version mapping.
185
186    Deleting a mapping causes the board to revert to the current default
187    mapping in the AFE.
188
189    @param afe          AFE object for RPC calls.
190    @param mode         Mode indicating whether to print a shell
191                        command, call an RPC, or do nothing.
192    @param board        Delete the mapping for this board.
193    """
194    assert board != _DEFAULT_BOARD
195    if mode == _SHELL_MODE:
196        print _DELETE_VERSION % board
197    elif mode == _NORMAL_MODE:
198        afe.run('delete_stable_version', board=board)
199
200
201def _apply_upgrades(afe, mode, afe_versions,
202                    upgrade_versions, new_default):
203    """Change stable version mappings in the AFE.
204
205    Update the `afe_stable_versions` database table to have the new
206    settings indicated by `upgrade_versions` and `new_default`.  Order
207    the changes so that at any moment, every board is mapped either
208    according to the old or the new mapping.
209
210    @param afe                  AFE object for RPC calls.
211    @param mode                 Mode indicating whether the action is to
212                                print shell commands, do nothing, or
213                                actually make RPC calls for changes.
214    @param afe_versions         The current board->version mappings in
215                                the AFE.
216    @param upgrade_versions     The current board->version mappings from
217                                Omaha for the Beta channel.
218    @param new_default          The new default build for the AFE.
219    """
220    old_default = afe_versions[_DEFAULT_BOARD]
221    if mode != _SHELL_MODE and new_default != old_default:
222        print 'Default %s -> %s' % (old_default, new_default)
223        print 'Applying stable version changes:'
224    # N.B. The ordering here matters:  Any board that will have a
225    # non-default stable version must be updated _before_ we change the
226    # default mapping, below.
227    for board, build in upgrade_versions.items():
228        if build == new_default:
229            continue
230        if board in afe_versions and build == afe_versions[board]:
231            if mode == _SHELL_MODE:
232                message = '# Leave board %s at %s'
233            else:
234                message = '    %-22s (no change) -> %s'
235            print message % (board, build)
236        else:
237            if mode != _SHELL_MODE:
238                old_build = afe_versions.get(board, '(default)')
239                print '    %-22s %s -> %s' % (board, old_build, build)
240            _set_stable_version(afe, mode, board, build)
241    # At this point, all non-default mappings have been installed.
242    # If there's a new default mapping, make that change now, and delete
243    # any non-default mappings made obsolete by the update.
244    if new_default != old_default:
245        _set_stable_version(afe, mode, _DEFAULT_BOARD, new_default)
246    for board, build in upgrade_versions.items():
247        if board in afe_versions and build == new_default:
248            if mode != _SHELL_MODE:
249                print ('    %-22s %s -> (default)' %
250                       (board, afe_versions[board]))
251            _delete_stable_version(afe, mode, board)
252
253
254def _parse_command_line(argv):
255    """Parse the command line arguments.
256
257    Create an argument parser for this command's syntax, parse the
258    command line, and return the result of the ArgumentParser
259    parse_args() method.
260
261    @param argv Standard command line argument vector; argv[0] is
262                assumed to be the command name.
263    @return Result returned by ArgumentParser.parse_args().
264
265    """
266    parser = argparse.ArgumentParser(
267            prog=argv[0],
268            description='Update the stable repair version for all '
269                        'boards')
270    mode_group = parser.add_mutually_exclusive_group()
271    mode_group.add_argument('-x', '--shell-mode', dest='mode',
272                            action='store_const', const=_SHELL_MODE,
273                            help='print shell commands to make the '
274                                 'changes')
275    mode_group.add_argument('-n', '--dry-run', dest='mode',
276                            action='store_const', const=_DRY_RUN,
277                            help='print changes without executing them')
278    parser.add_argument('extra_boards', nargs='*', metavar='BOARD',
279                        help='Names of additional boards to be updated.')
280    arguments = parser.parse_args(argv[1:])
281    if not arguments.mode:
282        arguments.mode = _NORMAL_MODE
283    return arguments
284
285
286def main(argv):
287    """Standard main routine.
288
289    @param argv  Command line arguments including `sys.argv[0]`.
290    """
291    arguments = _parse_command_line(argv)
292    if arguments.mode == _DRY_RUN:
293        print 'Dry run; no changes will be made.'
294    afe = frontend_wrappers.RetryingAFE(server=None)
295    boards = (set(arguments.extra_boards) |
296              lab_inventory.get_managed_boards(afe))
297    # The 'get_all_stable_versions' RPC returns a dictionary mapping
298    # `_DEFAULT_BOARD` to the current default version, plus a set of
299    # non-default board -> version mappings.
300    afe_versions = afe.run('get_all_stable_versions')
301    upgrade_versions, new_default = (
302        _get_upgrade_versions(afe_versions,
303                              _get_omaha_versions(),
304                              boards))
305    _apply_upgrades(afe, arguments.mode, afe_versions,
306                    upgrade_versions, new_default)
307
308
309if __name__ == '__main__':
310    main(sys.argv)
311