def _handle_simple_execute_request(self, request_context: RequestContext, params: SimpleExecuteRequest):
new_owner_uri = str(uuid.uuid4())
connection_service = self._service_provider[utils.constants.CONNECTION_SERVICE_NAME]
connection_info = connection_service.get_connection_info(params.owner_uri)
connection_service.connect(ConnectRequestParams(connection_info.details, new_owner_uri, ConnectionType.QUERY))
new_connection = self._get_connection(new_owner_uri, ConnectionType.QUERY)
execute_params = ExecuteStringParams()
execute_params.query = params.query_string
execute_params.owner_uri = new_owner_uri
def on_query_complete(query_complete_params):
subset_params = SubsetParams()
subset_params.owner_uri = new_owner_uri
subset_params.batch_index = 0
subset_params.result_set_index = 0
subset_params.rows_start_index = 0
resultset_summary = query_complete_params.batch_summaries[0].result_set_summaries[0]
subset_params.rows_count = resultset_summary.row_count
subset = self._get_result_subset(request_context, subset_params)
simple_execute_response = SimpleExecuteResponse(subset.result_subset.rows, subset.result_subset.row_count, resultset_summary.column_info)
request_context.send_response(simple_execute_response)
worker_args = ExecuteRequestWorkerArgs(new_owner_uri, new_connection, request_context, ResultSetStorageType.FILE_STORAGE,
on_query_complete=on_query_complete)
self._start_query_execution_thread(request_context, execute_params, worker_args)
def server_info_is_cloud_internal(self, host_suffix, is_cloud):
"""Test that the connection response handles cloud connections correctly"""
# Set up the parameters for the connection
connection_uri = 'someuri'
connection_details = ConnectionDetails()
connection_details.options = {
'user': '******',
'password': '******',
'host': f'myserver{host_suffix}',
'dbname': 'postgres'
}
connection_type = ConnectionType.DEFAULT
# Set up the mock connection for psycopg2's connect method to return
mock_connection = MockPsycopgConnection(
dsn_parameters={
'host': f'myserver{host_suffix}',
'dbname': 'postgres',
'user': '******'
})
# Set up the connection service and call its connect method with the
# supported options
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_connection)):
response = self.connection_service.connect(
ConnectRequestParams(connection_details, connection_uri,
connection_type))
# Verify that the response's serverInfo.isCloud attribute is set correctly
self.assertIsNotNone(response.connection_id)
self.assertIsNotNone(response.server_info.server_version)
self.assertEqual(response.server_info.is_cloud, is_cloud)
def test_get_connection_creates_connection(self):
"""Test that get_connection creates a new connection when none exists for the given URI and type"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type = ConnectionType.EDIT
mock_connection = MockPGServerConnection(cur=None,
host='myserver',
name='postgres',
user='******')
# Insert a ConnectionInfo object into the connection service's map
connection_details = ConnectionDetails.from_data({})
connection_info = ConnectionInfo(connection_uri, connection_details)
self.connection_service.owner_to_connection_map[
connection_uri] = connection_info
with mock.patch(
'ossdbtoolsservice.driver.connection_manager.ConnectionManager._create_connection',
new=mock.Mock(
return_value=mock_connection)) as mock_psycopg2_connect:
# Open the connection
self.connection_service.connect(
ConnectRequestParams(connection_details, connection_uri,
connection_type))
# Get the connection
connection = self.connection_service.get_connection(
connection_uri, connection_type)
self.assertEqual(connection, mock_connection)
mock_psycopg2_connect.assert_called_once()
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 get_connection(self, owner_uri: str, connection_type: ConnectionType) -> Optional[ServerConnection]:
"""
Get a connection for the given owner URI and connection type
:raises ValueError: If there is no connection associated with the provided URI
"""
connection_info = self.owner_to_connection_map.get(owner_uri)
if connection_info is None:
raise ValueError('No connection associated with given owner URI')
if not connection_info.has_connection(connection_type):
self.connect(ConnectRequestParams(connection_info.details, owner_uri, connection_type))
return connection_info.get_connection(connection_type)
def _create_connection(
self, connection_key: str,
conn_info: ConnectionInfo) -> Optional[ServerConnection]:
conn_service = self._connection_service
key_uri = INTELLISENSE_URI + connection_key
connect_request = ConnectRequestParams(conn_info.details, key_uri,
ConnectionType.INTELLISENSE)
connect_result = conn_service.connect(connect_request)
if connect_result.error_message is not None:
raise RuntimeError(connect_result.error_message)
connection = conn_service.get_connection(key_uri,
ConnectionType.INTELLISENSE)
return connection
def handle_change_database_request(self, request_context: RequestContext,
params: ChangeDatabaseRequestParams) -> bool:
"""change database of an existing connection or create a new connection
with default database from input"""
connection_info: ConnectionInfo = self.get_connection_info(params.owner_uri)
if connection_info is None:
return False
connection_info_params: Dict[str, str] = connection_info.details.options.copy()
connection_info_params["dbname"] = params.new_database
connection_details: ConnectionDetails = ConnectionDetails.from_data(connection_info_params)
connection_request_params: ConnectRequestParams = ConnectRequestParams(connection_details, params.owner_uri, ConnectionType.DEFAULT)
self.handle_connect_request(request_context, connection_request_params)
def _create_connection(self, session: ObjectExplorerSession,
database_name: str) -> Optional[ServerConnection]:
conn_service = self._service_provider[
utils.constants.CONNECTION_SERVICE_NAME]
options = session.connection_details.options.copy()
options['dbname'] = database_name
conn_details = ConnectionDetails.from_data(options)
key_uri = session.id + database_name
connect_request = ConnectRequestParams(conn_details, key_uri,
ConnectionType.OBJECT_EXLPORER)
connect_result = conn_service.connect(connect_request)
if connect_result.error_message is not None:
raise RuntimeError(connect_result.error_message)
connection = conn_service.get_connection(
key_uri, ConnectionType.OBJECT_EXLPORER)
return connection
def run_on_connect_callback(self, conn_type: ConnectionType,
expect_callback: bool) -> None:
"""Inner function for callback tests that verifies expected behavior given different connection types"""
callbacks = [MagicMock(), MagicMock()]
for callback in callbacks:
self.connection_service.register_on_connect_callback(callback)
# Set up the parameters for the connection
connection_uri = 'someuri'
connection_details = ConnectionDetails()
connection_details.options = {
'user': '******',
'password': '******',
'host': f'myserver',
'dbname': 'postgres'
}
connection_type = conn_type
# Set up the mock connection for psycopg2's connect method to return
mock_psycopg_connection = MockPsycopgConnection(
dsn_parameters={
'host': f'myserver',
'dbname': 'postgres',
'user': '******'
})
# Set up the connection service and call its connect method with the
# supported options
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_psycopg_connection)):
self.connection_service.connect(
ConnectRequestParams(connection_details, connection_uri,
connection_type))
self.connection_service.get_connection(connection_uri, conn_type)
# ... The mock config change callbacks should have been called
for callback in callbacks:
if (expect_callback):
callback.assert_called_once()
# Verify call args match expected
callargs: ConnectionInfo = callback.call_args[0][0]
self.assertEqual(callargs.owner_uri, connection_uri)
else:
callback.assert_not_called()