| #!/usr/bin/env python |
| |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| from collections import namedtuple |
| import glob |
| import sys |
| import os |
| import argparse |
| import traceback |
| import pydocumentdb.documents as documents |
| import pydocumentdb.errors as document_errors |
| import pydocumentdb.document_client as document_client |
| |
| try: |
| import argcomplete |
| except ImportError: |
| argcomplete = False |
| |
| CLI_DIR = os.path.dirname(os.path.realpath(sys.argv[0])) |
| # ROOT_DIR is the OpenWhisk repository root |
| ROOT_DIR = os.path.join(os.path.join(CLI_DIR, os.pardir), os.pardir) |
| |
| DbContext = namedtuple('DbContext', ['client', 'db', 'whisks', 'subjects', 'activations']) |
| verbose = False |
| |
| |
| def main(): |
| global verbose |
| exit_code = 0 |
| try: |
| args = parse_args() |
| verbose = args.verbose |
| client = init_client(args) |
| exit_code = { |
| 'init': init_cmd, |
| 'prune': prune_cmd, |
| 'drop': drop_cmd |
| }[args.cmd](args, client) |
| except Exception as e: |
| print('Exception: ', e) |
| traceback.print_exc() |
| exit_code = 1 |
| |
| sys.exit(exit_code) |
| |
| |
| def parse_args(): |
| parser = argparse.ArgumentParser(description='OpenWhisk CosmosDB bootstrap tool') |
| parser.add_argument('--endpoint', help='DB Endpoint url like https://example.documents.azure.com:443/', |
| required=True) |
| parser.add_argument('--key', help='DB access key', required=True) |
| parser.add_argument('-v', '--verbose', help='Verbose mode', action="store_true") |
| |
| subparsers = parser.add_subparsers(title='available commands', dest='cmd') |
| |
| propmenu = subparsers.add_parser('init', help='initialize database') |
| propmenu.add_argument('db', help='Database name under which the collections would be created') |
| propmenu.add_argument('--dir', help='Directory under which auth files are stored') |
| |
| propmenu = subparsers.add_parser('prune', help='remove stale databases created by test') |
| propmenu.add_argument('--prefix', help='Database name prefix which are matched for removal', default="travis-") |
| |
| propmenu = subparsers.add_parser('drop', help='drop database') |
| propmenu.add_argument('db', help='Database name to be removed') |
| |
| if argcomplete: |
| argcomplete.autocomplete(parser) |
| return parser.parse_args() |
| |
| |
| def init_cmd(args, client): |
| db = get_or_create_db(client, args.db) |
| |
| whisks = init_coll(client, db, "whisks") |
| subjects = init_coll(client, db, "subjects") |
| activations = init_coll(client, db, "activations") |
| |
| db_ctx = DbContext(client, db, whisks, subjects, activations) |
| init_auth(db_ctx) |
| return 0 |
| |
| |
| def prune_cmd(args, client): |
| # Remove database which are one day old |
| pass |
| |
| |
| def drop_cmd(args, client): |
| db = get_db(client, args.db) |
| if db is not None: |
| client.DeleteDatabase(db['_self']) |
| log("Removed database : %s" % args.db) |
| else: |
| log("Database %s not found" % args.db) |
| |
| |
| def init_auth(ctx): |
| for subject in find_default_subjects(): |
| link = create_link(ctx.db, ctx.subjects, subject['id']) |
| options = {'partitionKey': subject.get('id')} |
| try: |
| ctx.client.ReadDocument(link, options) |
| log('Subject already exists : ' + subject['id']) |
| except document_errors.HTTPFailure as e: |
| if e.status_code == 404: |
| ctx.client.CreateDocument(ctx.subjects['_self'], subject, options) |
| log('Created subject : ' + subject['id']) |
| else: |
| raise e |
| |
| |
| def create_link(db, coll, doc_id): |
| return 'dbs/' + db['id'] + '/colls/' + coll['id'] + '/docs/' + doc_id |
| |
| |
| def find_default_subjects(): |
| files_dir = os.path.join(ROOT_DIR, "ansible/files") |
| for name in glob.glob1(files_dir, "auth.*"): |
| auth_file = open(os.path.join(files_dir, name), 'r') |
| uuid, key = auth_file.read().strip().split(":") |
| subject = name[name.index('.') + 1:] |
| doc = { |
| 'id': subject, |
| 'subject': subject, |
| 'namespaces': [ |
| { |
| 'name': subject, |
| 'uuid': uuid, |
| 'key': key |
| } |
| ] |
| } |
| auth_file.close() |
| yield doc |
| |
| |
| def init_client(args): |
| return document_client.DocumentClient(args.endpoint, {'masterKey': args.key}) |
| |
| |
| def get_db(client, db_name): |
| query = client.QueryDatabases('SELECT * FROM root r WHERE r.id=\'' + db_name + '\'') |
| return next(iter(query), None) |
| |
| |
| def get_or_create_db(client, db_name): |
| db = get_db(client, db_name) |
| if db is None: |
| db = client.CreateDatabase({'id': db_name}) |
| log('Created database "%s"' % db_name) |
| return db |
| |
| |
| def init_coll(client, db, coll_name): |
| query = client.QueryCollections(db['_self'], 'SELECT * FROM root r WHERE r.id=\'' + coll_name + '\'') |
| it = iter(query) |
| coll = next(it, None) |
| if coll is None: |
| collection_definition = {'id': coll_name, |
| 'partitionKey': |
| { |
| 'paths': ['/id'], |
| 'kind': documents.PartitionKind.Hash |
| } |
| } |
| collection_options = {} # {'offerThroughput': 10100} |
| coll = client.CreateCollection(db['_self'], collection_definition, collection_options) |
| log('Created collection "%s"' % coll_name) |
| return coll |
| |
| |
| def log(msg): |
| if verbose: |
| print(msg) |
| |
| |
| if __name__ == '__main__': |
| main() |