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)
Example #4
0
    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)