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))