Merge "armband: adding pod2, and disabling jobs for pod1"
[releng.git] / utils / test / result_collection_api / update / update.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 utils import main, parse_mongodb_url
14 from changes import collections_old2New, fields_old2New, docs_old2New
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)