1 files changed, 88 insertions, 0 deletions
diff --git a/scripts/scalability/test_controller/server/user_dbs.py b/scripts/scalability/test_controller/server/user_dbs.py
new file mode 100755
index 00000000..a1a9c222
--- /dev/null
+++ b/scripts/scalability/test_controller/server/user_dbs.py
@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+
+# Handle creation of user databases for scalability tests.
+
+import argparse
+import treq
+
+from functools import partial
+from urlparse import urljoin
+
+from twisted.internet import reactor, defer
+from twisted.logger import Logger
+
+COUCH_URL = "http://127.0.0.1:5984"
+CREATE = 1000
+
+
+logger = Logger()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--couch-url', default=COUCH_URL,
+                        help='The URL to the CouchDB server.')
+    parser.add_argument('--create', default=CREATE,
+                        help='The number of databases to create.')
+    return parser.parse_args()
+
+
+def get_db_names(create):
+    dbs = []
+    for i in xrange(create):
+        dbname = 'user-%d' % i
+        dbs.append(dbname)
+    return dbs
+
+
+semaphore = defer.DeferredSemaphore(20)
+
+
+def _log(db, action, res):
+    logger.info('table %s %s' % (db, action))
+    return res
+
+
+@defer.inlineCallbacks
+def delete_dbs(dbs):
+    deferreds = []
+    for db in dbs:
+        d = semaphore.run(treq.delete, urljoin(COUCH_URL, db))
+        logfun = partial(_log, db, 'deleted')
+        d.addCallback(logfun)
+        deferreds.append(d)
+    responses = yield defer.gatherResults(deferreds)
+    codes = map(lambda r: r.code, responses)
+    assert all(map(lambda c: c == 200 or c == 404, codes))
+
+
+@defer.inlineCallbacks
+def create_dbs(dbs):
+    deferreds = []
+    for db in dbs:
+        d = semaphore.run(treq.put, urljoin(COUCH_URL, db))
+        logfun = partial(_log, db, 'created')
+        d.addCallback(logfun)
+        deferreds.append(d)
+    responses = yield defer.gatherResults(deferreds)
+    codes = map(lambda r: r.code, responses)
+    assert all(map(lambda c: c == 201, codes))
+
+
+@defer.inlineCallbacks
+def ensure_dbs(couch_url=COUCH_URL, create=CREATE):
+    dbs = get_db_names(create)
+    yield delete_dbs(dbs)
+    yield create_dbs(dbs)
+
+
+@defer.inlineCallbacks
+def main(couch_url, create):
+    yield ensure_dbs(couch_url, create)
+    reactor.stop()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    d = main(args.couch_url, args.create)
+    reactor.run()