def __init__(self,
cur: Optional[MockCursor] = None,
connection: Optional[MockPsycopgConnection] = None,
version: str = '90602',
name: str = 'postgres',
host: str = 'localhost',
port: str = '25565',
user: str = 'postgres'):
# Setup mocks for the connection
self.close = mock.MagicMock()
self.cursor = mock.MagicMock(return_value=cur)
# if no mock pyscopg connection passed, create default one
if not connection:
connection = MockPsycopgConnection(cursor=cur,
dsn_parameters={
'dbname': name,
'host': host,
'port': port,
'user': user
})
# mock psycopg2.connect call in PostgreSQLConnection.__init__ to return mock psycopg connection
with mock.patch('psycopg2.connect',
mock.Mock(return_value=connection)):
super().__init__({
"host_name": host,
"user_name": user,
"port": port,
"database_name": name
})
def test_execute_dict_success(self):
# Setup: Create a mock server connection that will return a result set
mock_cursor = utils.MockCursor(utils.get_mock_results())
mock_conn = MockPsycopgConnection(cursor=mock_cursor)
# noinspection PyTypeChecker
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_conn)):
server_conn = PostgreSQLConnection({})
# If: I execute a query as a dictionary
results = server_conn.execute_dict('SELECT * FROM pg_class')
# Then:
# ... Both the columns and the results should be returned
self.assertIsInstance(results, tuple)
self.assertEqual(len(results), 2)
# ... I should get a list of columns returned to me
cols = results[0]
self.assertIsInstance(cols, list)
self.assertListEqual(cols, mock_cursor.description)
# ... I should get the results formatted as a list of dictionaries
rows = results[1]
self.assertIsInstance(rows, list)
for idx, row in enumerate(rows):
self.assertDictEqual(row, mock_cursor._results[1][idx])
# ... The cursor should be closed
mock_cursor.close.assert_called_once()
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 = utils.get_mock_service_provider(
{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_psycopg_connection = MockPsycopgConnection(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_server_conn_init(self):
# Setup: Create a mock connection with an 'interesting' version
dsn_parameters = {'dbname': 'postgres', 'host': 'localhost', 'port': '25565', 'user': '******'}
mock_conn = MockPsycopgConnection(dsn_parameters=dsn_parameters)
mock_conn.server_version = '100216'
# If: I initialize a server connection
# noinspection PyTypeChecker
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_conn)):
server_conn = PostgreSQLConnection({})
# Then: The properties should be properly set
self.assertEqual(server_conn._conn, mock_conn)
self.assertEqual(server_conn.connection, mock_conn)
self.assertDictEqual(server_conn._dsn_parameters, dsn_parameters)
self.assertTupleEqual((10, 2, 16), server_conn.server_version)
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_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_non_default_database(self):
"""Test that if a database is given, the default database is not used"""
# Set up the connection params and default database name
default_db = 'test_db'
actual_db = 'postgres'
self.connection_service._service_provider[
WORKSPACE_SERVICE_NAME].configuration.pgsql.default_database = default_db
params: ConnectRequestParams = ConnectRequestParams.from_dict({
'ownerUri':
'someUri',
'type':
ConnectionType.DEFAULT,
'connection': {
'options': {
'user': '******',
'password': '******',
'host': 'myserver',
'dbname': actual_db
}
}
})
# If I connect with an empty database name
with mock.patch('ossdbtoolsservice.connection.connection_service._build_connection_response'), \
mock.patch('psycopg2.connect', return_value=MockPsycopgConnection()) as mock_psycopg2_connect:
self.connection_service.connect(params)
# Then psycopg2's connect method was called with the default database
calls = mock_psycopg2_connect.mock_calls
self.assertEqual(len(calls), 1)
self.assertNotEqual(calls[0][2]['dbname'], default_db)
self.assertEqual(calls[0][2]['dbname'], actual_db)
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 setUp(self):
self._metadata_factory = mock.MagicMock()
self._mock_cursor = MockCursor(None)
self._connection = MockPsycopgConnection({"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_execute_dict_fail(self):
# Setup: Create a mock psycopg connection that will raise an exception
mock_cursor = utils.MockCursor(None, throw_on_execute=True)
mock_conn = MockPsycopgConnection(cursor=mock_cursor)
# noinspection PyTypeChecker
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_conn)):
server_conn = PostgreSQLConnection({})
# If: I execute a query as a dictionary
# Then:
# ... I should get an exception
with self.assertRaises(Exception):
server_conn.execute_dict('SELECT * FROM pg_class')
# ... The cursor should be closed
mock_cursor.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_recovery_properties(self):
# Setup:
# NOTE: We're *not* mocking out the template rendering b/c this will verify that there's a template
# ... Create a mock query execution that will return the properties
mock_exec_dict = mock.MagicMock(return_value=([], [TestServer.CHECK_RECOVERY_ROW]))
# ... Create an instance of the class and override the connection
mock_connection = MockPsycopgConnection({'host': 'host', 'dbname': 'dbname'})
with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)):
pg_connection = PostgreSQLConnection({})
pg_connection.execute_dict = mock_exec_dict
obj = Server(pg_connection)
# If: I retrieve all the values in the recovery properties
# Then:
# ... The properties based on the properties should be availble
self.assertEqual(obj.in_recovery, TestServer.CHECK_RECOVERY_ROW['inrecovery'])
self.assertEqual(obj.wal_paused, TestServer.CHECK_RECOVERY_ROW['isreplaypaused'])
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()
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 = MockPsycopgConnection(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)._conn,
mock_connection)
self.assertIsNotNone(response.connection_id)
self.assertIsNotNone(response.server_info.server_version)
self.assertFalse(response.server_info.is_cloud)