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 = MockConnection(dsn_parameters={
'host': 'myserver1',
'dbname': 'postgres1',
'user': '******'
})
mock_connection_2 = MockConnection(dsn_parameters={
'host': 'myserver2',
'dbname': '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_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 = MockConnection({
'host': server_name,
'dbname': 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 = pgsqltoolsservice.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, mock_connection.server_version)
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 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 = MockConnection(
dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
},
cursor=mock_cursor)
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_handle_get_database_info_request(self):
"""Test that the database info handler responds with the correct database info"""
uri = 'test_uri'
db_name = 'test_db'
user_name = 'test_user'
# Set up the request parameters
params = GetDatabaseInfoParameters()
params.owner_uri = uri
request_context = MockRequestContext()
# Set up a mock connection and cursor for the test
mock_query_results = [(user_name,)]
mock_cursor = MockCursor(mock_query_results)
mock_connection = MockConnection({'dbname': db_name}, mock_cursor)
self.connection_service.get_connection = mock.Mock(return_value=mock_connection)
# If I send a get_database_info request
self.admin_service._handle_get_database_info_request(request_context, params)
# Then the service responded with the expected information
response = request_context.last_response_params
self.assertIsInstance(response, GetDatabaseInfoResponse)
expected_info = {'owner': user_name}
self.assertEqual(response.database_info.options, expected_info)
# And the service retrieved the owner name using a query with the database name as a parameter
mock_cursor.execute.assert_called_once_with(mock.ANY, (db_name,))
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_connection = MockConnection(
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_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_connect(self):
"""Test that the service connects to a PostgreSQL server"""
# Set up the parameters for the connection
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri': 'someUri',
'type': ConnectionType.DEFAULT,
'connection': {
'options': {
'user': '******',
'password': '******',
'host': 'myserver',
'dbname': 'postgres'
}
}
})
# Set up the mock connection for psycopg2's connect method to return
mock_connection = MockConnection(dsn_parameters={
'host': '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_connection)):
response = self.connection_service.connect(params)
# Verify that psycopg2's connection method was called and that the
# response has a connection id, indicating success.
self.assertIs(self.connection_service.owner_to_connection_map[params.owner_uri].get_connection(params.type),
mock_connection)
self.assertIsNotNone(response.connection_id)
self.assertIsNotNone(response.server_info.server_version)
self.assertFalse(response.server_info.is_cloud)
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_connection = MockConnection(dsn_parameters={
'host': 'myserver',
'dbname': '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
# 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_connection)) as mock_psycopg2_connect:
response = self.connection_service.connect(params)
mock_psycopg2_connect.assert_not_called()
mock_connection.close.assert_not_called()
self.assertIsNotNone(response.connection_id)
def setUp(self):
"""Set up the tests with common connection parameters"""
# Set up the mock connection service and connection info
self.connection_service = ConnectionService()
self.connection_service._service_provider = {
constants.WORKSPACE_SERVICE_NAME: WorkspaceService()
}
self.owner_uri = 'test_uri'
self.connection_type = ConnectionType.DEFAULT
self.connect_params: ConnectRequestParams = ConnectRequestParams.from_dict(
{
'ownerUri': self.owner_uri,
'type': self.connection_type,
'connection': {
'options': {}
}
})
self.mock_connection = MockConnection(dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
})
# Mock psycopg2's connect method to store the current cancellation token. This lets us
# capture the cancellation token state as it would be during a long-running connection.
self.token_store = []
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 = MockConnection(dsn_parameters={
'host': 'myserver1',
'dbname': '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 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 = MockConnection(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 = MockConnection(
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
with mock.patch('psycopg2.connect', 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_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 = MockConnection(dsn_parameters={
'host': 'myserver',
'dbname': '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 setUp(self):
self._smo_metadata_factory = SmoEditTableMetadataFactory()
self._connection = MockConnection({
"port": "8080",
"host": "test",
"dbname": "test",
"user": "******"
})
self._server = Server(self._connection)
self._schema_name = 'public'
self._table_name = 'Employee'
self._view_name = 'Vendor'
self._table_object_type = 'TABLE'
self._view_object_type = 'VIEW'
self._columns = [
Column(self._server, "testTable", 'testName', 'testDatatype')
]
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_connection = MockConnection(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_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()
def test_connect_with_access_token(self):
"""Test that the service connects to a PostgreSQL server using an access token as a password"""
# Set up the parameters for the connection
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri': 'someUri',
'type': ConnectionType.DEFAULT,
'connection': {
'options': {
'user': '******',
'azureAccountToken': 'exampleToken',
'host': 'myserver',
'dbname': 'postgres'
}
}
})
# Set up the mock connection for psycopg2's connect method to return
mock_connection = MockConnection(dsn_parameters={
'host': 'myserver',
'dbname': 'postgres',
'user': '******'
})
# Set up psycopg2 instance for connection service to call
mock_connect_method = mock.Mock(return_value=mock_connection)
# Set up the connection service and call its connect method with the supported options
with mock.patch('psycopg2.connect', new=mock_connect_method):
response = self.connection_service.connect(params)
# Verify that psycopg2's connection method was called with password set to account token.
mock_connect_method.assert_called_once_with(user='******', password='******', host='myserver', dbname='postgres')
# Verify that psycopg2's connection method was called and that the
# response has a connection id, indicating success.
self.assertIs(self.connection_service.owner_to_connection_map[params.owner_uri].get_connection(params.type),
mock_connection)
self.assertIsNotNone(response.connection_id)
self.assertIsNotNone(response.server_info.server_version)
self.assertFalse(response.server_info.is_cloud)
def setUp(self):
self._metadata_factory = mock.MagicMock()
self._mock_cursor = MockCursor(None)
self._connection = MockConnection(
{
"port": "8080",
"host": "test",
"dbname": "test"
}, self._mock_cursor)
self._initialize_edit_request = InitializeEditParams()
self._initialize_edit_request.schema_name = 'public'
self._initialize_edit_request.object_name = 'Employee'
self._initialize_edit_request.object_type = 'Table'
db_column = DbColumn()
column = EditColumnMetadata(db_column, None)
self._columns_metadata = [column]
self._schema_name = 'public'
self._table_name = 'table'
self._edit_table_metadata = EditTableMetadata(self._schema_name,
self._table_name,
self._columns_metadata)
self._query_executer = mock.MagicMock()
self._on_success = mock.MagicMock()
self._on_failure = mock.MagicMock()
self._data_editor_session = DataEditorSession(self._metadata_factory)
self._metadata_factory.get = mock.Mock(
return_value=self._edit_table_metadata)
self._query = 'SELECT TESTCOLUMN FROM TESTTABLE LIMIT 100'
self._data_editor_session._construct_initialize_query = mock.Mock(
return_value=self._query)
def test_metadata_list_request(self):
"""Test that the metadata list handler properly starts a thread to list metadata and responds with the list"""
# Set up the parameters and mocks for the request
expected_metadata = [
ObjectMetadata(schema='schema1',
name='table1',
metadata_type=MetadataType.TABLE),
ObjectMetadata(schema='schema1',
name='view1',
metadata_type=MetadataType.VIEW),
ObjectMetadata(schema='schema1',
name='function1',
metadata_type=MetadataType.FUNCTION),
ObjectMetadata(schema='schema1',
name='table2',
metadata_type=MetadataType.TABLE),
ObjectMetadata(schema='schema2',
name='view1',
metadata_type=MetadataType.VIEW),
ObjectMetadata(schema='schema2',
name='function1',
metadata_type=MetadataType.FUNCTION),
]
metadata_type_to_str_map = {
MetadataType.TABLE: 't',
MetadataType.VIEW: 'v',
MetadataType.FUNCTION: 'f'
}
# Query results have schema_name, object_name, and object_type columns in that order
list_query_result = [(metadata.schema, metadata.name,
metadata_type_to_str_map[metadata.metadata_type])
for metadata in expected_metadata]
mock_cursor = MockCursor(list_query_result)
mock_connection = MockConnection(cursor=mock_cursor)
self.connection_service.get_connection = mock.Mock(
return_value=mock_connection)
request_context = MockRequestContext()
params = MetadataListParameters()
params.owner_uri = self.test_uri
mock_thread = MockThread()
with mock.patch(
'threading.Thread',
new=mock.Mock(side_effect=mock_thread.initialize_target)):
# If I call the metadata list request handler
self.metadata_service._handle_metadata_list_request(
request_context, params)
# Then the worker thread was kicked off
self.assertEqual(mock_thread.target,
self.metadata_service._metadata_list_worker)
mock_thread.start.assert_called_once()
# And the worker retrieved the correct connection and executed a query on it
self.connection_service.get_connection.assert_called_once_with(
self.test_uri, ConnectionType.DEFAULT)
mock_cursor.execute.assert_called_once()
# And the handler responded with the expected results
self.assertIsNone(request_context.last_error_message)
self.assertIsNone(request_context.last_notification_method)
response = request_context.last_response_params
self.assertIsInstance(response, MetadataListResponse)
for index, actual_metadata in enumerate(response.metadata):
self.assertIsInstance(actual_metadata, ObjectMetadata)
self.assertEqual(actual_metadata.schema,
expected_metadata[index].schema)
self.assertEqual(actual_metadata.name,
expected_metadata[index].name)
self.assertEqual(actual_metadata.metadata_type,
expected_metadata[index].metadata_type)