|
| 1 | +#!/usr/bin/python |
| 2 | + |
| 3 | +# |
| 4 | +# Licensed to the Apache Software Foundation (ASF) under one or more |
| 5 | +# contributor license agreements. See the NOTICE file distributed with |
| 6 | +# this work for additional information regarding copyright ownership. |
| 7 | +# The ASF licenses this file to You under the Apache License, Version 2.0 |
| 8 | +# (the "License"); you may not use this file except in compliance with |
| 9 | +# the License. You may obtain a copy of the License at |
| 10 | +# |
| 11 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 12 | +# |
| 13 | +# Unless required by applicable law or agreed to in writing, software |
| 14 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 15 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 16 | +# See the License for the specific language governing permissions and |
| 17 | +# limitations under the License. |
| 18 | +# |
| 19 | + |
| 20 | +from __future__ import absolute_import, division, print_function |
| 21 | +__metaclass__ = type |
| 22 | + |
| 23 | + |
| 24 | +DOCUMENTATION = ''' |
| 25 | +--- |
| 26 | +module: mongodb |
| 27 | +short_description: A module which support some simple operations on MongoDB. |
| 28 | +description: |
| 29 | + - Including add user/insert document/create indexes in MongoDB |
| 30 | +options: |
| 31 | + connect_string: |
| 32 | + description: |
| 33 | + - The uri of mongodb server |
| 34 | + required: true |
| 35 | + database: |
| 36 | + description: |
| 37 | + - The name of the database you want to manipulate |
| 38 | + required: true |
| 39 | + user: |
| 40 | + description: |
| 41 | + - The name of the user to add or remove, required when use 'user' mode |
| 42 | + required: false |
| 43 | + default: null |
| 44 | + password: |
| 45 | + description: |
| 46 | + - The password to use for the user, required when use 'user' mode |
| 47 | + required: false |
| 48 | + default: null |
| 49 | + roles: |
| 50 | + description: |
| 51 | + - The roles of the user, it's a list of dict, each dict requires two fields: 'db' and 'role', required when use 'user' mode |
| 52 | + required: false |
| 53 | + default: null |
| 54 | + collection: |
| 55 | + required: false |
| 56 | + description: |
| 57 | + - The name of the collection you want to manipulate, required when use 'doc' or 'indexes' mode |
| 58 | + doc: |
| 59 | + required: false |
| 60 | + description: |
| 61 | + - The document you want to insert into MongoDB, required when use 'doc' mode |
| 62 | + indexes: |
| 63 | + required: false |
| 64 | + description: |
| 65 | + - The indexes you want to create in MongoDB, it's a list of dict, you can see the example for the usage, required when use 'index' mode |
| 66 | + force_update: |
| 67 | + required: false |
| 68 | + description: |
| 69 | + - Whether replace/update existing user or doc or raise DuplicateKeyError, default is false |
| 70 | + mode: |
| 71 | + required: false |
| 72 | + default: user |
| 73 | + choices: ['user', 'doc', 'index'] |
| 74 | + description: |
| 75 | + - use 'user' mode if you want to add user, 'doc' mode to insert document, 'index' mode to create indexes |
| 76 | +
|
| 77 | +requirements: [ "pymongo" ] |
| 78 | +author: |
| 79 | + - "Jinag PengCheng" |
| 80 | +''' |
| 81 | + |
| 82 | +EXAMPLES = ''' |
| 83 | +# add user |
| 84 | +- mongodb: |
| 85 | + connect_string: mongodb://localhost:27017 |
| 86 | + database: admin |
| 87 | + user: test |
| 88 | + password: 123456 |
| 89 | + roles: |
| 90 | + - db: test_database |
| 91 | + role: read |
| 92 | + force_update: true |
| 93 | +
|
| 94 | +# add doc |
| 95 | +- mongodb: |
| 96 | + connect_string: mongodb://localhost:27017 |
| 97 | + mode: doc |
| 98 | + database: admin |
| 99 | + collection: main |
| 100 | + doc: |
| 101 | + id: "id/document" |
| 102 | + title: "the name of document" |
| 103 | + content: "which doesn't matter" |
| 104 | + force_update: true |
| 105 | +
|
| 106 | +# add indexes |
| 107 | +- mongodb: |
| 108 | + connect_string: mongodb://localhost:27017 |
| 109 | + mode: index |
| 110 | + database: admin |
| 111 | + collection: main |
| 112 | + indexes: |
| 113 | + - index: |
| 114 | + - field: updated_at |
| 115 | + direction: 1 |
| 116 | + - field: name |
| 117 | + direction: -1 |
| 118 | + name: test-index |
| 119 | + unique: true |
| 120 | +''' |
| 121 | + |
| 122 | +import traceback |
| 123 | + |
| 124 | +from ansible.module_utils.basic import AnsibleModule |
| 125 | +from ansible.module_utils._text import to_native |
| 126 | + |
| 127 | +try: |
| 128 | + from pymongo import ASCENDING, DESCENDING, GEO2D, GEOHAYSTACK, GEOSPHERE, HASHED, TEXT |
| 129 | + from pymongo import IndexModel |
| 130 | + from pymongo import MongoClient |
| 131 | + from pymongo.errors import DuplicateKeyError |
| 132 | +except ImportError: |
| 133 | + pass |
| 134 | + |
| 135 | + |
| 136 | +# ========================================= |
| 137 | +# MongoDB module specific support methods. |
| 138 | +# |
| 139 | + |
| 140 | +class UnknownIndexPlugin(Exception): |
| 141 | + pass |
| 142 | + |
| 143 | + |
| 144 | +def check_params(params, mode, module): |
| 145 | + missed_params = [] |
| 146 | + for key in OPERATIONS[mode]['required']: |
| 147 | + if params[key] is None: |
| 148 | + missed_params.append(key) |
| 149 | + |
| 150 | + if missed_params: |
| 151 | + module.fail_json(msg="missing required arguments: %s" % (",".join(missed_params))) |
| 152 | + |
| 153 | + |
| 154 | +def _recreate_user(module, db, user, password, roles): |
| 155 | + try: |
| 156 | + db.command("dropUser", user) |
| 157 | + db.command("createUser", user, pwd=password, roles=roles) |
| 158 | + except Exception as e: |
| 159 | + module.fail_json(msg='Unable to create user: %s' % to_native(e), exception=traceback.format_exc()) |
| 160 | + |
| 161 | + |
| 162 | + |
| 163 | +def user(module, client, db_name, **kwargs): |
| 164 | + roles = kwargs['roles'] |
| 165 | + if roles is None: |
| 166 | + roles = [] |
| 167 | + db = client[db_name] |
| 168 | + |
| 169 | + try: |
| 170 | + db.command("createUser", kwargs['user'], pwd=kwargs['password'], roles=roles) |
| 171 | + except DuplicateKeyError as e: |
| 172 | + if kwargs['force_update']: |
| 173 | + _recreate_user(module, db, kwargs['user'], kwargs['password'], roles) |
| 174 | + else: |
| 175 | + module.fail_json(msg='Unable to create user: %s' % to_native(e), exception=traceback.format_exc()) |
| 176 | + except Exception as e: |
| 177 | + module.fail_json(msg='Unable to create user: %s' % to_native(e), exception=traceback.format_exc()) |
| 178 | + |
| 179 | + module.exit_json(changed=True, user=kwargs['user']) |
| 180 | + |
| 181 | + |
| 182 | +def doc(module, client, db_name, **kwargs): |
| 183 | + coll = client[db_name][kwargs['collection']] |
| 184 | + try: |
| 185 | + coll.insert_one(kwargs['doc']) |
| 186 | + except DuplicateKeyError as e: |
| 187 | + if kwargs['force_update']: |
| 188 | + try: |
| 189 | + coll.replace_one({'_id': kwargs['doc']['_id']}, kwargs['doc']) |
| 190 | + except Exception as e: |
| 191 | + module.fail_json(msg='Unable to insert doc: %s' % to_native(e), exception=traceback.format_exc()) |
| 192 | + else: |
| 193 | + module.fail_json(msg='Unable to insert doc: %s' % to_native(e), exception=traceback.format_exc()) |
| 194 | + except Exception as e: |
| 195 | + module.fail_json(msg='Unable to insert doc: %s' % to_native(e), exception=traceback.format_exc()) |
| 196 | + |
| 197 | + kwargs['doc']['_id'] = str(kwargs['doc']['_id']) |
| 198 | + module.exit_json(changed=True, doc=kwargs['doc']) |
| 199 | + |
| 200 | + |
| 201 | +def _clean_index_direction(direction): |
| 202 | + if direction in ["1", "-1"]: |
| 203 | + direction = int(direction) |
| 204 | + |
| 205 | + if direction not in [ASCENDING, DESCENDING, GEO2D, GEOHAYSTACK, GEOSPHERE, HASHED, TEXT]: |
| 206 | + raise UnknownIndexPlugin("Unable to create indexes: Unknown index plugin: %s" % direction) |
| 207 | + return direction |
| 208 | + |
| 209 | + |
| 210 | +def _clean_index_options(options): |
| 211 | + res = {} |
| 212 | + supported_options = set(['name', 'unique', 'background', 'sparse', 'bucketSize', 'min', 'max', 'expireAfterSeconds']) |
| 213 | + for key in set(options.keys()).intersection(supported_options): |
| 214 | + res[key] = options[key] |
| 215 | + if key in ['min', 'max', 'bucketSize', 'expireAfterSeconds']: |
| 216 | + res[key] = int(res[key]) |
| 217 | + |
| 218 | + return res |
| 219 | + |
| 220 | + |
| 221 | +def parse_indexes(idx): |
| 222 | + keys = [(k['field'], _clean_index_direction(k['direction'])) for k in idx.pop('index')] |
| 223 | + options = _clean_index_options(idx) |
| 224 | + return IndexModel(keys, **options) |
| 225 | + |
| 226 | + |
| 227 | +def index(module, client, db_name, **kwargs): |
| 228 | + parsed_indexes = map(parse_indexes, kwargs['indexes']) |
| 229 | + try: |
| 230 | + coll = client[db_name][kwargs['collection']] |
| 231 | + coll.create_indexes(parsed_indexes) |
| 232 | + except Exception as e: |
| 233 | + module.fail_json(msg='Unable to create indexes: %s' % to_native(e), exception=traceback.format_exc()) |
| 234 | + |
| 235 | + module.exit_json(changed=True, indexes=kwargs['indexes']) |
| 236 | + |
| 237 | + |
| 238 | +OPERATIONS = { |
| 239 | + 'user': { 'function': user, 'params': ['user', 'password', 'roles', 'force_update'], 'required': ['user', 'password']}, |
| 240 | + 'doc': {'function': doc, 'params': ['doc', 'collection', 'force_update'], 'required': ['doc', 'collection']}, |
| 241 | + 'index': {'function': index, 'params': ['indexes', 'collection'], 'required': ['indexes', 'collection']} |
| 242 | +} |
| 243 | + |
| 244 | + |
| 245 | +# ========================================= |
| 246 | +# Module execution. |
| 247 | +# |
| 248 | + |
| 249 | +def main(): |
| 250 | + module = AnsibleModule( |
| 251 | + argument_spec=dict( |
| 252 | + connect_string=dict(required=True), |
| 253 | + database=dict(required=True, aliases=['db']), |
| 254 | + mode=dict(default='user', choices=['user', 'doc', 'index']), |
| 255 | + user=dict(default=None), |
| 256 | + password=dict(default=None, no_log=True), |
| 257 | + roles=dict(default=None, type='list'), |
| 258 | + collection=dict(default=None), |
| 259 | + doc=dict(default=None, type='dict'), |
| 260 | + force_update=dict(default=False, type='bool'), |
| 261 | + indexes=dict(default=None, type='list'), |
| 262 | + ) |
| 263 | + ) |
| 264 | + |
| 265 | + mode = module.params['mode'] |
| 266 | + |
| 267 | + db_name = module.params['database'] |
| 268 | + |
| 269 | + params = {key: module.params[key] for key in OPERATIONS[mode]['params']} |
| 270 | + check_params(params, mode, module) |
| 271 | + |
| 272 | + try: |
| 273 | + client = MongoClient(module.params['connect_string']) |
| 274 | + except NameError: |
| 275 | + module.fail_json(msg='the python pymongo module is required') |
| 276 | + except Exception as e: |
| 277 | + module.fail_json(msg='unable to connect to database: %s' % to_native(e), exception=traceback.format_exc()) |
| 278 | + |
| 279 | + OPERATIONS[mode]['function'](module, client, db_name, **params) |
| 280 | + |
| 281 | + |
| 282 | +if __name__ == '__main__': |
| 283 | + main() |
0 commit comments