auto pagination (#465)

hca/cli.py

@@ -144,7 +144,7 @@ def main(args=None):
             try:
                 err_log_filename = os.path.join(get_config().user_config_dir, "error.log")
                 with open(err_log_filename, "ab") as fh:
-                    print(datetime.datetime.now().isoformat(), file=fh)
+                    print(datetime.datetime.now().isoformat(), file=fh)  # noqa
                     print(err_msg, file=fh)
                 exit("{}: {}. See {} for error details.".format(e.__class__.__name__, e, err_log_filename))
             except Exception:

hca/util/__init__.py

@@ -92,6 +92,7 @@ def special_cli_command(self, required_argument, optional_argument=""):
 import time
 import jwt
 import requests
+import jmespath
 
 from inspect import signature, Parameter
 from requests.adapters import HTTPAdapter, DEFAULT_POOLSIZE
@@ -200,6 +201,27 @@ def paginate(self, **kwargs):
         for page in self._get_raw_pages(**kwargs):
             yield page.json()
 
+    def _cli_call(self, cli_args):
+        if cli_args.paginate is not True:
+            return super()._cli_call(cli_args)
+        return self._auto_page(**vars(cli_args))
+
+    def _auto_page(self, **kwargs):
+        '''This method allows for autopaging in commands and bindings'''
+        response_data = None
+        for page in self._get_raw_pages(**kwargs):
+            page_data = page.json()
+            content_key = page.headers.get("X-OpenAPI-Paginated-Content-Key", "results")
+            if response_data is None:
+                response_data = page_data
+            else:
+                data_aggregrator = jmespath.search(content_key, response_data)
+                patch = jmespath.search(content_key, page_data)
+                if patch:
+                    data_aggregrator += patch
+                response_data[content_key] = data_aggregrator
+        return response_data
+
 
 class SwaggerClient(object):
     scheme = "https"
@@ -598,6 +620,9 @@ def build_argparse_subparsers(self, subparsers, help_menu=False):
                                        help=method_data["args"][param_name]["doc"],
                                        choices=method_data["args"][param_name]["choices"],
                                        required=method_data["args"][param_name]["required"])
+            if str(requests.codes.partial) in method_data["responses"]:
+                subparser.add_argument("--no-paginate", action="store_false", dest="paginate",
+                                       help='Do not automatically page the responses', default=True)
             subparser.set_defaults(entry_point=method_data["entry_point"])
 
         for command in self.commands:

test/integration/dss/test_dss_cli.py

@@ -2,12 +2,16 @@
 # coding: utf-8
 import contextlib
 import filecmp
+import io
 import json
 import os
 import sys
 import unittest
+import unittest.mock
 import uuid
 import tempfile
+import requests
+from requests.models import Response
 
 pkg_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))  # noqa
 sys.path.insert(0, pkg_root)  # noqa
@@ -23,7 +27,8 @@ class TestDssCLI(unittest.TestCase):
     def test_post_search_cli(self):
         query = json.dumps({})
         replica = "aws"
-        args = ["dss", "post-search", "--es-query", query, "--replica", replica, "--output-format", "raw"]
+        args = ["dss", "post-search", "--es-query", query, "--replica", replica,
+                "--output-format", "raw", "--no-paginate"]
         with CapturingIO('stdout') as stdout:
             hca.cli.main(args)
         result = json.loads(stdout.captured())
@@ -247,5 +252,45 @@ def test_collection_download(self):
                 self.assertFalse(diff.diff_files)
 
 
+class SessionMock:
+    def __init__(self):
+        self.pages = [{"results": ["result"] * i} for i in range(3)]
+
+    def request(self, *args, **kwargs):
+        res = Response()
+        res.status_code = requests.codes.ok
+        page = self.pages.pop()
+        res.headers = {"link": '<https:///>; rel="next"'} if self.pages else {}
+        res.headers['content-type'] = 'application/json'
+        res.headers['X-OpenAPI-Paginated-Content-Key'] = 'results'
+        res.raw = io.BytesIO(json.dumps(page).encode())
+        return res
+
+    def get(self, *args, **kwargs):
+        kwargs["method"] = "GET"
+        return self.request(*args, **kwargs)
+
+
+class MockTestCase(unittest.TestCase):
+    def _test_with_mock(self, expected_length, no_pagination=False):
+        paged_args = ['dss', 'get-bundles-all', '--replica', 'aws']
+        if no_pagination is True:
+            paged_args.append('--no-paginate')
+        with CapturingIO('stdout') as stdout:
+            hca.cli.main(args=paged_args)
+        output = json.loads(stdout.captured())
+        self.assertEqual(expected_length,len(output['results']))
+
+    @unittest.mock.patch("hca.util._ClientMethodFactory._request")
+    def test_no_page(self, request_mock):
+        request_mock.side_effect = SessionMock().request
+        self._test_with_mock(2, no_pagination=True)
+
+    @unittest.mock.patch("hca.util._ClientMethodFactory._request")
+    def test_page(self, request_mock):
+        request_mock.side_effect = SessionMock().request
+        self._test_with_mock(3, no_pagination=False)
+
+
 if __name__ == "__main__":
     unittest.main()