def test_getStatus(self, mock_isfile):
worker_id = 'any worker id'
mock_isfile.return_value = True
result = WorkerResponse.get_status(worker_id)
self.assertEqual(result, {
'id': worker_id,
'status': 'finished',
'size': mock.ANY
})
mock_isfile.side_effect = [False, True, False]
result = WorkerResponse.get_status(worker_id)
self.assertEqual(result, {
'id': worker_id,
'status': 'working',
'size': mock.ANY
})
mock_isfile.side_effect = [False, False, True]
result = WorkerResponse.get_status(worker_id)
self.assertEqual(result, {
'id': worker_id,
'status': 'aborting',
'size': mock.ANY
})
mock_isfile.side_effect = [False, False, False]
result = WorkerResponse.get_status(worker_id)
self.assertIsNone(result)
def test_isAborting(self, mock_isfile):
mock_isfile.return_value = True
result = WorkerResponse.is_aborting('any id')
self.assertTrue(result)
mock_isfile.return_value = False
result = WorkerResponse.is_aborting('any id')
self.assertFalse(result)
def test_isFinished(self, mock_isfile):
mock_isfile.return_value = True
result = WorkerResponse.is_finished('any id')
self.assertTrue(result)
mock_isfile.return_value = False
result = WorkerResponse.is_finished('any id')
self.assertFalse(result)
def test_isWorking(self, mock_isfile):
mock_isfile.side_effect = [True, False]
result = WorkerResponse.is_working('any id')
self.assertTrue(result)
# working and is aborting
mock_isfile.side_effect = [True, True]
result = WorkerResponse.is_working('any id')
self.assertFalse(result)
def test_kill(self, mock_Path, mock_isfile):
worker_id = 'any worker id'
mock_isfile.return_value = True
WorkerResponse.kill(worker_id)
mock_Path.return_value.touch.assert_not_called()
mock_isfile.side_effect = [False, True, False]
WorkerResponse.kill(worker_id)
mock_Path.return_value.touch.assert_called()
def test_yield_progress(self):
worker = WorkerResponse()
self.assertIsNone(worker._last_progress)
result = [
progress for progress in worker.yield_progress('any filename')
]
self.assertEqual(result, ["0\n"])
self.assertIsNotNone(worker._last_progress)
result = [
progress for progress in worker.yield_progress('any filename')
]
self.assertEqual(result, [])
def test_stream_with_context(self, mock_stream_with_context, mock_response,
mock_request):
mock_request.headers = {}
result = WorkerResponse.stream_with_context([], 'any mimetype')
self.assertEqual(result, mock_response.return_value)
mock_response.assert_called_with(mock_stream_with_context.return_value,
mimetype='any mimetype')
mock_request.headers = {WorkerResponse._WORKER_REQUEST: True}
result = WorkerResponse.stream_with_context([], 'any mimetype')
self.assertEqual(result, mock_response.return_value)
mock_response.assert_called_with(mock_stream_with_context.return_value,
mimetype='text/plain')
def worker_end(worker_id):
"""
End a running worker
:param worker_id:
:return:
"""
status = WorkerResponse.get_status(worker_id)
if status:
WorkerResponse.kill(worker_id)
return "", 204 # No Content
else:
return _worker_not_found(worker_id)
def entrypoint(self):
# Compatible with plain GraphQL endpoint
query = request.args.get('query')
if not query:
# Compatible with existing GOB export code
request_data = json.loads(request.data.decode('utf-8'))
query = request_data['query']
graphql2sql = GraphQL2SQL(query)
try:
sql = graphql2sql.sql()
except NoAccessException as e:
return "Forbidden", 403
session = get_session()
# use an ad-hoc Connection and stream results (instead of pre-buffered)
result_rows = session.connection().execution_options(
stream_results=True).execute(text(sql))
response_builder = \
GraphQLCustomStreamingResponseBuilder(result_rows,
graphql2sql.relations_hierarchy,
graphql2sql.selections,
request_args=request.args)
return WorkerResponse.stream_with_context(
response_builder, mimetype='application/x-ndjson')
def worker_result(worker_id):
"""
Get the result of a worker.
If the worker has finished the worker file is returned.
:param worker_id:
:return:
"""
filename = WorkerResponse.get_response_file(worker_id)
if filename:
return send_file(filename)
elif WorkerResponse.is_working(worker_id):
return f"Worker {worker_id} not finished", 204 # No Content
else:
return _worker_not_found(worker_id)
def entrypoint(self):
query = self._get_query()
if not query:
logger.warning("No query passed in")
return jsonify({'error': str("no query given")}), 400
flat_query = query.replace("\n", "")
logger.info(f"Running GraphQL for {get_request_id()}: {flat_query}")
graphql2sql = GraphQL2SQL(query)
try:
sql = graphql2sql.sql()
except NoAccessException:
return jsonify({'error': 'Forbidden'}), 403
except InvalidQueryException as e:
return jsonify({'error': str(e)}), 400
session = get_session()
# use an ad-hoc Connection and stream results (instead of pre-buffered)
result_rows = session.connection().execution_options(
stream_results=True).execute(text(sql))
response_builder = GraphQLCustomStreamingResponseBuilder(
result_rows,
graphql2sql.relations_hierarchy,
graphql2sql.selections,
request_args=request.args)
return WorkerResponse.stream_with_context(
response_builder, mimetype='application/x-ndjson')
def _collection(catalog_name, collection_name):
"""Returns the list of entities within the specified collection
A list of entities is returned. This output is paged, default page 1 page size 100
:param catalog_name: e.g. meetbouten
:param collection_name: e.g. meting
:return:
"""
if GOBModel().get_collection(catalog_name, collection_name):
page = int(request.args.get('page', 1))
page_size = int(request.args.get('page_size', 100))
view = request.args.get('view', None)
stream = request.args.get('stream', None) == "true"
ndjson = request.args.get('ndjson', None) == "true"
# If a view is requested and doesn't exist return a 404
if view and not GOBViews().get_view(catalog_name, collection_name,
view):
return not_found(
f'{catalog_name}.{collection_name}?view={view} not found')
view_name = GOBViews().get_view(catalog_name, collection_name,
view)['name'] if view else None
if stream:
entities, convert = query_entities(catalog_name, collection_name,
view_name)
result = stream_entities(entities, convert)
return WorkerResponse.stream_with_context(
result, mimetype='application/json')
elif ndjson:
entities, convert = query_entities(catalog_name, collection_name,
view_name)
result = ndjson_entities(entities, convert)
return WorkerResponse.stream_with_context(
result, mimetype='application/x-ndjson')
else:
result, links = _entities(catalog_name, collection_name, page,
page_size, view_name)
return hal_response(data=result, links=links)
else:
return not_found(f'{catalog_name}.{collection_name} not found')
def worker_status(worker_id):
"""
Returns the status of a worker: dict(id, status, size)
:param worker_id:
:return:
"""
status = WorkerResponse.get_status(worker_id)
if status:
return jsonify(status)
else:
return _worker_not_found(worker_id)
def test_writeResponse(self):
worker_response = WorkerResponse()
WorkerResponse._YIELD_PROGRESS_INTERVAL = 1
worker_response._write_response = lambda rows: sleep(3)
worker_response.is_finished = lambda id: False
result = [r for r in worker_response.write_response([])]
self.assertTrue("FAILURE" in result)
worker_response.is_finished = lambda id: True
result = [r for r in worker_response.write_response([])]
self.assertTrue("OK" in result)
def _dump(catalog_name, collection_name):
"""
Dump all entities in the requested format. Currently only csv
:param catalog_name:
:param collection_name:
:return: Streaming response of all entities in csv format with header
"""
method = request.method
if method == 'GET':
format = request.args.get('format')
exclude_deleted = request.args.get('exclude_deleted') == 'true'
filter = (lambda table: getattr(table, FIELD.DATE_DELETED).is_(None)
) if exclude_deleted else None
entities, model = dump_entities(catalog_name,
collection_name,
filter=filter)
if format == "csv":
result = csv_entities(entities, model)
return WorkerResponse.stream_with_context(result,
mimetype='text/csv')
elif format == "sql":
return Response(sql_entities(catalog_name, collection_name, model),
mimetype='application/sql')
else:
return f"Unrecognised format parameter '{format}'" if format else "Format parameter not set", 400
elif method == 'POST':
content_type = request.content_type
if content_type == 'application/json':
config = json.loads(request.data)
result = dump_to_db(catalog_name, collection_name, config)
return WorkerResponse.stream_with_context(result,
mimetype='text/plain')
else:
return f"Unrecognised content type '{content_type}'", 400
def test_writeResponse(self, mock_rename, mock_isfile, mock_open):
mock_isfile.return_value = True
worker = WorkerResponse()
result = [r for r in worker.write_response(['row'])]
mock_rename.assert_not_called()
mock_isfile.return_value = False
worker = WorkerResponse()
result = [r for r in worker.write_response(['row'])]
mock_rename.assert_called()
def test_constructor(self):
worker_response = WorkerResponse()
self.assertIsNotNone(worker_response.id)
def test_getResponseFile(self, mock_isfile):
worker_id = 'any worker id'
mock_isfile.return_value = True
result = WorkerResponse.get_response_file(worker_id)
self.assertEqual(result, WorkerResponse._get_filename(worker_id))