Include headers in initial client load (#38)

* Include headers in initial client load

To do this, we had to copy over the `STACObject.from_file` method from
pystac, becuase that method did some URL absolutization before
using STAC_IO.read_text_method, which breaks our custom
urllib.request.Request handling. It's not a perfect solution -- it might
be better to make the pystac method smarter -- but it's a quick fix.

This commit also adds:
- Return codes for the cli
- A simple cli test
- An extra dev dependency (pytest-console-scripts) to test to cli

* Switch to a monkeypath

* Remove print

* Add headers to the stac_io logging

* add header logging to search

* Allow headers to be a file

Author: gadomski

pystac_client/cli.py

@@ -5,7 +5,6 @@
 import sys
 
 from .item_collection import ItemCollection
-
 from .client import Client
 from .version import __version__
 
@@ -35,6 +34,7 @@ def search(url=STAC_URL, matched=False, save=None, headers=None, **kwargs):
     except Exception as e:
         logger.error(e, exc_info=True)
         print(e)
+        return 1
 
 
 def parse_args(args):
@@ -51,7 +51,7 @@ def parse_args(args):
     parent.add_argument('--url', help='Root Catalog URL', default=os.getenv('STAC_URL', None))
     parent.add_argument('--limit', help='Page size limit', type=int, default=500)
     parent.add_argument('--headers',
-                        help='Additional request headers (JSON string)',
+                        help='Additional request headers (JSON string or file)',
                         default=None)
 
     subparsers = parser0.add_subparsers(dest='command')
@@ -108,7 +108,12 @@ def parse_args(args):
 
     # if headers provided, parse it
     if 'headers' in parsed_args:
-        parsed_args['headers'] = json.loads(parsed_args['headers'])
+        headers = parsed_args['headers']
+        if os.path.exists(headers):
+            with open(headers) as headers_file:
+                parsed_args['headers'] = json.load(headers_file)
+        else:
+            parsed_args['headers'] = json.loads(headers)
 
     return parsed_args
 
@@ -129,8 +134,10 @@ def cli():
 
     cmd = args.pop('command')
     if cmd == 'search':
-        search(**args)
+        return search(**args)
 
 
 if __name__ == "__main__":
-    cli()
+    return_code = cli()
+    if return_code and return_code != 0:
+        sys.exit(return_code)

pystac_client/client.py

@@ -1,9 +1,11 @@
 from copy import deepcopy
 from typing import Callable, Optional
+from urllib.request import Request
 
 import pystac
 import pystac.stac_object
 import pystac.validation
+from pystac import STAC_IO
 
 from pystac_client.conformance import ConformanceClasses
 from pystac_client.exceptions import ConformanceError
@@ -67,13 +69,13 @@ def __init__(self,
                 'API does not conform to {ConformanceClasses.STAC_API_CORE}. Must contain one of the following '
                 f'URIs to conform (preferably the first):\n\t{allowed_uris}.')
 
-        self.headers = headers
+        self.headers = headers or {}
 
     def __repr__(self):
         return '<Catalog id={}>'.format(self.id)
 
     @classmethod
-    def open(cls, url, headers={}):
+    def open(cls, url, headers=None):
         """Alias for PySTAC's STAC Object `from_file` method
 
         Parameters
@@ -85,7 +87,16 @@ def open(cls, url, headers={}):
         -------
         catalog : Client
         """
+        import pystac_client.stac_io
+
+        def read_text_method(url):
+            request = Request(url, headers=headers or {})
+            return pystac_client.stac_io.read_text_method(request)
+
+        old_read_text_method = STAC_IO.read_text_method
+        STAC_IO.read_text_method = read_text_method
         catalog = cls.from_file(url)
+        STAC_IO.read_text_method = old_read_text_method
         catalog.headers = headers
         return catalog
 

pystac_client/item_search.py

@@ -171,7 +171,7 @@ def __init__(
     ):
         self.conformance = conformance
         self.session = Session()
-        self.session.headers.update(headers)
+        self.session.headers.update(headers or {})
         self.request = Request(method=method, url=url)
 
         self._next_resolver = next_resolver or simple_stac_resolver

pystac_client/stac_io.py

@@ -18,7 +18,7 @@ def read_text_method(uri):
     instances as input. This method also raises any :exc:`urllib.error.HTTPError` exceptions rather than catching
     them to allow us to handle different response status codes as needed."""
     if isinstance(uri, Request):
-        logger.debug(f"Requesting {uri.get_full_url()}")
+        logger.debug(f"Requesting {uri.get_full_url()} with headers {uri.headers}")
         with urlopen(uri) as response:
             resp = response.read()
         return resp.decode("utf-8")
@@ -34,10 +34,14 @@ def make_request(session, request, additional_parameters={}):
     _request = deepcopy(request)
     if _request.method == 'POST':
         _request.json.update(additional_parameters)
-        logger.debug(f"Requesting {_request.url}, Payload: {json.dumps(_request.json)}")
+        logger.debug(
+            f"Requesting {_request.url}, Payload: {json.dumps(_request.json)}, Headers: {session.headers}"
+        )
     else:
         _request.params.update(additional_parameters)
-        logger.debug(f"Requesting {_request.url}, Payload: {json.dumps(_request.params)}")
+        logger.debug(
+            f"Requesting {_request.url}, Payload: {json.dumps(_request.params)}, Headers: {session.headers}"
+        )
     prepped = session.prepare_request(_request)
     resp = session.send(prepped)
     if resp.status_code != 200:

requirements-dev.txt

@@ -4,6 +4,7 @@ flake8==3.8.*
 pytest~=6.2.3
 pytest-cov~=2.11.1
 pytest-recording~=0.11.0
+pytest-console-scripts~=1.1.0
 recommonmark~=0.7.1
 Sphinx~=3.5.1
 toml~=0.10.2