Merge pull request #39 from stac-utils/mah/fix_headers

Fix headers

Author: matthewhanson

pystac_client/cli.py

@@ -5,7 +5,6 @@
 import sys
 
 from .item_collection import ItemCollection
-
 from .client import Client
 from .version import __version__
 
@@ -14,11 +13,11 @@
 logger = logging.getLogger(__name__)
 
 
-def search(url=STAC_URL, matched=False, save=None, **kwargs):
+def search(url=STAC_URL, matched=False, save=None, headers=None, **kwargs):
     """ Main function for performing a search """
 
     try:
-        catalog = Client.open(url)
+        catalog = Client.open(url, headers=headers)
         search = catalog.search(**kwargs)
 
         if matched:
@@ -35,6 +34,7 @@ def search(url=STAC_URL, matched=False, save=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 file or string)',
+                        help='Additional request headers (JSON string or file)',
                         default=None)
 
     subparsers = parser0.add_subparsers(dest='command')
@@ -106,6 +106,15 @@ def parse_args(args):
                 else:
                     parsed_args['intersects'] = data
 
+    # if headers provided, parse it
+    if 'headers' in parsed_args:
+        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
 
 
@@ -125,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
@@ -48,7 +50,8 @@ def __init__(self,
                  extra_fields=None,
                  href=None,
                  catalog_type=None,
-                 conformance=None):
+                 conformance=None,
+                 headers=None):
         super().__init__(id=id,
                          description=description,
                          title=title,
@@ -66,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 = {}
+        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
@@ -84,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