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, \
  14     fields_old2New, docs_old2New
  15 from utils import main, parse_mongodb_url
  16
  17 parser = argparse.ArgumentParser(description='Update MongoDBs')
  18
  19 parser.add_argument('-u', '--url',
  20                     type=str,
  21                     required=False,
  22                     default='mongodb://127.0.0.1:27017/',
  23                     help='Mongo DB URL for Backups')
  24
  25 parser.add_argument('-d', '--db',
  26                     type=str,
  27                     required=False,
  28                     default='test_results_collection',
  29                     help='database for the update.')
  30
  31
  32 def assert_collections(a_dict):
  33     if a_dict is not None:
  34         collections = eval_db('collection_names')
  35         no_collections = []
  36         for collection in a_dict.keys():
  37             if collection not in collections:
  38                 no_collections.append(collection)
  39         assert len(no_collections) == 0, \
  40             'collections {} not exist'.format(no_collections)
  41
  42
  43 def rename_collections(a_dict):
  44     if a_dict is not None:
  45         for collection, new_name in a_dict.iteritems():
  46             eval_collection(collection, 'rename', new_name)
  47
  48
  49 def rename_fields(a_dict):
  50     collection_update(a_dict, '$rename')
  51
  52
  53 def change_docs(a_dict):
  54     collection_update(a_dict, '$set')
  55
  56
  57 def eval_db(method, *args, **kwargs):
  58     exec_db = db.__getattribute__(method)
  59     return exec_db(*args, **kwargs)
  60
  61
  62 def eval_collection(collection, method, *args, **kwargs):
  63     exec_collection = db.__getattr__(collection)
  64     return exec_collection.__getattribute__(method)(*args, **kwargs)
  65
  66
  67 def collection_update(a_dict, operator):
  68     if a_dict is not None:
  69         for collection, updates in a_dict.iteritems():
  70             for (query, doc) in updates:
  71                 doc_dict = {operator: doc}
  72                 eval_collection(collection, 'update', query,
  73                                 doc_dict, upsert=False, multi=True)
  74
  75
  76 def update(args):
  77     parse_mongodb_url(args.url)
  78     client = MongoClient(args.url)
  79     global db
  80     db = client[args.db]
  81     assert_collections(docs_old2New)
  82     assert_collections(fields_old2New)
  83     assert_collections(collections_old2New)
  84     change_docs(docs_old2New)
  85     rename_fields(fields_old2New)
  86     rename_collections(collections_old2New)
  87
  88 if __name__ == '__main__':
  89     main(update, parser)