All rows flag for obj plugin (#504)

Co-authored-by: Youjung Kim <[email protected]>

Author: zliang-akamai

linodecli/arg_helpers.py

@@ -16,7 +16,7 @@
 from linodecli import plugins
 
 from .completion import bake_completions
-from .helpers import register_args_shared
+from .helpers import pagination_args_shared, register_args_shared
 
 
 def register_args(parser):
@@ -76,21 +76,6 @@ def register_args(parser):
         action="store_true",
         help="If set, does not display headers in output.",
     )
-    parser.add_argument(
-        "--page",
-        metavar="PAGE",
-        default=1,
-        type=int,
-        help="For listing actions, specifies the page to request",
-    )
-    parser.add_argument(
-        "--page-size",
-        metavar="PAGESIZE",
-        default=100,
-        type=int,
-        help="For listing actions, specifies the number of items per page, "
-        "accepts any value between 25 and 500",
-    )
     parser.add_argument(
         "--all",
         action="store_true",
@@ -148,6 +133,7 @@ def register_args(parser):
         "--debug", action="store_true", help="Enable verbose HTTP debug output."
     )
 
+    pagination_args_shared(parser)
     register_args_shared(parser)
 
     return parser

linodecli/helpers.py

@@ -63,6 +63,33 @@ def filter_markdown_links(text):
     return result
 
 
+def pagination_args_shared(parser: ArgumentParser):
+    """
+    Add pagination related arguments to the given
+    ArgumentParser that may be shared across the CLI and plugins.
+    """
+    parser.add_argument(
+        "--page",
+        metavar="PAGE",
+        default=1,
+        type=int,
+        help="For listing actions, specifies the page to request",
+    )
+    parser.add_argument(
+        "--page-size",
+        metavar="PAGESIZE",
+        default=100,
+        type=int,
+        help="For listing actions, specifies the number of items per page, "
+        "accepts any value between 25 and 500",
+    )
+    parser.add_argument(
+        "--all-rows",
+        action="store_true",
+        help="Output all possible rows in the results with pagination",
+    )
+
+
 def register_args_shared(parser: ArgumentParser):
     """
     Adds certain arguments to the given ArgumentParser that may be shared across
@@ -87,12 +114,6 @@ def register_args_shared(parser: ArgumentParser):
         "This is useful for scripting the CLI's behavior.",
     )
 
-    parser.add_argument(
-        "--all-rows",
-        action="store_true",
-        help="Output all possible rows in the results with pagination",
-    )
-
     return parser
 
 

linodecli/plugins/obj/__init__.py

@@ -14,16 +14,15 @@
 from contextlib import suppress
 from datetime import datetime
 from math import ceil
-from pathlib import Path
-from typing import List
+from typing import Iterable, List
 
 from rich import print as rprint
 from rich.table import Table
 
 from linodecli.cli import CLI
 from linodecli.configuration import _do_get_request
 from linodecli.configuration.helpers import _default_thing_input
-from linodecli.helpers import expand_globs
+from linodecli.helpers import expand_globs, pagination_args_shared
 from linodecli.plugins import PluginContext, inherit_plugin_args
 from linodecli.plugins.obj.buckets import create_bucket, delete_bucket
 from linodecli.plugins.obj.config import (
@@ -70,14 +69,36 @@
 except ImportError:
     HAS_BOTO = False
 
+TRUNCATED_MSG = (
+    "Notice: Not all results were shown. If your would "
+    "like to get more results, you can add the '--all-row' "
+    "flag to the command or use the built-in pagination flags."
+)
+
+INVALID_PAGE_MSG = "No result to show in this page."
+
+
+def flip_to_page(iterable: Iterable, page: int = 1):
+    """Given a iterable object and return a specific iteration (page)"""
+    iterable = iter(iterable)
+    for _ in range(page - 1):
+        try:
+            next(iterable)
+        except StopIteration:
+            print(INVALID_PAGE_MSG)
+            sys.exit(2)
+
+    return next(iterable)
+
 
 def list_objects_or_buckets(
     get_client, args, **kwargs
-):  # pylint: disable=too-many-locals,unused-argument
+):  # pylint: disable=too-many-locals,unused-argument,too-many-branches
     """
     Lists buckets or objects
     """
     parser = inherit_plugin_args(ArgumentParser(PLUGIN_BASE + " ls"))
+    pagination_args_shared(parser)
 
     parser.add_argument(
         "bucket",
@@ -106,16 +127,30 @@ def list_objects_or_buckets(
             prefix = ""
 
         data = []
+        objects = []
+        sub_directories = []
+        pages = client.get_paginator("list_objects_v2").paginate(
+            Prefix=prefix,
+            Bucket=bucket_name,
+            Delimiter="/",
+            PaginationConfig={"PageSize": parsed.page_size},
+        )
         try:
-            response = client.list_objects_v2(
-                Prefix=prefix, Bucket=bucket_name, Delimiter="/"
-            )
+            if parsed.all_rows:
+                results = pages
+            else:
+                page = flip_to_page(pages, parsed.page)
+                if page.get("IsTruncated", False):
+                    print(TRUNCATED_MSG)
+
+                results = [page]
         except client.exceptions.NoSuchBucket:
             print("No bucket named " + bucket_name)
             sys.exit(2)
 
-        objects = response.get("Contents", [])
-        sub_directories = response.get("CommonPrefixes", [])
+        for item in results:
+            objects.extend(item.get("Contents", []))
+            sub_directories.extend(item.get("CommonPrefixes", []))
 
         for d in sub_directories:
             data.append((" " * 16, "DIR", d.get("Prefix")))
@@ -329,17 +364,31 @@ def list_all_objects(
     """
     # this is for printing help when --help is in the args
     parser = inherit_plugin_args(ArgumentParser(PLUGIN_BASE + " la"))
+    pagination_args_shared(parser)
 
-    parser.parse_args(args)
+    parsed = parser.parse_args(args)
 
     client = get_client()
 
-    # all buckets
     buckets = [b["Name"] for b in client.list_buckets().get("Buckets", [])]
 
     for b in buckets:
         print()
-        objects = client.list_objects_v2(Bucket=b).get("Contents", [])
+        objects = []
+        pages = client.get_paginator("list_objects_v2").paginate(
+            Bucket=b, PaginationConfig={"PageSize": parsed.page_size}
+        )
+        if parsed.all_rows:
+            results = pages
+        else:
+            page = flip_to_page(pages, parsed.page)
+            if page.get("IsTruncated", False):
+                print(TRUNCATED_MSG)
+
+            results = [page]
+
+        for page in results:
+            objects.extend(page.get("Contents", []))
 
         for obj in objects:
             size = obj.get("Size", 0)

tests/integration/helpers.py

@@ -159,3 +159,7 @@ def exec_failing_test_command(args: List[str]):
     )
     assert process.returncode == 1
     return process
+
+
+def count_lines(text: str):
+    return len(list(filter(len, text.split("\n"))))

tests/integration/obj/test_obj_plugin.py

@@ -1,16 +1,19 @@
 import logging
-import subprocess
 from dataclasses import dataclass
-from typing import Callable, List, Optional
+from typing import Callable, Optional
 
 import pytest
 import requests
 from pytest import MonkeyPatch
 
 from linodecli.configuration.auth import _do_request
-from linodecli.plugins.obj import ENV_ACCESS_KEY_NAME, ENV_SECRET_KEY_NAME
+from linodecli.plugins.obj import (
+    ENV_ACCESS_KEY_NAME,
+    ENV_SECRET_KEY_NAME,
+    TRUNCATED_MSG,
+)
 from tests.integration.fixture_types import GetTestFilesType, GetTestFileType
-from tests.integration.helpers import BASE_URL
+from tests.integration.helpers import BASE_URL, count_lines, exec_test_command
 
 REGION = "us-southeast-1"
 BASE_CMD = ["linode-cli", "obj", "--cluster", REGION]
@@ -77,15 +80,6 @@ def patch_keys(keys: Keys, monkeypatch: MonkeyPatch):
     monkeypatch.setenv(ENV_SECRET_KEY_NAME, keys.secret_key)
 
 
-def exec_test_command(args: List[str]):
-    process = subprocess.run(
-        args,
-        stdout=subprocess.PIPE,
-    )
-    assert process.returncode == 0
-    return process
-
-
 @pytest.fixture
 def create_bucket(
     name_generator: Callable, keys: Keys, monkeypatch: MonkeyPatch
@@ -176,6 +170,49 @@ def test_multi_files_multi_bucket(
             assert "Done" in output
 
 
+def test_all_rows(
+    create_bucket: Callable[[Optional[str]], str],
+    generate_test_files: GetTestFilesType,
+    keys: Keys,
+    monkeypatch: MonkeyPatch,
+):
+    patch_keys(keys, monkeypatch)
+    number = 5
+    bucket_name = create_bucket()
+    file_paths = generate_test_files(number)
+
+    process = exec_test_command(
+        BASE_CMD
+        + ["put"]
+        + [str(file.resolve()) for file in file_paths]
+        + [bucket_name]
+    )
+    output = process.stdout.decode()
+    assert "100.0%" in output
+    assert "Done" in output
+
+    process = exec_test_command(
+        BASE_CMD + ["ls", bucket_name, "--page-size", "2", "--page", "1"]
+    )
+    output = process.stdout.decode()
+    assert TRUNCATED_MSG in output
+    assert count_lines(output) == 3
+
+    process = exec_test_command(
+        BASE_CMD + ["ls", bucket_name, "--page-size", "999"]
+    )
+    output = process.stdout.decode()
+    assert TRUNCATED_MSG not in output
+    assert count_lines(output) == 5
+
+    process = exec_test_command(
+        BASE_CMD + ["ls", bucket_name, "--page-size", "2", "--all-rows"]
+    )
+    output = process.stdout.decode()
+    assert TRUNCATED_MSG not in output
+    assert count_lines(output) == 5
+
+
 def test_modify_access_control(
     keys: Keys,
     monkeypatch: MonkeyPatch,
@@ -273,7 +310,7 @@ def test_show_usage(
 
     process = exec_test_command(BASE_CMD + ["du"])
     output = process.stdout.decode()
-    assert "40.0 MB Total" in output
+    assert "MB Total" in output
 
     process = exec_test_command(BASE_CMD + ["du", bucket1])
     output = process.stdout.decode()

tests/unit/test_helpers.py

@@ -1,4 +1,10 @@
-from linodecli.helpers import filter_markdown_links
+from argparse import ArgumentParser
+
+from linodecli.helpers import (
+    filter_markdown_links,
+    pagination_args_shared,
+    register_args_shared,
+)
 
 
 class TestHelpers:
@@ -14,3 +20,24 @@ def test_markdown_links(self):
         )
 
         assert filter_markdown_links(original_text) == expected_text
+
+    def test_pagination_args_shared(self):
+        parser = ArgumentParser()
+        pagination_args_shared(parser)
+
+        args = parser.parse_args(
+            ["--page", "2", "--page-size", "50", "--all-rows"]
+        )
+        assert args.page == 2
+        assert args.page_size == 50
+        assert args.all_rows
+
+    def test_register_args_shared(self):
+        parser = ArgumentParser()
+        register_args_shared(parser)
+
+        args = parser.parse_args(
+            ["--as-user", "linode-user", "--suppress-warnings"]
+        )
+        assert args.as_user == "linode-user"
+        assert args.suppress_warnings