def db_drop(res_id):
DB = get_db()
collections = get_collection_names(res_id)
with UseResId(res_id):
for c in collections:
DB.drop_collection(c)
return empty_success()
def test_updates_collection_list(self):
with self.real_app.app_context():
db = get_db()
res_id = 'myresid.'
# Setup resource id record
clients_collection = db[CLIENTS_COLLECTION]
clients_collection.remove({'res_id': res_id})
clients_collection.insert({'res_id': res_id, 'collections': []})
with UseResId(res_id) as db:
self.assertItemsEqual(get_collection_names(res_id), [])
db.foo.insert({'message': 'test'})
self.assertItemsEqual(get_collection_names(res_id), ['foo'])
self.assertItemsEqual(list(db.foo.find({}, {'_id': 0})),
[{
'message': 'test'
}])
db.bar.update({}, {'message': 'test'}, upsert=True)
self.assertItemsEqual(get_collection_names(res_id),
['foo', 'bar'])
self.assertItemsEqual(list(db.bar.find({}, {'_id': 0})),
[{
'message': 'test'
}])
db.foo.drop()
self.assertItemsEqual(get_collection_names(res_id), ['bar'])
self.assertNotIn(res_id + 'foo', get_collection_names(res_id))
def load_data_from_mongoexport(res_id,
export_location,
collection_name,
remove_id=False):
"""
This file should come from mongoexport, with or without the --jsonArray
flag. That is to say, it should either be a series of documents, each on
its own line, or a single array of documents. All documents will be
inserted into the given collection.
"""
export_location = _data_file_path(export_location)
with open(export_location) as export:
first_char = export.read(1)
export.seek(0, SEEK_SET)
if first_char == '[':
# Data is already in an array
documents = loads(export.read())
else:
# Each line of data is an object
documents = []
for line in export:
documents.append(loads(line))
if remove_id:
_remove_id(documents)
with UseResId(res_id) as db:
db[collection_name].insert(documents)
def db_collection_insert(res_id, collection_name):
parse_get_json()
document = request.json.get('document')
if document is None:
raise MWSServerError(400,
"no object passed to insert!")
validate_document_or_list(document)
req_size = calculate_document_size(document)
# Insert document
with UseResId(res_id) as db:
# Check quota
size = db[collection_name].size()
if size + req_size > current_app.config['QUOTA_COLLECTION_SIZE']:
raise MWSServerError(403, 'Collection size exceeded')
# Attempt Insert
try:
res = db[collection_name].insert(document)
except (DuplicateKeyError, OperationFailure) as e:
raise MWSServerError(400, str(e))
if isinstance(res, list):
pretty_response = pretty_bulk_insert.format(len(res))
else:
pretty_response = pretty_insert.format(1)
return to_json({'pretty': pretty_response})
def db_cursor_next(res_id, collection_name):
parse_get_json()
result = {}
batch_size = current_app.config['CURSOR_BATCH_SIZE']
with UseResId(res_id, db=get_keepalive_db()) as db:
coll = db[collection_name]
cursor_id = int(request.json.get('cursor_id'))
retrieved = request.json.get('retrieved', 0)
drain_cursor = request.json.get('drain_cursor', False)
batch_size = -1 if drain_cursor else current_app.config['CURSOR_BATCH_SIZE']
cursor = recreate_cursor(coll, cursor_id, retrieved, batch_size)
try:
result['result'] = []
for i in range(batch_size):
try:
result['result'].append(cursor.next())
except StopIteration:
result['empty_cursor'] = True
break
except OperationFailure as e:
return MWSServerError(400, 'Cursor not found')
# kill cursor on server if all results are returned
if result.get('empty_cursor'):
kill_cursor(coll, long(cursor_id))
return to_json(result)
def db_collection_insert(res_id, collection_name):
# TODO: Ensure request.json is not None.
if 'document' in request.json:
document = request.json['document']
else:
error = '\'document\' argument not found in the insert request.'
raise MWSServerError(400, error)
# Check quota
size = get_collection_size(res_id, collection_name)
# Handle inserting both a list of docs or a single doc
if isinstance(document, list):
req_size = 0
for d in document:
req_size += len(BSON.encode(d))
else:
req_size = len(BSON.encode(document))
if size + req_size > current_app.config['QUOTA_COLLECTION_SIZE']:
raise MWSServerError(403, 'Collection size exceeded')
# Insert document
with UseResId(res_id):
try:
get_db()[collection_name].insert(document)
return empty_success()
except InvalidDocument as e:
raise MWSServerError(400, e.message)
def db_count(res_id, collection_name):
parse_get_json()
with UseResId(res_id) as db:
query = request.json.get('query')
coll = db[collection_name]
count = coll.find(query).count()
return to_json(count)
def db_collection_save(res_id, collection_name):
parse_get_json()
document = request.json.get('document')
if document is None:
raise MWSServerError(400,
"'document' argument not found in the save request.")
validate_document(document)
req_size = calculate_document_size(document)
# Get database
with UseResId(res_id) as db:
# Check quota
size = db[collection_name].size()
if size + req_size > current_app.config['QUOTA_COLLECTION_SIZE']:
raise MWSServerError(403, 'Collection size exceeded')
# Save document
try:
db[collection_name].save(document)
return empty_success()
except (InvalidId, TypeError, InvalidDocument, DuplicateKeyError) as e:
raise MWSServerError(400, str(e))
def db_collection_update(res_id, collection_name):
query = update = None
if request.json:
query = request.json.get('query')
update = request.json.get('update')
upsert = request.json.get('upsert', False)
multi = request.json.get('multi', False)
if query is None or update is None:
error = 'update requires spec and document arguments'
raise MWSServerError(400, error)
# Check quota
size = get_collection_size(res_id, collection_name)
with UseResId(res_id):
# Computation of worst case size increase - update size * docs affected
# It would be nice if we were able to make a more conservative estimate
# of the space difference that an update will cause. (especially if it
# results in smaller documents)
db = get_db()
affected = db[collection_name].find(query).count()
req_size = len(BSON.encode(update)) * affected
if size + req_size > current_app.config['QUOTA_COLLECTION_SIZE']:
raise MWSServerError(403, 'Collection size exceeded')
try:
db[collection_name].update(query, update, upsert, multi=multi)
return empty_success()
except OperationFailure as e:
raise MWSServerError(400, e.message)
def db_find_one(res_id, collection_name):
parse_get_json()
with UseResId(res_id) as db:
query = request.json.get('query')
projection = request.json.get('projection')
coll = db[collection_name]
doc = coll.find_one(query, projection)
return to_json(doc)
def __precompute(self, collection, data_only, data, check_id):
with UseResId(self.res_id):
query = {'$or': data} if data_only else {}
projection = None if check_id else {'_id': 0}
result = self.db[collection].find(query, projection)
data = (self.__hashable(x) for x in data)
result = (self.__hashable(x) for x in result)
return data, result
def test_mangles_collection_names_automatically(self):
with self.real_app.app_context():
db = get_db()
with UseResId('myresid.'):
coll = db.foo
self.assertEqual(coll.name, 'myresid.foo')
coll = db.foo
self.assertEqual(coll.name, 'foo')
def setUp(self):
super(QuotaCollectionsTestCase, self).setUp()
self.old_quota = self.real_app.config['QUOTA_NUM_COLLECTIONS']
self.res_id = 'myresid.'
with self.real_app.app_context():
collections = get_collection_names(self.res_id)
with UseResId(self.res_id) as db:
for c in collections:
db.drop_collection(c)
def db_collection_aggregate(res_id, collection_name):
parse_get_json(request)
try:
with UseResId(res_id):
coll = get_db()[collection_name]
result = coll.aggregate(request.json)
return to_json(result)
except OperationFailure as e:
raise MWSServerError(400, e.message)
def test_quota_collections_zero(self):
self.real_app.config['QUOTA_NUM_COLLECTIONS'] = 0
with self.real_app.app_context(), UseResId(self.res_id):
with self.assertRaises(MWSServerError) as cm:
self.db.a.insert({'a': 1})
self.assertEqual(cm.exception.error, 429)
self.db.drop_collection('a')
def db_collection_remove(res_id, collection_name):
constraint = request.json.get('constraint') if request.json else {}
just_one = request.json and request.json.get('just_one', False)
with UseResId(res_id):
collection = get_db()[collection_name]
if just_one:
collection.find_and_modify(constraint, remove=True)
else:
collection.remove(constraint)
return empty_success()
def db_collection_aggregate(res_id, collection_name):
parse_get_json()
with UseResId(res_id) as db:
try:
result = db[collection_name].aggregate(request.json)
except (InvalidId,
TypeError,
InvalidDocument,
OperationFailure) as e:
raise MWSServerError(400, str(e))
return to_json(result)
def db_collection_count(res_id, collection_name):
parse_get_json()
query = request.json.get('query')
with UseResId(res_id) as db:
coll = db[collection_name]
try:
count = coll.find(query).count()
return to_json({'count': count})
except InvalidDocument as e:
raise MWSServerError(400, str(e))
def db_collection_count(res_id, collection_name):
parse_get_json(request)
query = request.json.get('query')
skip = request.json.get('skip', 0)
limit = request.json.get('limit', 0)
use_skip_limit = bool(skip or limit)
with UseResId(res_id):
coll = get_db()[collection_name]
cursor = coll.find(query, skip=skip, limit=limit)
count = cursor.count(use_skip_limit)
return to_json({'count': count})
def load_data_from_json(res_id, file_name, remove_id=False):
"""
The top level of this file should be an object who's keys are collection
names which map to an array of documents to be inserted into the collection
"""
file_name = _data_file_path(file_name)
with open(file_name) as json_file:
collections = loads(json_file.read())
with UseResId(res_id) as db:
for collection, documents in collections.iteritems():
if remove_id:
_remove_id(documents)
db[collection].insert(documents)
def test_quota_collections(self):
self.real_app.config['QUOTA_NUM_COLLECTIONS'] = 2
with self.real_app.app_context():
with UseResId(self.res_id) as db:
db.a.insert({'a': 1})
db.b.insert({'b': 1})
with self.assertRaises(MWSServerError) as cm:
db.c.insert({'c': 1})
self.assertEqual(cm.exception.error, 429)
for c in ['a', 'b']:
db.drop_collection(c)
def db_collection_remove(res_id, collection_name):
parse_get_json()
constraint = request.json.get('constraint') if request.json else {}
options = request.json and request.json.get('options', False)
multi = not options.get('justOne')
with UseResId(res_id) as db:
collection = db[collection_name]
try:
res = collection.remove(constraint, multi=multi)
pretty_response = pretty_remove.format(res.get('n'))
except (InvalidDocument, InvalidId, TypeError, OperationFailure) as e:
raise MWSServerError(400, str(e))
return to_json({'pretty': pretty_response})
def db_collection_remove(res_id, collection_name):
parse_get_json()
constraint = request.json.get('constraint') if request.json else {}
just_one = request.json and request.json.get('just_one', False)
with UseResId(res_id) as db:
collection = db[collection_name]
try:
if just_one:
collection.find_and_modify(constraint, remove=True)
else:
collection.remove(constraint)
except (InvalidDocument, InvalidId, TypeError, OperationFailure) as e:
raise MWSServerError(400, str(e))
return empty_success()
def load_data_from_mongodump(res_id, dump_location, collection_name):
"""
The dump location should point to a .bson file, not a directory structure
as created by mongodump. Instead, use the .bson files inside this
directory structure.
"""
dump_location = _data_file_path(dump_location)
if not os.path.exists(dump_location):
raise NotFound('Unable to find dump file')
p = Popen(('mongorestore', '-d', 'mws', '-c',
'%s%s' % (res_id, collection_name), dump_location))
p.communicate() # Wait for process to finish
if p.poll() != 0:
raise InternalServerError('Loading dumped data failed')
UseResId(res_id).insert_client_collection(collection_name)
def db_collection_update(res_id, collection_name):
parse_get_json()
query = request.json.get('query')
update = request.json.get('update')
upsert = request.json.get('upsert', False)
multi = request.json.get('multi', False)
if query is None or update is None:
error = 'update requires spec and document arguments'
raise MWSServerError(400, error)
with UseResId(res_id) as db:
# Check quota
coll = db[collection_name]
# Computation of worst case size increase - update size * docs affected
# It would be nice if we were able to make a more conservative estimate
# of the space difference that an update will cause. (especially if it
# results in smaller documents)
# TODO: Make this more intelligent. I'm not sure that this even makes sense.
affected = coll.find(query).count()
req_size = calculate_document_size(update) * affected
size = db[collection_name].size()
if size + req_size > current_app.config['QUOTA_COLLECTION_SIZE']:
raise MWSServerError(403, 'Collection size exceeded')
# Attempt Update
try:
res = db[collection_name].update(query, update, upsert, multi=multi)
_logger.info("res: {0}".format(res))
n_matched = 0 if res.get('upserted') else res.get('n')
n_upserted = 1 if res.get('upserted') else 0
n_modified = res.get('nModified', 0)
if n_upserted:
_id = res.get('upserted')[0].get('_id')
pretty_response = pretty_upsert.format(n_matched, n_upserted, n_modified, _id)
else:
pretty_response = pretty_update.format(n_matched, n_modified)
return to_json({'pretty': pretty_response})
except (DuplicateKeyError,
InvalidDocument,
InvalidId,
TypeError,
OperationFailure) as e:
raise MWSServerError(400, str(e))
def db_collection_find(res_id, collection_name):
parse_get_json()
result = {}
batch_size = current_app.config['CURSOR_BATCH_SIZE']
with UseResId(res_id, db=get_keepalive_db()) as db:
limit = request.json.get('limit', 0)
coll = db[collection_name]
query = request.json.get('query')
projection = request.json.get('projection')
skip = request.json.get('skip', 0)
sort = request.json.get('sort', {})
sort = sort.items()
cursor = coll.find(spec=query, fields=projection, skip=skip,
limit=limit)
cursor.batch_size(batch_size)
if len(sort) > 0:
cursor.sort(sort)
# count is only available before cursor is read so we include it
# in the first response
result['count'] = cursor.count(with_limit_and_skip=True)
count = result['count']
num_to_return = min(limit, batch_size) if limit else batch_size
try:
result['result'] = []
for i in range(num_to_return):
try:
result['result'].append(cursor.next())
except StopIteration:
break
except OperationFailure as e:
return MWSServerError(400, 'Cursor not found')
# cursor_id is too big as a number, use a string instead
result['cursor_id'] = str(cursor.cursor_id)
# close the Cursor object, but keep the cursor alive on the server
del cursor
return to_json(result)
def db_collection_find(res_id, collection_name):
parse_get_json()
query = request.json.get('query')
projection = request.json.get('projection')
skip = request.json.get('skip', 0)
limit = request.json.get('limit', 0)
sort = request.json.get('sort', {})
sort = sort.items()
with UseResId(res_id) as db:
coll = db[collection_name]
try:
cursor = coll.find(query, projection, skip, limit)
if len(sort) > 0:
cursor.sort(sort)
documents = list(cursor)
except (InvalidId, TypeError, OperationFailure) as e:
raise MWSServerError(400, str(e))
return to_json({'result': documents})
def db_collection_find(res_id, collection_name):
# TODO: Should we specify a content type? Then we have to use an options
# header, and we should probably get the return type from the content-type
# header.
parse_get_json(request)
query = request.json.get('query')
projection = request.json.get('projection')
skip = request.json.get('skip', 0)
limit = request.json.get('limit', 0)
sort = request.json.get('sort', {})
sort = sort.items()
with UseResId(res_id):
coll = get_db()[collection_name]
cursor = coll.find(query, projection, skip, limit)
if len(sort) > 0:
cursor.sort(sort)
documents = list(cursor)
return to_json({'result': documents})
def test_updates_collection_list(self):
with self.real_app.app_context():
db = get_db()
res_id = 'myresid.'
# Setup resource id record
clients_collection = db[CLIENTS_COLLECTION]
clients_collection.remove({'res_id': res_id})
clients_collection.insert({'res_id': res_id, 'collections': []})
def get_collections():
# Can't use the util function because we would be using it
# inside the with, so the collection name would be mangled
return clients_collection.find({'res_id': res_id}, {
'_id': 0,
'collections': 1
})[0]['collections']
with UseResId(res_id):
self.assertItemsEqual(get_collections(), [])
db.foo.insert({'message': 'test'})
self.assertItemsEqual(get_collections(), ['foo'])
self.assertItemsEqual(list(db.foo.find({}, {'_id': 0})),
[{
'message': 'test'
}])
db.bar.update({}, {'message': 'test'})
self.assertItemsEqual(get_collections(), ['foo'])
db.bar.update({}, {'message': 'test'}, upsert=True)
self.assertItemsEqual(get_collections(), ['foo', 'bar'])
self.assertItemsEqual(list(db.bar.find({}, {'_id': 0})),
[{
'message': 'test'
}])
db.foo.drop()
self.assertItemsEqual(get_collections(), ['bar'])
self.assertNotIn(res_id + 'foo', db.collection_names())
db.drop_collection('bar')
self.assertItemsEqual(get_collections(), [])
self.assertNotIn(res_id + 'bar', db.collection_names())
def db_collection_save(res_id, collection_name):
parse_get_json()
document = request.json.get('document')
if document is None:
raise MWSServerError(400,
"'document' argument not found in the save request.")
validate_document(document)
req_size = calculate_document_size(document)
# Get database
with UseResId(res_id) as db:
# Check quota
size = db[collection_name].size()
if size + req_size > current_app.config['QUOTA_COLLECTION_SIZE']:
raise MWSServerError(403, 'Collection size exceeded')
# Save document
try:
if "_id" not in document:
res = db[collection_name].insert(document)
if res:
res_len = len(res) if isinstance(res, list) else 1
pretty_response = pretty_insert.format(res_len)
else:
res = db[collection_name].update({"_id": document["_id"]},
document, True)
n_matched = 0 if res.get('upserted') else 1
n_upserted = 1 if res.get('upserted') else 0
n_modified = res.get('nModified', 0)
if n_upserted:
_id = res.get('upserted')[0].get('_id')
pretty_response = pretty_upsert.format(n_matched, n_upserted, n_modified, _id)
else:
pretty_response = pretty_update.format(n_matched, n_modified)
return to_json({'pretty': pretty_response})
except (InvalidId, TypeError, InvalidDocument, DuplicateKeyError) as e:
raise MWSServerError(400, str(e))
