paginate.py

# Copyright 2013 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""This module has customizations to unify paging parameters.

For any operation that can be paginated, we will:

    * Hide the service specific pagination params.  This can vary across
    services and we're going to replace them with a consistent set of
    arguments.  The arguments will still work, but they are not
    documented.  This allows us to add a pagination config after
    the fact and still remain backwards compatible with users that
    were manually doing pagination.
    * Add a ``--starting-token`` and a ``--max-items`` argument.

"""
import logging
import sys
from functools import partial

from awscli.customizations.utils import uni_print
from botocore import xform_name
from botocore.exceptions import DataNotFoundError, PaginationError
from botocore import model

from awscli.arguments import BaseCLIArgument
from awscli.utils import resolve_v2_debug_mode

logger = logging.getLogger(__name__)


STARTING_TOKEN_HELP = """
<p>A token to specify where to start paginating.  This is the
<code>NextToken</code> from a previously truncated response.</p>
<p>For usage examples, see <a
href="https://docs.aws.amazon.com/cli/latest/userguide/pagination.html"
>Pagination</a> in the <i>AWS Command Line Interface User
Guide</i>.</p>
"""

MAX_ITEMS_HELP = """
<p>The total number of items to return in the command's output.
If the total number of items available is more than the value
specified, a <code>NextToken</code> is provided in the command's
output.  To resume pagination, provide the
<code>NextToken</code> value in the <code>starting-token</code>
argument of a subsequent command.  <b>Do not</b> use the
<code>NextToken</code> response element directly outside of the
AWS CLI.</p>
<p>For usage examples, see <a
href="https://docs.aws.amazon.com/cli/latest/userguide/pagination.html"
>Pagination</a> in the <i>AWS Command Line Interface User
Guide</i>.</p>
"""

PAGE_SIZE_HELP = """
<p>The size of each page to get in the AWS service call.  This
does not affect the number of items returned in the command's
output.  Setting a smaller page size results in more calls to
the AWS service, retrieving fewer items in each call.  This can
help prevent the AWS service calls from timing out.</p>
<p>For usage examples, see <a
href="https://docs.aws.amazon.com/cli/latest/userguide/pagination.html"
>Pagination</a> in the <i>AWS Command Line Interface User
Guide</i>.</p>
"""


def register_pagination(event_handlers):
    event_handlers.register('building-argument-table', unify_paging_params)
    event_handlers.register_last('doc-description', add_paging_description)


def get_paginator_config(session, service_name, operation_name):
    try:
        paginator_model = session.get_paginator_model(service_name)
    except DataNotFoundError:
        return None
    try:
        operation_paginator_config = paginator_model.get_paginator(
            operation_name)
    except ValueError:
        return None
    return operation_paginator_config


def add_paging_description(help_command, **kwargs):
    # This customization is only applied to the description of
    # Operations, so we must filter out all other events.
    if not isinstance(help_command.obj, model.OperationModel):
        return
    service_name = help_command.obj.service_model.service_name
    paginator_config = get_paginator_config(
        help_command.session, service_name, help_command.obj.name)
    if not paginator_config:
        return
    help_command.doc.style.new_paragraph()
    help_command.doc.writeln(
        ('``%s`` is a paginated operation. Multiple API calls may be issued '
         'in order to retrieve the entire data set of results. You can '
         'disable pagination by providing the ``--no-paginate`` argument.')
        % help_command.name)
    # Only include result key information if it is present.
    if paginator_config.get('result_key'):
        queries = paginator_config['result_key']
        if type(queries) is not list:
            queries = [queries]
        queries = ", ".join([('``%s``' % s) for s in queries])
        help_command.doc.writeln(
            ('When using ``--output text`` and the ``--query`` argument on a '
             'paginated response, the ``--query`` argument must extract data '
             'from the results of the following query expressions: %s')
            % queries)


def unify_paging_params(argument_table, operation_model, event_name,
                        session, **kwargs):
    paginator_config = get_paginator_config(
        session, operation_model.service_model.service_name,
        operation_model.name)
    if paginator_config is None:
        # We only apply these customizations to paginated responses.
        return
    logger.debug("Modifying paging parameters for operation: %s",
                 operation_model.name)
    _remove_existing_paging_arguments(argument_table, paginator_config)
    parsed_args_event = event_name.replace('building-argument-table.',
                                           'operation-args-parsed.')
    call_parameters_event = event_name.replace(
        'building-argument-table', 'calling-command'
    )
    shadowed_args = {}
    add_paging_argument(argument_table, 'starting-token',
                        PageArgument('starting-token', STARTING_TOKEN_HELP,
                                     parse_type='string',
                                     serialized_name='StartingToken'),
                        shadowed_args)
    input_members = operation_model.input_shape.members
    type_name = 'integer'
    if 'limit_key' in paginator_config:
        limit_key_shape = input_members[paginator_config['limit_key']]
        type_name = limit_key_shape.type_name
        if type_name not in PageArgument.type_map:
            raise TypeError(
                ('Unsupported pagination type {0} for operation {1}'
                 ' and parameter {2}').format(
                    type_name, operation_model.name,
                    paginator_config['limit_key']))
        add_paging_argument(argument_table, 'page-size',
                            PageArgument('page-size', PAGE_SIZE_HELP,
                                         parse_type=type_name,
                                         serialized_name='PageSize'),
                            shadowed_args)

    add_paging_argument(argument_table, 'max-items',
                        PageArgument('max-items', MAX_ITEMS_HELP,
                                     parse_type=type_name,
                                     serialized_name='MaxItems'),
                        shadowed_args)
    session.register(
        parsed_args_event,
        partial(check_should_enable_pagination,
                list(_get_all_cli_input_tokens(paginator_config)),
                shadowed_args, argument_table))
    session.register(
        call_parameters_event,
        partial(
            check_should_enable_pagination_call_parameters,
            session,
            list(_get_all_input_tokens(paginator_config)),
        ),
    )


def add_paging_argument(argument_table, arg_name, argument, shadowed_args):
    if arg_name in argument_table:
        # If there's already an entry in the arg table for this argument,
        # this means we're shadowing an argument for this operation.  We
        # need to store this later in case pagination is turned off because
        # we put these arguments back.
        # See the comment in check_should_enable_pagination() for more info.
        shadowed_args[arg_name] = argument_table[arg_name]
    argument_table[arg_name] = argument


def check_should_enable_pagination(input_tokens, shadowed_args, argument_table,
                                   parsed_args, parsed_globals, **kwargs):
    normalized_paging_args = ['start_token', 'max_items']
    for token in input_tokens:
        py_name = token.replace('-', '_')
        if getattr(parsed_args, py_name) is not None and \
                py_name not in normalized_paging_args:
            # The user has specified a manual (undocumented) pagination arg.
            # We need to automatically turn pagination off.
            logger.debug("User has specified a manual pagination arg. "
                         "Automatically setting --no-paginate.")
            parsed_globals.paginate = False

    if not parsed_globals.paginate:
        ensure_paging_params_not_set(parsed_args, shadowed_args)
        # Because pagination is now disabled, there's a chance that
        # we were shadowing arguments.  For example, we inject a
        # --max-items argument in unify_paging_params().  If the
        # the operation also provides its own MaxItems (which we
        # expose as --max-items) then our custom pagination arg
        # was shadowing the customers arg.  When we turn pagination
        # off we need to put back the original argument which is
        # what we're doing here.
        for key, value in shadowed_args.items():
            argument_table[key] = value


def ensure_paging_params_not_set(parsed_args, shadowed_args):
    paging_params = ['starting_token', 'page_size', 'max_items']
    shadowed_params = [p.replace('-', '_') for p in shadowed_args.keys()]
    params_used = [p for p in paging_params if
                   p not in shadowed_params and getattr(parsed_args, p, None)]

    if len(params_used) > 0:
        converted_params = ', '.join(
            ["--" + p.replace('_', '-') for p in params_used])
        raise PaginationError(
            message="Cannot specify --no-paginate along with pagination "
                    "arguments: %s" % converted_params)


def _remove_existing_paging_arguments(argument_table, pagination_config):
    for cli_name in _get_all_cli_input_tokens(pagination_config):
        argument_table[cli_name]._UNDOCUMENTED = True


def _get_all_cli_input_tokens(pagination_config):
    # Get all input tokens including the limit_key
    # if it exists.
    tokens = _get_input_tokens(pagination_config)
    for token_name in tokens:
        cli_name = xform_name(token_name, '-')
        yield cli_name
    if 'limit_key' in pagination_config:
        key_name = pagination_config['limit_key']
        cli_name = xform_name(key_name, '-')
        yield cli_name


# Get all tokens but return them in API namespace rather than CLI namespace
def _get_all_input_tokens(pagination_config):
    # Get all input tokens including the limit_key
    # if it exists.
    tokens = _get_input_tokens(pagination_config)
    for token_name in tokens:
        yield token_name
    if 'limit_key' in pagination_config:
        key_name = pagination_config['limit_key']
        yield key_name


def _get_input_tokens(pagination_config):
    tokens = pagination_config['input_token']
    if not isinstance(tokens, list):
        return [tokens]
    return tokens


def _get_cli_name(param_objects, token_name):
    for param in param_objects:
        if param.name == token_name:
            return param.cli_name.lstrip('-')


def check_should_enable_pagination_call_parameters(
        session,
        input_tokens,
        call_parameters,
        parsed_args,
        parsed_globals,
        **kwargs
):
    """
    Check for pagination args in the actual calling arguments passed to
    the function.

    If the user is using the --cli-input-json parameter to provide JSON
    parameters they are all in the API naming space rather than the CLI
    naming space and would be missed by the processing above. This function
    gets called on the calling-command event.
    """
    if resolve_v2_debug_mode(parsed_globals):
        cli_input_json_data = session.emit_first_non_none_response(
            f"get-cli-input-json-data",
        )
        if cli_input_json_data is None:
            cli_input_json_data = {}
        pagination_params_in_input_tokens = [
            param for param in cli_input_json_data if param in input_tokens
        ]
        if pagination_params_in_input_tokens:
            uni_print(
                '\nAWS CLI v2 UPGRADE WARNING: In AWS CLI v2, if you specify '
                'pagination parameters by using a file with the '
                '`--cli-input-json` parameter, automatic pagination will be '
                'turned off. This is different from v1 behavior, where '
                'pagination parameters specified via the `--cli-input-json` '
                'parameter are ignored. To retain AWS CLI v1 behavior in '
                'AWS CLI v2, remove all pagination parameters from the input '
                'JSON. See https://docs.aws.amazon.com/cli/latest/userguide/'
                'cliv2-migration-changes.html'
                '#cliv2-migration-skeleton-paging.\n',
                out_file=sys.stderr
            )


class PageArgument(BaseCLIArgument):
    type_map = {
        'string': str,
        'integer': int,
        'long': int,
    }

    def __init__(self, name, documentation, parse_type, serialized_name):
        self.argument_model = model.Shape('PageArgument', {'type': 'string'})
        self._name = name
        self._serialized_name = serialized_name
        self._documentation = documentation
        self._parse_type = parse_type
        self._required = False

    def _emit_non_positive_max_items_warning(self):
        uni_print(
            "warning: Non-positive values for --max-items may result in undefined behavior.\n",
            sys.stderr)

    @property
    def cli_name(self):
        return '--' + self._name

    @property
    def cli_type_name(self):
        return self._parse_type

    @property
    def required(self):
        return self._required

    @required.setter
    def required(self, value):
        self._required = value

    @property
    def documentation(self):
        return self._documentation

    def add_to_parser(self, parser):
        parser.add_argument(self.cli_name, dest=self.py_name,
                            type=self.type_map[self._parse_type])

    def add_to_params(self, parameters, value):
        if value is not None:
            if self._serialized_name == 'MaxItems' and int(value) <= 0:
                self._emit_non_positive_max_items_warning()
            pagination_config = parameters.get('PaginationConfig', {})
            pagination_config[self._serialized_name] = value
            parameters['PaginationConfig'] = pagination_config