def test_build_connection_response(self):
"""Test that the connection response is built correctly"""
# Set up the test with mock data
server_name = 'testserver'
db_name = 'testdb'
user = '******'
mock_connection = MockPGServerConnection(cur=None,
host=server_name,
name=db_name,
user=user)
connection_type = ConnectionType.EDIT
connection_details = ConnectionDetails.from_data(opts={})
owner_uri = 'test_uri'
connection_info = ConnectionInfo(owner_uri, connection_details)
connection_info._connection_map = {connection_type: mock_connection}
# If I build a connection response for the connection
response = ossdbtoolsservice.connection.connection_service._build_connection_response(
connection_info, connection_type)
# Then the response should have accurate information about the connection
self.assertEqual(response.owner_uri, owner_uri)
self.assertEqual(
response.server_info.server_version,
str(mock_connection.server_version[0]) + "." +
str(mock_connection.server_version[1]) + "." +
str(mock_connection.server_version[2]))
self.assertEqual(response.server_info.is_cloud, False)
self.assertEqual(response.connection_summary.server_name, server_name)
self.assertEqual(response.connection_summary.database_name, db_name)
self.assertEqual(response.connection_summary.user_name, user)
self.assertEqual(response.type, connection_type)
def setUp(self):
"""Constructor"""
self.default_connection_key = 'server_db_user'
self.mock_connection_service = ConnectionService()
self.mock_server = JSONRPCServer(None, None)
self.mock_service_provider = ServiceProvider(self.mock_server, {},
PG_PROVIDER_NAME, None)
self.mock_service_provider._services[
constants.CONNECTION_SERVICE_NAME] = self.mock_connection_service
self.mock_service_provider._is_initialized = True
# Create connection information for use in the tests
self.connection_details = ConnectionDetails.from_data({})
self.connection_details.server_name = 'test_host'
self.connection_details.database_name = 'test_db'
self.connection_details.user_name = 'user'
self.expected_context_key = 'test_host|test_db|user'
self.expected_connection_uri = INTELLISENSE_URI + self.expected_context_key
self.test_uri = 'test_uri'
self.connection_info = ConnectionInfo(self.test_uri,
self.connection_details)
# Create mock CompletionRefresher to avoid calls to create separate thread
self.refresher_mock = mock.MagicMock()
self.refresh_method_mock = mock.MagicMock()
self.refresher_mock.refresh = self.refresh_method_mock
def test_list_databases_handles_query_failure(self):
"""Test that the list databases handler returns an error if the list databases query fails for any reason"""
# Set up the test with mock data
mock_query_results = [('database1', ), ('database2', )]
connection_uri = 'someuri'
mock_cursor = MockCursor(mock_query_results)
mock_cursor.fetchall.side_effect = psycopg2.ProgrammingError('')
mock_connection = MockPGServerConnection(cur=mock_cursor,
host='myserver',
name='postgres',
user='******')
mock_request_context = utils.MockRequestContext()
# 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
# Verify that calling the listdatabases handler returns the expected
# databases
params = ListDatabasesParams()
params.owner_uri = connection_uri
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_connection)):
self.connection_service.handle_list_databases(
mock_request_context, params)
self.assertIsNone(mock_request_context.last_notification_method)
self.assertIsNone(mock_request_context.last_notification_params)
self.assertIsNone(mock_request_context.last_response_params)
self.assertIsNotNone(mock_request_context.last_error_message)
def test_get_connection_for_existing_connection(self):
"""Test that get_connection returns a connection that already exists for the given URI and type"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type = ConnectionType.EDIT
mock_connection = MockPsycopgConnection(dsn_parameters={
'host': 'myserver',
'dbname': '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
# Get the connection without first creating it
with mock.patch(
'psycopg2.connect', new=mock.Mock(
return_value=mock_connection)) as mock_psycopg2_connect:
connection = self.connection_service.get_connection(
connection_uri, connection_type)
mock_psycopg2_connect.assert_called_once()
self.assertEqual(connection._conn, mock_connection)
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 test_list_databases(self):
"""Test that the list databases handler correctly lists the connection's databases"""
# Set up the test with mock data
mock_query_results = [('database1', ), ('database2', )]
connection_uri = 'someuri'
mock_psycopg_connection = MockPsycopgConnection(
dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
},
cursor=MockCursor(mock_query_results))
mock_request_context = utils.MockRequestContext()
# 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
# Verify that calling the listdatabases handler returns the expected databases
params = ListDatabasesParams()
params.owner_uri = connection_uri
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_psycopg_connection)):
self.connection_service.handle_list_databases(
mock_request_context, params)
expected_databases = [result[0] for result in mock_query_results]
self.assertEqual(
mock_request_context.last_response_params.database_names,
expected_databases)
def test_changing_options_disconnects_existing_connection(self):
"""
Test that the connect method disconnects an existing connection when trying to open the same connection with
different options
"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type = ConnectionType.DEFAULT
mock_connection = MockPGServerConnection(cur=None,
host='myserver',
name='postgres',
user='******')
# Insert a ConnectionInfo object into the connection service's map
old_connection_details = ConnectionDetails.from_data({
'host': 'myserver',
'dbname': 'postgres',
'user': '******',
'abc': 123
})
old_connection_info = ConnectionInfo(connection_uri,
old_connection_details)
old_connection_info.add_connection(connection_type, mock_connection)
self.connection_service.owner_to_connection_map[
connection_uri] = old_connection_info
# Create a different request with the same owner uri
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
connection_uri,
'type':
connection_type,
'connection': {
'options': {
'host': 'myserver',
'dbname': 'postgres',
'user': '******',
'abc': 234
}
}
})
# Connect with different options, and verify that disconnect was called
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_connection)):
self.connection_service.connect(params)
mock_connection.close.assert_called_once()
def test_same_options_uses_existing_connection(self):
"""Test that the connect method uses an existing connection when connecting again with the same options"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type = ConnectionType.DEFAULT
mock_psycopg_connection = MockPsycopgConnection(dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
})
mock_server_connection = MockPGServerConnection(cur=None,
host='myserver',
name='postgres',
user='******')
# Insert a ConnectionInfo object into the connection service's map
old_connection_details = ConnectionDetails.from_data({
'host': 'myserver',
'dbname': 'postgres',
'user': '******',
'abc': 123
})
old_connection_info = ConnectionInfo(connection_uri,
old_connection_details)
old_connection_info.add_connection(connection_type,
mock_server_connection)
self.connection_service.owner_to_connection_map[
connection_uri] = old_connection_info
# Connect with identical options, and verify that disconnect was not called
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
connection_uri,
'type':
connection_type,
'connection': {
'options': old_connection_details.options
}
})
with mock.patch('psycopg2.connect',
new=mock.Mock(return_value=mock_psycopg_connection)
) as mock_psycopg2_connect:
response = self.connection_service.connect(params)
mock_psycopg2_connect.assert_not_called()
mock_psycopg_connection.close.assert_not_called()
self.assertIsNotNone(response.connection_id)
def test_on_connect_sends_notification(self):
"""
Test that the service sends an intellisense ready notification after handling an on connect notification from the connection service.
This is a slightly more end-to-end test that verifies calling through to the queue layer
"""
# If: I create a new language service
service: LanguageService = self._init_service_with_flow_validator()
conn_info = ConnectionInfo(
'file://msuri.sql',
ConnectionDetails.from_data({
'host': None,
'dbname': 'TEST_DBNAME',
'user': '******'
}))
connect_result = mock.MagicMock()
connect_result.error_message = None
self.mock_connection_service.get_connection = mock.Mock(
return_value=mock.MagicMock())
self.mock_connection_service.connect = mock.MagicMock(
return_value=connect_result)
def validate_success_notification(response: IntelliSenseReadyParams):
self.assertEqual(response.owner_uri, conn_info.owner_uri)
# When: I notify of a connection complete for a given URI
self.flow_validator.add_expected_notification(
IntelliSenseReadyParams, INTELLISENSE_READY_NOTIFICATION,
validate_success_notification)
refresher_mock = mock.MagicMock()
refresh_method_mock = mock.MagicMock()
refresher_mock.refresh = refresh_method_mock
patch_path = 'ossdbtoolsservice.language.operations_queue.CompletionRefresher'
with mock.patch(patch_path) as refresher_patch:
refresher_patch.return_value = refresher_mock
task: threading.Thread = service.on_connect(conn_info)
# And when refresh is "complete"
refresh_method_mock.assert_called_once()
callback = refresh_method_mock.call_args[0][0]
self.assertIsNotNone(callback)
callback(None)
# Wait for task to return
task.join()
# Then:
# an intellisense ready notification should be sent for that URI
self.flow_validator.validate()
# ... and the scriptparseinfo should be created
info: ScriptParseInfo = service.get_script_parse_info(
conn_info.owner_uri)
self.assertIsNotNone(info)
# ... and the info should have the connection key set
self.assertEqual(info.connection_key,
OperationsQueue.create_key(conn_info))
def test_disconnect_all_types(self):
"""Test that the disconnect method calls close on a all open connection types when no type is given"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type_1 = ConnectionType.DEFAULT
connection_type_2 = ConnectionType.EDIT
mock_connection_1 = MockPGServerConnection(cur=None,
host='myserver1',
name='postgres1',
user='******')
mock_connection_2 = MockPGServerConnection(cur=None,
host='myserver2',
name='postgres2',
user='******')
# Insert a ConnectionInfo object into the connection service's map
old_connection_details = ConnectionDetails.from_data({'abc': 123})
old_connection_info = ConnectionInfo(connection_uri,
old_connection_details)
old_connection_info.add_connection(connection_type_1,
mock_connection_1)
old_connection_info.add_connection(connection_type_2,
mock_connection_2)
self.connection_service.owner_to_connection_map[
connection_uri] = old_connection_info
# Close the connection by calling disconnect
response = self.connection_service._close_connections(
old_connection_info)
mock_connection_1.close.assert_called_once()
mock_connection_2.close.assert_called_once()
self.assertTrue(response)
def test_get_connection_info(self):
"""Test that get_connection_info returns the ConnectionInfo object corresponding to a connection"""
# Set up the test with mock data
connection_uri = 'someuri'
# 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
# Get the connection info
actual_connection_info = self.connection_service.get_connection_info(
connection_uri)
self.assertIs(actual_connection_info, connection_info)
def test_disconnect_for_invalid_connection(self):
"""Test that the disconnect method returns false when called on a connection that does not exist"""
# Set up the test with mock data
connection_uri = 'someuri'
connection_type_1 = ConnectionType.DEFAULT
mock_connection_1 = MockPGServerConnection(cur=None,
host='myserver1',
name='postgres1',
user='******')
# Insert a ConnectionInfo object into the connection service's map
old_connection_details = ConnectionDetails.from_data({'abc': 123})
old_connection_info = ConnectionInfo(connection_uri,
old_connection_details)
old_connection_info.add_connection(connection_type_1,
mock_connection_1)
self.connection_service.owner_to_connection_map[
connection_uri] = old_connection_info
# Close the connection by calling disconnect
response = self.connection_service._close_connections(
old_connection_info, ConnectionType.EDIT)
mock_connection_1.close.assert_not_called()
self.assertFalse(response)
def do_test():
# When I add context for 2 URIs with same connection details
operations_queue = OperationsQueue(self.mock_service_provider)
with mock.patch(COMPLETIONREFRESHER_PATH_PATH) as refresher_patch:
refresher_patch.return_value = self.refresher_mock
operations_queue.add_connection_context(self.connection_info)
conn_info2 = ConnectionInfo('newuri',
self.connection_info.details)
operations_queue.add_connection_context(conn_info2)
# Then I expect to only have connection
connect_mock: mock.MagicMock = self.mock_connection_service.connect
connect_mock.assert_called_once()
connect_params: ConnectRequestParams = connect_mock.call_args[
0][0]
self.assertEqual(connect_params.owner_uri,
self.expected_connection_uri)
def _perform_restore(connection_info: ConnectionInfo, params: RestoreParams,
task: Task) -> TaskResult:
"""Call out to pg_restore to restore from a backup"""
try:
connection = connection_info.get_connection(ConnectionType.DEFAULT)
pg_restore_location = _get_pg_exe_path('pg_restore',
connection.server_version)
except ValueError as e:
return TaskResult(TaskStatus.FAILED, str(e))
pg_restore_args = [pg_restore_location]
pg_restore_args += _get_backup_restore_connection_params(
connection_info.details.options)
pg_restore_args.append(params.options.path)
# Remove the options that were already used, and pass the rest so that they can be automatically serialized
options = params.options.__dict__.copy()
del options['path']
return _perform_backup_restore(connection_info, pg_restore_args, options,
task)
def _perform_backup(connection_info: ConnectionInfo, params: BackupParams,
task: Task) -> TaskResult:
"""Call out to pg_dump to do a backup"""
try:
connection = connection_info.get_connection(ConnectionType.DEFAULT)
pg_dump_location = _get_pg_exe_path('pg_dump',
connection.server_version)
except ValueError as e:
return TaskResult(TaskStatus.FAILED, str(e))
pg_dump_args = [
pg_dump_location, f'--file={params.backup_info.path}',
f'--format={_BACKUP_FORMAT_MAP[params.backup_info.type]}'
]
pg_dump_args += _get_backup_restore_connection_params(
connection_info.details.options)
# Remove the options that were already used, and pass the rest so that they can be automatically serialized
options = params.backup_info.__dict__.copy()
del options['path']
del options['type']
return _perform_backup_restore(connection_info, pg_dump_args, options,
task)
def do_test():
# When I add context for 2 URIs with same connection details
operations_queue = OperationsQueue(self.mock_service_provider)
with mock.patch(COMPLETIONREFRESHER_PATH_PATH) as refresher_patch:
refresher_patch.return_value = self.refresher_mock
operations_queue.add_connection_context(self.connection_info)
conn_info2 = ConnectionInfo('newuri',
self.connection_info.details)
operations_queue.add_connection_context(conn_info2,
overwrite=True)
# Then I expect to only have 1 connection
# and I expect disconnect and reconnect to have been called
connect_mock: mock.MagicMock = self.mock_connection_service.connect
self.assertEqual(connect_mock.call_count, 2)
self.assertEqual(
connect_mock.call_args_list[0][0][0].owner_uri,
self.expected_connection_uri)
self.assertEqual(
connect_mock.call_args_list[1][0][0].owner_uri,
self.expected_connection_uri)
disconnect_mock: mock.MagicMock = self.mock_connection_service.disconnect
disconnect_mock.assert_called_once()
def setUp(self):
"""Set up the tests with a disaster recovery service and connection service with mock connection info"""
self.disaster_recovery_service = DisasterRecoveryService()
self.connection_service = ConnectionService()
self.task_service = TaskService()
self.disaster_recovery_service._service_provider = utils.get_mock_service_provider(
{
constants.CONNECTION_SERVICE_NAME: self.connection_service,
constants.TASK_SERVICE_NAME: self.task_service
})
# Create connection information for use in the tests
self.connection_details = ConnectionDetails()
self.host = 'test_host'
self.dbname = 'test_db'
self.username = '******'
self.connection_details.options = {
'host': self.host,
'dbname': self.dbname,
'user': self.username,
'port': 5432
}
self.test_uri = 'test_uri'
self.connection_info = ConnectionInfo(self.test_uri,
self.connection_details)
# Create backup parameters for the tests
self.request_context = utils.MockRequestContext()
self.backup_path = 'mock/path/test.sql'
self.backup_type = 'sql'
self.data_only = False
self.no_owner = True
self.schema = 'test_schema'
self.backup_params = BackupParams.from_dict({
'ownerUri': self.test_uri,
'backupInfo': {
'type': self.backup_type,
'path': self.backup_path,
'data_only': self.data_only,
'no_owner': self.no_owner,
'schema': self.schema
}
})
self.restore_path = 'mock/path/test.dump'
self.restore_params = RestoreParams.from_dict({
'ownerUri': self.test_uri,
'options': {
'path': self.restore_path,
'data_only': self.data_only,
'no_owner': self.no_owner,
'schema': self.schema
}
})
self.pg_dump_exe = 'pg_dump'
self.pg_restore_exe = 'pg_restore'
# Create the mock task for the tests
self.mock_action = mock.Mock()
self.mock_task = Task(None, None, None, None, None,
self.request_context, self.mock_action)
self.mock_task.start = mock.Mock()