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_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 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 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_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[constants.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('pgsqltoolsservice.connection.connection_service._build_connection_response'), \
                mock.patch('psycopg2.connect') 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 get_connection(self, owner_uri: str, connection_type: ConnectionType) -> Optional[psycopg2.extensions.connection]:
        """
        Get a psycopg2 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)
Example #8
0
    def _create_connection(
            self, connection_key: str, conn_info: ConnectionInfo
    ) -> Optional[psycopg2.extensions.connection]:
        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 _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 PGSMO Server object for the session and create the root node for the server
            session.server = 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.options['dbname']
            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 _create_connection(self, session: ObjectExplorerSession, database_name: str) -> Optional[psycopg2.extensions.connection]:
        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 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 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()
Example #13
0
    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 test_response_when_connect_fails(self):
     """Test that the proper response is given when a connection fails"""
     error_message = 'some error'
     params: ConnectRequestParams = ConnectRequestParams.from_dict({
         'ownerUri': 'someUri',
         'type': ConnectionType.DEFAULT,
         'connection': {
             'options': {
                 'connectionString': ''
             }
         }
     })
     with mock.patch('psycopg2.connect', new=mock.Mock(side_effect=Exception(error_message))):
         response = self.connection_service.connect(params)
     # The response should not have a connection ID and should contain the error message
     self.assertIsNone(response.connection_id)
     self.assertEqual(response.error_message, error_message)
    def test_handle_connect_request(self):
        """Test that the handle_connect_request method kicks off a new thread to do the connection"""
        # Setup: Create a mock request context to handle output
        rc = utils.MockRequestContext()
        connect_response = ConnectionCompleteParams()
        self.connection_service.connect = Mock(return_value=connect_response)

        # If: I make a request to connect
        params: ConnectRequestParams = ConnectRequestParams.from_dict({
            'ownerUri':
            'someUri',
            'type':
            ConnectionType.QUERY,
            'connection': {
                'server_name': 'someserver',
                'user_name': 'someuser',
                'database_name': 'somedb',
                'options': {
                    'password': '******'
                }
            }
        })

        # Connect and wait for the thread to finish executing, then verify the connection information
        self.connection_service.handle_connect_request(rc, params)
        connection_thread = self.connection_service.owner_to_thread_map[
            params.owner_uri]
        self.assertIsNotNone(connection_thread)
        connection_thread.join()

        # Then:
        # ... Connect should have been called once
        self.connection_service.connect.assert_called_once_with(params)

        # ... A True should have been sent as the response to the request
        rc.send_response.assert_called_once_with(True)

        # ... A connection complete notification should have been sent back as well
        rc.send_notification.assert_called_once_with(
            CONNECTION_COMPLETE_METHOD, connect_response)

        # ... An error should not have been called
        rc.send_error.assert_not_called()
    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)