def _handle_cancel_query_request(self, request_context: RequestContext,
params: QueryCancelParams):
"""Handles a 'query/cancel' request"""
try:
if params.owner_uri in self.query_results:
query = self.query_results[params.owner_uri]
else:
request_context.send_response(
QueryCancelResult(NO_QUERY_MESSAGE)) # TODO: Localize
return
# Only cancel the query if we're in a cancellable state
if query.execution_state is ExecutionState.EXECUTED:
request_context.send_response(
QueryCancelResult(
'Query already executed')) # TODO: Localize
return
query.is_canceled = True
# Only need to do additional work to cancel the query
# if it's currently running
if query.execution_state is ExecutionState.EXECUTING:
self.cancel_query(params.owner_uri)
request_context.send_response(QueryCancelResult())
except Exception as e:
if self._service_provider.logger is not None:
self._service_provider.logger.exception(str(e))
request_context.send_unhandled_error_response(e)
def _handle_execute_query_request(
self, request_context: RequestContext,
params: ExecuteRequestParamsBase) -> None:
"""Kick off thread to execute query in response to an incoming execute query request"""
def before_query_initialize(before_query_initialize_params):
# Send a response to indicate that the query was kicked off
request_context.send_response(before_query_initialize_params)
def on_batch_start(batch_event_params):
request_context.send_notification(BATCH_START_NOTIFICATION,
batch_event_params)
def on_message_notification(notice_message_params):
request_context.send_notification(MESSAGE_NOTIFICATION,
notice_message_params)
def on_resultset_complete(result_set_params):
request_context.send_notification(
RESULT_SET_AVAILABLE_NOTIFICATION, result_set_params)
request_context.send_notification(RESULT_SET_COMPLETE_NOTIFICATION,
result_set_params)
def on_batch_complete(batch_event_params):
request_context.send_notification(BATCH_COMPLETE_NOTIFICATION,
batch_event_params)
def on_query_complete(query_complete_params):
request_context.send_notification(QUERY_COMPLETE_NOTIFICATION,
query_complete_params)
# Get a connection for the query
try:
conn = self._get_connection(params.owner_uri, ConnectionType.QUERY)
except Exception as e:
if self._service_provider.logger is not None:
self._service_provider.logger.exception(
'Encountered exception while handling query request'
) # TODO: Localize
request_context.send_unhandled_error_response(e)
return
worker_args = ExecuteRequestWorkerArgs(
params.owner_uri, conn, request_context,
ResultSetStorageType.FILE_STORAGE, before_query_initialize,
on_batch_start, on_message_notification, on_resultset_complete,
on_batch_complete, on_query_complete)
self._start_query_execution_thread(request_context, params,
worker_args)
def _handle_dispose_request(self, request_context: RequestContext,
params: QueryDisposeParams):
try:
if params.owner_uri not in self.query_results:
request_context.send_error(NO_QUERY_MESSAGE) # TODO: Localize
return
# Make sure to cancel the query first if it's not executed.
# If it's not started, then make sure it never starts. If it's executing, make sure
# that we stop it
if self.query_results[
params.
owner_uri].execution_state is not ExecutionState.EXECUTED:
self.cancel_query(params.owner_uri)
del self.query_results[params.owner_uri]
request_context.send_response({})
except Exception as e:
request_context.send_unhandled_error_response(e)
