tools/vocabulary-migrate/vocabulary-migrate.py

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