diff options
Diffstat (limited to 'scripts/db_access')
| -rw-r--r-- | scripts/db_access/reset_db.py | 132 |
1 files changed, 95 insertions, 37 deletions
diff --git a/scripts/db_access/reset_db.py b/scripts/db_access/reset_db.py index 80871856..7c6d281b 100644 --- a/scripts/db_access/reset_db.py +++ b/scripts/db_access/reset_db.py @@ -5,20 +5,21 @@ # WARNING: running this script over a database will delete all documents but # the one with id u1db_config (which contains db metadata) and design docs # needed for couch backend. +# +# Run it like this to get some help: +# +# ./reset_db.py --help -import sys -from ConfigParser import ConfigParser import threading import logging -from couchdb import Database as CouchDatabase - +import argparse +import re -if len(sys.argv) != 2: - print 'Usage: %s <uuid>' % sys.argv[0] - exit(1) -uuid = sys.argv[1] +from ConfigParser import ConfigParser +from couchdb import Database as CouchDatabase +from couchdb import Server as CouchServer # create a logger @@ -27,23 +28,6 @@ LOG_FORMAT = '%(asctime)s %(message)s' logging.basicConfig(format=LOG_FORMAT, level=logging.INFO) -# get couch url -cp = ConfigParser() -cp.read('/etc/leap/soledad-server.conf') -url = cp.get('soledad-server', 'couch_url') - - -# confirm -yes = raw_input("Are you sure you want to reset the database for user %s " - "(type YES)? " % uuid) -if yes != 'YES': - print 'Bailing out...' - exit(2) - - -db = CouchDatabase('%s/user-%s' % (url, uuid)) - - class _DeleterThread(threading.Thread): def __init__(self, db, doc_id, release_fun): @@ -59,21 +43,95 @@ class _DeleterThread(threading.Thread): self._release_fun() -semaphore_pool = threading.BoundedSemaphore(value=20) - - -threads = [] -for doc_id in db: - if doc_id != 'u1db_config' and not doc_id.startswith('_design'): +def get_confirmation(noconfirm, uuid, shared): + msg = "Are you sure you want to reset %s (type YES)? " + if shared: + msg = msg % "the shared database" + elif uuid: + msg = msg % ("the database for user %s" % uuid) + else: + msg = msg % "all databases" + if noconfirm is False: + yes = raw_input(msg) + if yes != 'YES': + print 'Bailing out...' + exit(2) + + +def get_url(empty): + url = None + if empty is False: + # get couch url + cp = ConfigParser() + cp.read('/etc/leap/soledad-server.conf') + url = cp.get('soledad-server', 'couch_url') + else: + with open('/etc/couchdb/couchdb.netrc') as f: + netrc = f.read() + admin_password = re.match('^.* password (.*)$', netrc).groups()[0] + url = 'http://admin:%s@127.0.0.1:5984' % admin_password + return url + + +def reset_all_dbs(url, empty): + server = CouchServer('%s' % (url)) + for dbname in server: + if dbname.startswith('user-') or dbname == 'shared': + reset_db(url, dbname, empty) + + +def reset_db(url, dbname, empty): + db = CouchDatabase('%s/%s' % (url, dbname)) + semaphore_pool = threading.BoundedSemaphore(value=20) + + # launch threads for deleting docs + threads = [] + for doc_id in db: + if empty is False: + if doc_id == 'u1db_config' or doc_id.startswith('_design'): + continue semaphore_pool.acquire() logger.info('[main] launching thread for doc: %s' % doc_id) t = _DeleterThread(db, doc_id, semaphore_pool.release) t.start() threads.append(t) - -logger.info('[main] waiting for threads.') -map(lambda thread: thread.join(), threads) - - -logger.info('[main] done.') + # wait for threads to finish + logger.info('[main] waiting for threads.') + map(lambda thread: thread.join(), threads) + logger.info('[main] done.') + + +def _parse_args(): + parser = argparse.ArgumentParser() + group = parser.add_mutually_exclusive_group() + group.add_argument('-u', dest='uuid', default=False, + help='Reset database of given user.') + group.add_argument('-s', dest='shared', action='store_true', default=False, + help='Reset the shared database.') + group.add_argument('-a', dest='all', action='store_true', default=False, + help='Reset all user databases.') + parser.add_argument( + '-e', dest='empty', action='store_true', required=False, default=False, + help='Empty database (do not preserve minimal set of u1db documents).') + parser.add_argument( + '-y', dest='noconfirm', action='store_true', required=False, + default=False, + help='Do not ask for confirmation.') + return parser.parse_args(), parser + + +if __name__ == '__main__': + args, parser = _parse_args() + if not (args.uuid or args.shared or args.all): + parser.print_help() + exit(1) + + url = get_url(args.empty) + get_confirmation(args.noconfirm, args.uuid, args.shared) + if args.uuid: + reset_db(url, "user-%s" % args.uuid, args.empty) + elif args.shared: + reset_db(url, "shared", args.empty) + elif args.all: + reset_all_dbs(url, args.empty) |