Source code for kuha_document_store.db_setup

#!/usr/bin/env python3
# Author(s): Toni Sissala
# Copyright 2020 Finnish Social Science Data Archive FSD / University of Tampere
# Licensed under the EUPL. See LICENSE.txt for full license.
"""Script to help setup Document Store database.

Database administrator may use this script to setup MongoDB
instance for usage with Document Store.
"""

import os
import sys
import getpass
from pprint import pprint

import pymongo
from kuha_common import cli_setup
from kuha_document_store.configure import add_database_configs
from kuha_document_store import database


[docs]def setup_admin_user(admin_username, admin_password, db): """Setup administrator credentials. :note: authentication must be disabled in MongoDB to use this operation. :param admin_username: administrator username. :type admin_username: str :param admin_password: administrator password. :type admin_password: str :param db: MongoDB database :type db: :obj:`pymongo.database.Database` :returns: MongoDB database command response """ return db.command('createUser', admin_username, pwd=admin_password, roles=['root'])
[docs]def setup_users(settings, client): """Setup database users for Document Store. :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: list of MongoDB database command responses """ result = [] db = client[settings.database_name] result.append(db.command('createUser', settings.database_user_reader, pwd=settings.database_pass_reader, roles=['read'])) result.append(db.command('createUser', settings.database_user_editor, pwd=settings.database_pass_editor, roles=['readWrite'])) return result
[docs]def remove_users(settings, client): """Remove Document Store users from database. :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: list of MongoDB database command responses """ result = [] db = client[settings.database_name] result.append(db.command('dropUser', settings.database_user_reader)) result.append(db.command('dropUser', settings.database_user_editor)) return result
[docs]def setup_database(settings, client): """Create Document Store database. :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: PyMongo Database object """ return client.get_database(settings.database_name)
[docs]def delete_database(settings, client): """Delete Document Store database. :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: None """ return client.drop_database(settings.database_name)
[docs]def list_databases(settings, client): """List (print) databases. :note: Database won't show in list before it has a collection :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: list of database names """ return client.list_database_names()
[docs]def setup_collections(settings, client): """Setup Document Store collections (tables). :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: list of results """ results = {} db_name = settings.database_name for _coll in database.RECORD_COLLECTIONS: result = [] collection = client[db_name].create_collection( _coll.collection_name, validator=_coll.get_validator() ) for _index in _coll.indexes_unique: result.append(collection.create_index(_index, unique=True)) for _index in _coll.indexes: result.append(collection.create_index(_index)) results.update({_coll.collection_name: result}) return results
[docs]def delete_collections(settings, client): """Delete Document Store collections (tables). :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: list of drop_collection results """ db = client[settings.database_name] result = [] for _coll in database.RECORD_COLLECTIONS: result.append(db.drop_collection(_coll.collection_name)) return result
[docs]def list_collections(settings, client): """List Document Store collections (tables). :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: List of mongodb collections """ db_name = settings.database_name return client[db_name].list_collection_names()
[docs]def list_db_users(settings, client): """List (print) database users. :param settings: Document Store settings. :type settings: :obj:`argparse.Namespace` :param client: MongoDB client :type client: :obj:`pymongo.mongo_client.MongoClient` :returns: dictionary containing database users and their properties. """ db_name = settings.database_name return client[db_name].command('usersInfo')
#: Supported operations. OPERATIONS = { 'setup_users': setup_users, 'setup_admin_user': setup_admin_user, 'setup_database': setup_database, 'delete_database': delete_database, 'remove_users': remove_users, 'setup_collections': setup_collections, 'delete_collections': delete_collections, 'list_database_users': list_db_users, 'list_collections': list_collections, 'list_databases': list_databases } def _configure(): if cli_setup.settings.is_settings_loaded(): return cli_setup.get_settings() cli_setup.load(os.path.dirname(os.path.realpath(__file__)), description='Setup Kuha Document Store database.', config_file='kuha_document_store_db_setup.ini') cli_setup.add( 'operations', nargs='+', help='Operations for perform', choices=list(OPERATIONS) ) add_database_configs() cli_setup.add_print_configuration(cli_setup.settings.parser) return cli_setup.setup() def _client(settings, credentials=None): uri = database.mongodburi('%s:%s' % (settings.database_host, settings.database_port), database='admin', credentials=credentials) return pymongo.MongoClient(uri)
[docs]def main(): """Script main entry point. """ settings = _configure() if settings.print_configuration: print("Print active configuration and exit\n") for key, value in settings._get_kwargs(): print("{} ... {}".format(key, value)) return 0 print("Give database credentials") admin_username = input("Admin username:") admin_password = getpass.getpass("Admin password:") client = _client(settings, (admin_username, admin_password)) unauth_client = None try: for op_str in settings.operations: print("Running operation %s..." % op_str) op_call = OPERATIONS[op_str] if op_str == 'setup_admin_user': unauth_client = _client(settings) result = op_call(admin_username, admin_password, unauth_client.admin) else: result = op_call(settings, client) print("Done") if result is not None: pprint(result) finally: client.close() if unauth_client: unauth_client.close() return 0
if __name__ == '__main__': sys.exit(main())