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 opts = argparse.ArgumentParser(add_help=False)
  22 opts.add_argument('--vocabulary-file', type=str, metavar='PATH', default=None,
  23                   help="""
  24 Use vocabulary definition file at PATH for migration decisions.
  25 """)
  26 opts.add_argument('--dry-run', action='store_true', default=False,
  27                   help="""
  28 Don't actually migrate properties, but only check if any collection/project
  29 should be migrated.
  30 """)
  31 opts.add_argument('--debug', action='store_true', default=False,
  32                   help="""
  33 Sets logging level to DEBUG.
  34 """)
  35 arg_parser = argparse.ArgumentParser(
  36     description='Migrate collections & projects properties to the new vocabulary format.',
  37     parents=[opts])
  38
  39 def parse_arguments(arguments):
  40     args = arg_parser.parse_args(arguments)
  41     if args.debug:
  42         logger.setLevel(logging.DEBUG)
  43     if args.vocabulary_file is None:
  44         arg_parser.error('Please specify the vocabulary file with --vocabulary-file')
  45     elif 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                 logger.warning('{} label "{}" for {} ID "{}" already seen at {} ID "{}". Skipping.'.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                 select=['uuid', 'properties'],
  97                 filters=[['properties', 'exists', key_label]]
  98             )
  99             for o in objs:
 100                 props = copy.copy(o['properties'])
 101                 migrated_props = migrate_properties(props, key_label_to_id_map, vocab)
 102                 if not args.dry_run:
 103                     logger.debug('Migrating {}: {} -> {}'.format(o['uuid'], props, migrated_props))
 104                     arv.collections().update(uuid=o['uuid'], body={
 105                         'properties': migrated_props
 106                     }).execute()
 107                 else:
 108                     logger.info('Should migrate {}: {} -> {}'.format(o['uuid'], props, migrated_props))
 109                 migrated_counter += 1
 110                 if not args.dry_run and migrated_counter % 100 == 0:
 111                     logger.info('Migrating {} objects...'.format(migrated_counter))
 112
 113     if not args.dry_run:
 114         logger.info('Done, total object migrated: {}.'.format(migrated_counter))
 115     return 0
 116
 117 if __name__ == "__main__":
 118     sys.exit(main())