Merge branch '21258-flaky-adc-test'
[arvados.git] / tools / vocabulary-migrate / vocabulary-migrate.py
1 #!/usr/bin/env python3
2 #
3 # Copyright (C) The Arvados Authors. All rights reserved.
4 #
5 # SPDX-License-Identifier: CC-BY-SA-3.0
6
7 import argparse
8 import copy
9 import json
10 import logging
11 import os
12 import sys
13
14 import arvados
15 import arvados.util
16
17 logger = logging.getLogger('arvados.vocabulary_migrate')
18 logger.setLevel(logging.INFO)
19
20 class VocabularyError(Exception):
21     pass
22
23 opts = argparse.ArgumentParser(add_help=False)
24 opts.add_argument('--vocabulary-file', type=str, metavar='PATH', required=True,
25                   help="""
26 Use vocabulary definition file at PATH for migration decisions.
27 """)
28 opts.add_argument('--dry-run', action='store_true', default=False,
29                   help="""
30 Don't actually migrate properties, but only check if any collection/project
31 should be migrated.
32 """)
33 opts.add_argument('--debug', action='store_true', default=False,
34                   help="""
35 Sets logging level to DEBUG.
36 """)
37 arg_parser = argparse.ArgumentParser(
38     description='Migrate collections & projects properties to the new vocabulary format.',
39     parents=[opts])
40
41 def parse_arguments(arguments):
42     args = arg_parser.parse_args(arguments)
43     if args.debug:
44         logger.setLevel(logging.DEBUG)
45     if not os.path.isfile(args.vocabulary_file):
46         arg_parser.error("{} doesn't exist or isn't a file.".format(args.vocabulary_file))
47     return args
48
49 def _label_to_id_mappings(data, obj_name):
50     result = {}
51     for obj_id, obj_data in data.items():
52         for lbl in obj_data['labels']:
53             obj_lbl = lbl['label']
54             if obj_lbl not in result:
55                 result[obj_lbl] = obj_id
56             else:
57                 raise VocabularyError('{} label "{}" for {} ID "{}" already seen at {} ID "{}".'.format(obj_name, obj_lbl, obj_name, obj_id, obj_name, result[obj_lbl]))
58     return result
59
60 def key_labels_to_ids(vocab):
61     return _label_to_id_mappings(vocab['tags'], 'key')
62
63 def value_labels_to_ids(vocab, key_id):
64     if key_id in vocab['tags'] and 'values' in vocab['tags'][key_id]:
65         return _label_to_id_mappings(vocab['tags'][key_id]['values'], 'value')
66     return {}
67
68 def migrate_properties(properties, key_map, vocab):
69     result = {}
70     for k, v in properties.items():
71         key = key_map.get(k, k)
72         value = value_labels_to_ids(vocab, key).get(v, v)
73         result[key] = value
74     return result
75
76 def main(arguments=None):
77     args = parse_arguments(arguments)
78     vocab = None
79     with open(args.vocabulary_file, 'r') as f:
80         vocab = json.load(f)
81     arv = arvados.api('v1')
82     if 'tags' not in vocab or vocab['tags'] == {}:
83         logger.warning('Empty vocabulary file, exiting.')
84         return 1
85     if not arv.users().current().execute()['is_admin']:
86         logger.error('Admin privileges required.')
87         return 1
88     key_label_to_id_map = key_labels_to_ids(vocab)
89     migrated_counter = 0
90
91     for key_label in key_label_to_id_map:
92         logger.debug('Querying objects with property key "{}"'.format(key_label))
93         for resource in [arv.collections(), arv.groups()]:
94             objs = arvados.util.list_all(
95                 resource.list,
96                 order=['created_at'],
97                 select=['uuid', 'properties'],
98                 filters=[['properties', 'exists', key_label]]
99             )
100             for o in objs:
101                 props = copy.copy(o['properties'])
102                 migrated_props = migrate_properties(props, key_label_to_id_map, vocab)
103                 if not args.dry_run:
104                     logger.debug('Migrating {}: {} -> {}'.format(o['uuid'], props, migrated_props))
105                     arv.collections().update(uuid=o['uuid'], body={
106                         'properties': migrated_props
107                     }).execute()
108                 else:
109                     logger.info('Should migrate {}: {} -> {}'.format(o['uuid'], props, migrated_props))
110                 migrated_counter += 1
111                 if not args.dry_run and migrated_counter % 100 == 0:
112                     logger.info('Migrating {} objects...'.format(migrated_counter))
113
114     if args.dry_run and migrated_counter == 0:
115         logger.info('Nothing to do.')
116     elif not args.dry_run:
117         logger.info('Done, total objects migrated: {}.'.format(migrated_counter))
118     return 0
119
120 if __name__ == "__main__":
121     sys.exit(main())