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 _connect_and_respond(self, request_context: RequestContext, params: ConnectRequestParams) -> None:
"""Open a connection and fire the connection complete notification"""
response = self.connect(params)
# Send the connection complete response unless the connection was canceled
if response is not None:
request_context.send_notification(CONNECTION_COMPLETE_METHOD, response)
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 handle_cancellation_request(self, request_context: RequestContext, params: CancelConnectParams) -> None:
"""Cancel a connection attempt in response to a cancellation request"""
cancellation_key = (params.owner_uri, params.type)
with self._cancellation_lock:
connection_found = cancellation_key in self._cancellation_map
if connection_found:
self._cancellation_map[cancellation_key].cancel()
request_context.send_response(connection_found)
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 _initialize_session(self, request_context: RequestContext,
session: ObjectExplorerSession):
conn_service = self._service_provider[
utils.constants.CONNECTION_SERVICE_NAME]
connection = None
try:
# Step 1: Connect with the provided connection details
connect_request = ConnectRequestParams(
session.connection_details, session.id,
ConnectionType.OBJECT_EXLPORER)
connect_result = conn_service.connect(connect_request)
if connect_result is None:
raise RuntimeError(
'Connection was cancelled during connect') # TODO Localize
if connect_result.error_message is not None:
raise RuntimeError(connect_result.error_message)
# Step 2: Get the connection to use for object explorer
connection = conn_service.get_connection(
session.id, ConnectionType.OBJECT_EXLPORER)
# Step 3: Create the Server object for the session and create the root node for the server
session.server = self._server(
connection, functools.partial(self._create_connection,
session))
metadata = ObjectMetadata(session.server.urn_base, None,
'Database',
session.server.maintenance_db_name)
node = NodeInfo()
node.label = session.connection_details.database_name
node.is_leaf = False
node.node_path = session.id
node.node_type = 'Database'
node.metadata = metadata
# Step 4: Send the completion notification to the server
response = SessionCreatedParameters()
response.success = True
response.session_id = session.id
response.root_node = node
response.error_message = None
request_context.send_notification(SESSION_CREATED_METHOD, response)
# Mark the session as complete
session.is_ready = True
except Exception as e:
# Return a notification that an error occurred
message = f'Failed to initialize object explorer session: {str(e)}' # TODO Localize
self._session_created_error(request_context, session, message)
# Attempt to clean up the connection
if connection is not None:
conn_service.disconnect(session.id,
ConnectionType.OBJECT_EXLPORER)
def _edit_subset(self, request_context: RequestContext, params: EditSubsetParams) -> None:
session: DataEditorSession = self._active_sessions.get(params.owner_uri)
rows = session.get_rows(params.owner_uri, params.row_start_index, params.row_start_index + params.row_count)
self._handle_create_row_default_values(rows, session)
edit_subset_result = EditSubsetResponse(len(rows), rows)
request_context.send_response(edit_subset_result)
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 handle_connect_request(self, request_context: RequestContext,
params: ConnectRequestParams) -> None:
"""Kick off a connection in response to an incoming connection request"""
thread = threading.Thread(target=self._connect_and_respond,
args=(request_context, params))
thread.daemon = True
thread.start()
self.owner_to_thread_map[params.owner_uri] = thread
request_context.send_response(True)
def _expand_node_error(self, request_context: RequestContext,
params: ExpandParameters, message: str):
if self._service_provider.logger is not None:
self._service_provider.logger.warning(
f'OE service errored while expanding node: {message}')
response = ExpandCompletedParameters(params.session_id,
params.node_path)
response.error_message = f'Failed to expand node: {message}' # TODO: Localize
request_context.send_notification(EXPAND_COMPLETED_METHOD, response)
def send_definition_using_connected_completions(
self, request_context: RequestContext,
scriptparseinfo: ScriptParseInfo, params: TextDocumentPosition,
context: ConnectionContext) -> bool:
if not context or not context.is_connected:
return False
definition_result: DefinitionResult = None
completer: Completer = context.completer
completions: List[Completion] = completer.get_completions(
scriptparseinfo.document, None)
if completions:
word_under_cursor = scriptparseinfo.document.get_word_under_cursor(
)
matching_completion = next(
completion for completion in completions
if completion.display == word_under_cursor)
if matching_completion:
connection = self._connection_service.get_connection(
params.text_document.uri, ConnectionType.QUERY)
scripter_instance = scripter.Scripter(connection)
object_metadata = ObjectMetadata(
None, None, matching_completion.display_meta,
matching_completion.display, matching_completion.schema)
create_script = scripter_instance.script(
ScriptOperation.CREATE, object_metadata)
if create_script:
with tempfile.NamedTemporaryFile(
mode='wt',
delete=False,
encoding='utf-8',
suffix='.sql',
newline=None) as namedfile:
namedfile.write(create_script)
if namedfile.name:
file_uri = "file:///" + namedfile.name.strip('/')
location_in_script = Location(
file_uri, Range(Position(0, 1), Position(1,
1)))
definition_result = DefinitionResult(
False, None, [
location_in_script,
])
request_context.send_response(
definition_result.locations)
return True
if definition_result is None:
request_context.send_response(DefinitionResult(True, '', []))
return False
def __init__(self):
RequestContext.__init__(self, None, None)
self.last_response_params = None
self.last_notification_method = None
self.last_notification_params = None
self.last_error_message = None
self.send_response = mock.Mock(side_effect=self.send_response_impl)
self.send_notification = mock.Mock(
side_effect=self.send_notification_impl)
self.send_error = mock.Mock(side_effect=self.send_error_impl)
self.send_unhandled_error_response = mock.Mock(
side_effect=self.send_unhandled_error_response_impl)
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 _send_default_completions(self, request_context: RequestContext,
script_file: ScriptFile,
params: TextDocumentPosition) -> bool:
response = []
line: str = script_file.get_line(params.position.line)
(token_text,
text_range) = TextUtilities.get_text_and_range(params.position, line)
if token_text:
completions = self._completion_helper.get_matches(
token_text, text_range, self.should_lowercase)
response = completions
request_context.send_response(response)
return True
def _expand_node_thread(self, is_refresh: bool,
request_context: RequestContext,
params: ExpandParameters,
session: ObjectExplorerSession):
try:
response = ExpandCompletedParameters(session.id, params.node_path)
response.nodes = self._route_request(is_refresh, session,
params.node_path)
request_context.send_notification(EXPAND_COMPLETED_METHOD,
response)
except Exception as e:
self._expand_node_error(request_context, params, str(e))
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_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 handle_definition_request(
self, request_context: RequestContext,
text_document_position: TextDocumentPosition) -> None:
request_context.send_notification(
STATUS_CHANGE_NOTIFICATION,
StatusChangeParams(
owner_uri=text_document_position.text_document.uri,
status="DefinitionRequested"))
def do_send_default_empty_response():
request_context.send_response([])
if self.should_skip_intellisense(
text_document_position.text_document.uri):
do_send_default_empty_response()
return
script_file: ScriptFile = self._workspace_service.workspace.get_file(
text_document_position.text_document.uri)
if script_file is None:
do_send_default_empty_response()
return
script_parse_info: ScriptParseInfo = self.get_script_parse_info(
text_document_position.text_document.uri,
create_if_not_exists=False)
if not script_parse_info or not script_parse_info.can_queue():
do_send_default_empty_response()
return
cursor_position: int = len(
script_file.get_text_in_range(
Range.from_data(0, 0, text_document_position.position.line,
text_document_position.position.character)))
text: str = script_file.get_all_text()
script_parse_info.document = Document(text, cursor_position)
operation = QueuedOperation(
script_parse_info.connection_key,
functools.partial(self.send_definition_using_connected_completions,
request_context, script_parse_info,
text_document_position),
functools.partial(do_send_default_empty_response))
self.operations_queue.add_operation(operation)
request_context.send_notification(
STATUS_CHANGE_NOTIFICATION,
StatusChangeParams(
owner_uri=text_document_position.text_document.uri,
status="DefinitionRequestCompleted"))
def _handle_dmp_capabilities_request(
self, request_context: RequestContext,
params: Optional[CapabilitiesRequestParams]) -> None:
"""
Sends the capabilities of the tools service data protocol features
:param request_context: Context of the request
:param params: Parameters for the capabilities request
"""
provider: str = self._service_provider.provider
capabilities = SERVER_CAPABILITIES_MAP[provider]
result = CapabilitiesResult(capabilities)
# Send the response
request_context.send_response(result)
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 _session_created_error(self, request_context: RequestContext,
session: ObjectExplorerSession, message: str):
if self._service_provider.logger is not None:
self._service_provider.logger.warning(
f'OE service errored while creating session: {message}')
# Create error notification
response = SessionCreatedParameters()
response.success = False
response.session_id = session.id
response.root_node = None
response.error_message = message
request_context.send_notification(SESSION_CREATED_METHOD, response)
# Clean up the session from the session map
self._session_map.pop(session.id)
def test_initialization_with_none_as_schema_param(self):
request_context = RequestContext(None, None)
self._initialize_edit_request.schema_name = None
errormsg = 'Parameter schema_name contains a None, empty, or whitespace string'
self.assert_exception_on_method_call(
ValueError, errormsg, self._service_under_test._edit_initialize,
request_context, self._initialize_edit_request)
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 test_initialization_with_empty_object_type(self):
request_context = RequestContext(None, None)
self._initialize_edit_request.object_type = ' '
errormsg = 'Parameter object_type contains a None, empty, or whitespace string'
self.assert_exception_on_method_call(
ValueError, errormsg, self._service_under_test._edit_initialize,
request_context, self._initialize_edit_request)
def send_connected_completions(self, request_context: RequestContext,
scriptparseinfo: ScriptParseInfo,
params: TextDocumentPosition,
context: ConnectionContext) -> bool:
if not context or not context.is_connected:
return False
# Else use the completer to query for completions
completer: Completer = context.completer
completions: List[Completion] = completer.get_completions(
scriptparseinfo.document, None)
if completions:
response = [
LanguageService.to_completion_item(completion, params)
for completion in completions
]
request_context.send_response(response)
return True
# Else return false so the timeout task can be sent instead
return False
def test_initialization(self, mockdataeditorsession):
queue = Queue()
message = JSONRPCMessage.from_dictionary({
'id': '123',
'method': 'edit/initialize',
'params': {}
})
request_context = RequestContext(message, queue)
self._service_under_test._edit_initialize(
request_context, self._initialize_edit_request)
mockdataeditorsession.assert_called()
def _handle_get_database_info_request(
self, request_context: RequestContext,
params: GetDatabaseInfoParameters) -> None:
# Retrieve the connection from the connection service
connection_service = self._service_provider[
constants.CONNECTION_SERVICE_NAME]
connection: ServerConnection = connection_service.get_connection(
params.owner_uri, ConnectionType.DEFAULT)
# Get database owner
owner_result = connection.get_database_owner()
size_result = connection.get_database_size(connection.database_name)
# Set up and send the response
options = {
DatabaseInfo.DBNAME: connection.database_name,
DatabaseInfo.OWNER: owner_result,
DatabaseInfo.SIZE: size_result
}
request_context.send_response(
GetDatabaseInfoResponse(DatabaseInfo(options)))
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_initialize_request(
request_context: RequestContext,
params: Optional[InitializeRequestParams]) -> None:
"""
Sends the capabilities of the tools service language features
:param request_context: Context for the request
:param params: Initialization request parameters
"""
capabilities = ServerCapabilities(
text_document_sync=TextDocumentSyncKind.INCREMENTAL,
definition_provider=True,
references_provider=False,
document_formatting_provider=True,
document_range_formatting_provider=True,
document_highlight_provider=False,
hover_provider=False,
completion_provider=CompletionOptions(
True, ['.', '-', ':', '\\', '[', '"']))
result = InitializeResult(capabilities)
# Send the request
request_context.send_response(result)
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({})
