utils/test/result_collection_api/update/templates/update_mongodb.py

   1 ##############################################################################
   2 # Copyright (c) 2016 ZTE Corporation
   3 # feng.xiaowei@zte.com.cn
   4 # All rights reserved. This program and the accompanying materials
   5 # are made available under the terms of the Apache License, Version 2.0
   6 # which accompanies this distribution, and is available at
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 ##############################################################################
   9 import argparse
  10
  11 from pymongo import MongoClient
  12
  13 from changes_in_mongodb import collections_old2New, fields_old2New, docs_old2New
  14 from utils import main, parse_mongodb_url
  15
  16 parser = argparse.ArgumentParser(description='Update MongoDBs')
  17
  18 parser.add_argument('-u', '--url',
  19                     type=str,
  20                     required=False,
  21                     default='mongodb://127.0.0.1:27017/',
  22                     help='Mongo DB URL for Backups')
  23
  24 parser.add_argument('-d', '--db',
  25                     type=str,
  26                     required=False,
  27                     default='test_results_collection',
  28                     help='database for the update.')
  29
  30
  31 def assert_collections(a_dict):
  32     if a_dict is not None:
  33         collections = eval_db('collection_names')
  34         no_collections = []
  35         for collection in a_dict.keys():
  36             if collection not in collections:
  37                 no_collections.append(collection)
  38         assert len(no_collections) == 0, \
  39             'collections {} not exist'.format(no_collections)
  40
  41
  42 def rename_collections(a_dict):
  43     if a_dict is not None:
  44         for collection, new_name in a_dict.iteritems():
  45             eval_collection(collection, 'rename', new_name)
  46
  47
  48 def rename_fields(a_dict):
  49     collection_update(a_dict, '$rename')
  50
  51
  52 def change_docs(a_dict):
  53     collection_update(a_dict, '$set')
  54
  55
  56 def eval_db(method, *args, **kwargs):
  57     return eval('db.%s(*args, **kwargs)' % method)
  58
  59
  60 def eval_collection(collection, method, *args, **kwargs):
  61     return eval('db.%s.%s(*args, **kwargs)' % (collection, method))
  62
  63
  64 def collection_update(a_dict, operator):
  65     if a_dict is not None:
  66         for collection, updates in a_dict.iteritems():
  67             for (query, doc) in updates:
  68                 doc_dict = {operator: doc}
  69                 eval_collection(collection, 'update', query,
  70                                 doc_dict, upsert=False, multi=True)
  71
  72
  73 def update(args):
  74     parse_mongodb_url(args.url)
  75     client = MongoClient(args.url)
  76     global db
  77     db = client[args.db]
  78     assert_collections(docs_old2New)
  79     assert_collections(fields_old2New)
  80     assert_collections(collections_old2New)
  81     change_docs(docs_old2New)
  82     rename_fields(fields_old2New)
  83     rename_collections(collections_old2New)
  84
  85 if __name__ == '__main__':
  86     main(update, parser)