def handle_restore_request(self, request_context: RequestContext,
params: RestoreParams) -> None:
"""
Respond to restore/restore requests by performing a restore
"""
connection_info: ConnectionInfo = self._service_provider[
constants.CONNECTION_SERVICE_NAME].get_connection_info(
params.owner_uri)
if connection_info is None:
request_context.send_error(
'No connection corresponding to the given owner URI'
) # TODO: Localize
return
provider: str = self._service_provider.provider
host = connection_info.details.options['host']
database = connection_info.details.options['dbname']
task = Task(
'Restore',
f'Host: {host}, Database: {database}',
provider,
host,
database,
request_context, # TODO: Localize
functools.partial(_perform_restore, connection_info, params))
self._service_provider[constants.TASK_SERVICE_NAME].start_task(task)
request_context.send_response({})
def _get_session(
self, request_context: RequestContext,
params: ExpandParameters) -> Optional[ObjectExplorerSession]:
try:
utils.validate.is_not_none('params', params)
utils.validate.is_not_none_or_whitespace('params.node_path',
params.node_path)
utils.validate.is_not_none_or_whitespace('params.session_id',
params.session_id)
session = self._session_map.get(params.session_id)
if session is None:
raise ValueError(
f'OE session with ID {params.session_id} does not exist'
) # TODO: Localize
if not session.is_ready:
raise ValueError(
f'Object Explorer session with ID {params.session_id} is not ready, yet.'
) # TODO: Localize
request_context.send_response(True)
return session
except Exception as e:
message = f'Failed to expand node: {str(e)}' # TODO: Localize
if self._service_provider.logger is not None:
self._service_provider.logger.error(message)
request_context.send_error(message)
return
def _handle_close_session_request(self, request_context: RequestContext,
params: CloseSessionParameters) -> None:
"""Handle close Object Explorer" sessions request"""
try:
utils.validate.is_not_none('params', params)
# Try to remove the session
session = self._session_map.pop(params.session_id, None)
if session is not None:
self._close_database_connections(session)
conn_service = self._service_provider[
utils.constants.CONNECTION_SERVICE_NAME]
connect_result = conn_service.disconnect(
session.id, ConnectionType.OBJECT_EXLPORER)
if not connect_result:
if self._service_provider.logger is not None:
self._service_provider.logger.info(
f'Could not close the OE session with Id {session.id}'
)
request_context.send_response(False)
else:
request_context.send_response(True)
else:
request_context.send_response(False)
except Exception as e:
message = f'Failed to close OE session: {str(e)}' # TODO: Localize
if self._service_provider.logger is not None:
self._service_provider.logger.error(message)
request_context.send_error(message)
def _dispose(self, request_context: RequestContext, params: DisposeRequest) -> None:
try:
self._active_sessions.pop(params.owner_uri)
except KeyError:
request_context.send_error('Edit data session not found')
request_context.send_response(DisposeResponse())
def _handle_session_request(self, session_operation_request: SessionOperationRequest,
request_context: RequestContext, session_operation: Callable):
edit_session = self._get_active_session(session_operation_request.owner_uri)
try:
result = session_operation(edit_session)
request_context.send_response(result)
except Exception as ex:
request_context.send_error(str(ex))
self._logger.error(str(ex))
def _metadata_list_worker(self, request_context: RequestContext,
params: MetadataListParameters) -> None:
try:
metadata = self._list_metadata(params.owner_uri)
request_context.send_response(MetadataListResponse(metadata))
except Exception as e:
if self._service_provider.logger is not None:
self._service_provider.logger.exception(
'Unhandled exception while executing the metadata list worker thread'
)
request_context.send_error(
'Unhandled exception while listing metadata: ' +
str(e)) # TODO: Localize
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)
def _start_query_execution_thread(self, request_context: RequestContext, params: ExecuteRequestParamsBase, worker_args: ExecuteRequestWorkerArgs = None):
# Set up batch execution callback methods for sending notifications
def _batch_execution_started_callback(batch: Batch) -> None:
batch_event_params = BatchNotificationParams(batch.batch_summary, worker_args.owner_uri)
_check_and_fire(worker_args.on_batch_start, batch_event_params)
def _batch_execution_finished_callback(batch: Batch) -> None:
# Send back notices as a separate message to avoid error coloring / highlighting of text
notices = batch.notices
if notices:
notice_message_params = self.build_message_params(worker_args.owner_uri, batch.id, ''.join(notices), False)
_check_and_fire(worker_args.on_message_notification, notice_message_params)
batch_summary = batch.batch_summary
# send query/resultSetComplete response
result_set_params = self.build_result_set_complete_params(batch_summary, worker_args.owner_uri)
_check_and_fire(worker_args.on_resultset_complete, result_set_params)
# If the batch was successful, send a message to the client
if not batch.has_error:
rows_message = _create_rows_affected_message(batch)
message_params = self.build_message_params(worker_args.owner_uri, batch.id, rows_message, False)
_check_and_fire(worker_args.on_message_notification, message_params)
# send query/batchComplete and query/complete response
batch_event_params = BatchNotificationParams(batch_summary, worker_args.owner_uri)
_check_and_fire(worker_args.on_batch_complete, batch_event_params)
# Create a new query if one does not already exist or we already executed the previous one
if params.owner_uri not in self.query_results or self.query_results[params.owner_uri].execution_state is ExecutionState.EXECUTED:
query_text = self._get_query_text_from_execute_params(params)
execution_settings = QueryExecutionSettings(params.execution_plan_options, worker_args.result_set_storage_type)
query_events = QueryEvents(None, None, BatchEvents(_batch_execution_started_callback, _batch_execution_finished_callback))
self.query_results[params.owner_uri] = Query(params.owner_uri, query_text, execution_settings, query_events)
elif self.query_results[params.owner_uri].execution_state is ExecutionState.EXECUTING:
request_context.send_error('Another query is currently executing.') # TODO: Localize
return
thread = threading.Thread(
target=self._execute_query_request_worker,
args=(worker_args,)
)
self.owner_to_thread_map[params.owner_uri] = thread
thread.daemon = True
thread.start()
def _handle_scriptas_request(self, request_context: RequestContext, params: ScriptAsParameters) -> None:
try:
utils.validate.is_not_none('params', params)
scripting_operation = params.operation
connection_service = self._service_provider[utils.constants.CONNECTION_SERVICE_NAME]
connection = connection_service.get_connection(params.owner_uri, ConnectionType.QUERY)
object_metadata = self.create_metadata(params)
scripter = Scripter(connection)
script = scripter.script(scripting_operation, object_metadata)
request_context.send_response(ScriptAsResponse(params.owner_uri, script))
except Exception as e:
if self._service_provider.logger is not None:
self._service_provider.logger.exception('Scripting operation failed')
request_context.send_error(str(e), params)
def handle_list_databases(self, request_context: RequestContext, params: ListDatabasesParams):
"""List all databases on the server that the given URI has a connection to"""
connection = None
try:
connection = self.get_connection(params.owner_uri, ConnectionType.DEFAULT)
except ValueError as err:
request_context.send_error(str(err))
return
query_results = None
try:
query_results = connection.list_databases()
except Exception as err:
if self._service_provider is not None and self._service_provider.logger is not None:
self._service_provider.logger.exception('Error listing databases')
request_context.send_error(str(err))
return
database_names = [result[0] for result in query_results]
request_context.send_response(ListDatabasesResponse(database_names))
def _edit_initialize(self, request_context: RequestContext, params: InitializeEditParams) -> None:
utils.validate.is_object_params_not_none_or_whitespace('params', params, 'owner_uri', 'schema_name', 'object_name', 'object_type')
connection = self._connection_service.get_connection(params.owner_uri, ConnectionType.QUERY)
session = DataEditorSession(SmoEditTableMetadataFactory())
self._active_sessions[params.owner_uri] = session
if params.query_string is not None:
request_context.send_error('Edit data with custom query is not supported currently.')
return
def query_executer(query: str, columns: List[DbColumn], on_query_execution_complete: Callable):
def on_resultset_complete(result_set_params: ResultSetNotificationParams):
result_set_params.result_set_summary.column_info = columns
request_context.send_notification(RESULT_SET_UPDATED_NOTIFICATION, 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_query_complete(query_complete_params: QueryCompleteNotificationParams):
on_query_execution_complete(DataEditSessionExecutionState(self._query_execution_service.get_query(params.owner_uri)))
request_context.send_notification(QUERY_COMPLETE_NOTIFICATION, query_complete_params)
worker_args = ExecuteRequestWorkerArgs(params.owner_uri, connection, request_context, ResultSetStorageType.IN_MEMORY,
on_resultset_complete=on_resultset_complete, on_query_complete=on_query_complete)
execution_params = ExecuteStringParams()
execution_params.query = query
execution_params.owner_uri = params.owner_uri
self._query_execution_service._start_query_execution_thread(request_context, execution_params, worker_args)
def on_success():
request_context.send_notification(SESSION_READY_NOTIFICATION, SessionReadyNotificationParams(params.owner_uri, True, None))
def on_failure(error: str):
request_context.send_notification(SESSION_READY_NOTIFICATION, SessionReadyNotificationParams(params.owner_uri, False, error))
session.initialize(params, connection, query_executer, on_success, on_failure)
request_context.send_response({})
def _handle_create_session_request(self, request_context: RequestContext,
params: ConnectionDetails) -> None:
"""Handle a create object explorer session request"""
# Step 1: Create the session
try:
# Make sure we have the appropriate session params
utils.validate.is_not_none('params', params)
# Use the provider's default db if db name was not specified
if params.database_name is None or params.database_name == '':
if self._provider == utils.constants.MYSQL_PROVIDER_NAME:
params.database_name = self._service_provider[
utils.constants.
WORKSPACE_SERVICE_NAME].configuration.my_sql.default_database
elif self._provider == utils.constants.PG_PROVIDER_NAME:
params.database_name = self._service_provider[
utils.constants.
WORKSPACE_SERVICE_NAME].configuration.pgsql.default_database
# Use the provider's default port if port number was not specified
if not params.port:
params.port = utils.constants.DEFAULT_PORT[self._provider]
# Generate the session ID and create/store the session
session_id = self._generate_session_uri(params, self._provider)
session: ObjectExplorerSession = ObjectExplorerSession(
session_id, params)
# Add the session to session map in a lock to prevent race conditions between check and add
with self._session_lock:
if session_id in self._session_map:
# Removed the exception for now. But we need to investigate why we would get this
if self._service_provider.logger is not None:
self._service_provider.logger.error(
f'Object explorer session for {session_id} already exists!'
)
request_context.send_response(False)
return
self._session_map[session_id] = session
# Respond that the session was created
response = CreateSessionResponse(session_id)
request_context.send_response(response)
except Exception as e:
message = f'Failed to create OE session: {str(e)}'
if self._service_provider.logger is not None:
self._service_provider.logger.error(message)
request_context.send_error(message)
return
# Step 2: Connect the session and lookup the root node asynchronously
try:
session.init_task = threading.Thread(
target=self._initialize_session,
args=(request_context, session))
session.init_task.daemon = True
session.init_task.start()
except Exception as e:
# TODO: Localize
self._session_created_error(
request_context, session,
f'Failed to start OE init task: {str(e)}')
