Make signing STAC objects more convenient (#12)

* Make sign work with URL, Asset, Item, or ItemCollection

Also added a search_and_sign function that takes an ItemSearch,
performs the search, and signs the resulting ItemCollection

* Rename test_sign_assets file to better reflect its purpose

* Update README

* Update test class name

* Switch to using singledispatch for sign functions

* Add missing docstring

Author: kshepard

README.md

@@ -42,28 +42,40 @@ pip install -e .
 
 ## Usage
 
-This library currently assists with signing Azure Blob Storage URLs, both within PySTAC assets, and by providing raw URLs. The following examples demonstrate both of these use cases:
+This library currently assists with signing Azure Blob Storage URLs. The `sign` function operates directly on an HREF string, as well as several [PySTAC](https://github.com/stac-utils/pystac) objects: `Asset`, `Item`, and `ItemCollection`. In addition, the `sign` function accepts a [STAC API Client](https://github.com/stac-utils/pystac-client) `ItemSearch`, which performs a search and returns the resulting `ItemCollection` with all assets signed. The following example demonstrates these use cases:
 
 ```python
-import pystac
+from pystac import Asset, Item
+from pystac_client import ItemCollection, ItemSearch
 import planetary_computer as pc
 
-raw_item: pystac.Item = ...
-item: pystac.Item = pc.sign_assets(raw_item)
 
+# The sign function may be called directly on the Item
+raw_item: Item = ...
+item: Item = pc.sign(raw_item)
 # Now use the item however you want. All appropriate assets are signed for read access.
-```
-
-```python
-import planetary_computer as pc
-import pystac
-
-item: pystac.Item = ...  # Landsat item
-
-b4_href = pc.sign(item.assets['SR_B4'].href)
 
-with rasterio.open(b4_href) as ds:
-   ...
+# The sign function also works with an Asset
+raw_asset: Asset = raw_item.assets['SR_B4']
+asset = pc.sign(raw_asset)
+
+# The sign function also works with an HREF
+raw_href: str = raw_asset.href
+href = pc.sign(raw_href)
+
+# The sign function also works with an ItemCollection
+raw_item_collection = ItemCollection([raw_item])
+item_collection = pc.sign(raw_item_collection)
+
+# The sign function also accepts an ItemSearch, and signs the resulting ItemCollection
+search = ItemSearch(
+    url=...,
+    bbox=...,
+    collections=...,
+    limit=...,
+    max_items=...,
+)
+signed_item_collection = pc.sign(search)
 ```
 
 

planetary_computer/__init__.py

@@ -1,5 +1,5 @@
 """Planetary Computer Python SDK"""
 # flake8:noqa
 
-from planetary_computer.sas import sign, sign_assets  # type:ignore
+from planetary_computer.sas import sign  # type:ignore
 from planetary_computer.settings import set_subscription_key  # type:ignore

planetary_computer/sas.py

@@ -1,14 +1,15 @@
 from datetime import datetime, timezone
-from typing import Dict
+from typing import Any, Dict
 
+from functools import singledispatch
+import requests
 from pydantic import BaseModel, Field
-import pystac
+from pystac import Asset, Item
 from pystac.utils import datetime_to_str
-import requests
+from pystac_client import ItemCollection, ItemSearch
 
-
-from planetary_computer.utils import parse_blob_url
 from planetary_computer.settings import Settings
+from planetary_computer.utils import parse_blob_url
 
 
 class SASBase(BaseModel):
@@ -50,24 +51,24 @@ def ttl(self) -> float:
 TOKEN_CACHE: Dict[str, SASToken] = {}
 
 
-def sign(url: str) -> str:
-    """Sign a URL with a Shared Access (SAS)
-    Token, which allows for read access.
-
-    Parameters
-    ----------
-    url (str): The HREF of the asset in the format of a URL.
-        This can be found on STAC Item's Asset 'href' value.
+@singledispatch
+def sign(obj: Any) -> Any:
+    """Sign the relevant URLs belonging to any supported object with a
+    Shared Access (SAS) Token, which allows for read access.
 
-    Returns
-    -------
-    The signed HREF that permits read access to the asset.
+    Args:
+        obj (Any): The object to sign. Must be one of:
+            str (URL), Asset, Item, ItemCollection, or ItemSearch
+    Returns:
+        Any: A copy of the object where all relevant URLs have been signed
     """
-    link = sign_link(url)
-    return link.href
+    raise TypeError(
+        "Invalid type, must be one of: str, Asset, Item, ItemCollection, or ItemSearch"
+    )
 
 
-def sign_link(url: str) -> SignedLink:
+@sign.register(str)
+def _sign_url(url: str) -> str:
     """Sign a URL with a Shared Access (SAS) Token, which allows for read access.
 
     Args:
@@ -76,9 +77,7 @@ def sign_link(url: str) -> SignedLink:
             value.
 
     Returns:
-        SignedLink: An object that contains the signed HREF
-        in the format of a URL and the expiry time, which
-        is when the HREF will no longer permit read access.
+        str: The signed HREF
     """
     settings = Settings.get()
     account, container = parse_blob_url(url)
@@ -99,22 +98,76 @@ def sign_link(url: str) -> SignedLink:
         if not token:
             raise ValueError(f"No token found in response: {response.json()}")
         TOKEN_CACHE[token_request_url] = token
-    return token.sign(url)
+    return token.sign(url).href
 
 
-def sign_assets(item: pystac.Item) -> pystac.Item:
+@sign.register(Item)
+def _sign_item(item: Item) -> Item:
     """Sign all assets within a PySTAC item
 
     Args:
-        item (pystac.Item): The Item whose assets that will be signed
+        item (Item): The Item whose assets that will be signed
 
     Returns:
-        pystac.Item: A new copy of the Item where all assets HREFs have
+        Item: A new copy of the Item where all assets' HREFs have
         been replaced with a signed version. In addition, a "msft:expiry"
         property is added to the Item properties indicating the earliest
         expiry time for any assets that were signed.
     """
     signed_item = item.clone()
     for key in signed_item.assets:
-        signed_item.assets[key].href = sign(signed_item.assets[key].href)
+        signed_item.assets[key] = sign(signed_item.assets[key])
     return signed_item
+
+
+@sign.register(Asset)
+def _sign_asset(asset: Asset) -> Asset:
+    """Sign a PySTAC asset
+
+    Args:
+        asset (Asset): The Asset to sign
+
+    Returns:
+        Asset: A new copy of the Asset where the HREF is replaced with a
+        signed version.
+    """
+    signed_asset = asset.clone()
+    signed_asset.href = sign(signed_asset.href)
+    return signed_asset
+
+
+@sign.register(ItemCollection)
+def _sign_item_collection(item_collection: ItemCollection) -> ItemCollection:
+    """Sign a PySTAC item collection
+
+    Args:
+        item_collection (ItemCollection): The ItemCollection whose assets will be signed
+
+    Returns:
+        ItemCollection: A new copy of the ItemCollection where all assets'
+        HREFs for each item have been replaced with a signed version. In addition,
+        a "msft:expiry" property is added to the Item properties indicating the
+        earliest expiry time for any assets that were signed.
+    """
+    return ItemCollection.from_dict(
+        {
+            "type": "FeatureCollection",
+            "features": [sign(item).to_dict() for item in item_collection],
+        }
+    )
+
+
+@sign.register(ItemSearch)
+def _search_and_sign(search: ItemSearch) -> ItemCollection:
+    """Perform a PySTAC Client search, and sign the resulting item collection
+
+    Args:
+        search (ItemSearch): The ItemSearch whose resulting item assets will be signed
+
+    Returns:
+        ItemCollection: The resulting ItemCollection of the search where all assets'
+        HREFs for each item have been replaced with a signed version. In addition,
+        a "msft:expiry" property is added to the Item properties indicating the
+        earliest expiry time for any assets that were signed.
+    """
+    return sign(search.items_as_collection())

setup.cfg

@@ -14,6 +14,7 @@ install_requires =
     click>=7.1
     pydantic[dotenv]>=1.7.3
     pystac>=0.5.6,<0.6
+    pystac-client>=0.1.1,<0.2.0
     pytz>=2020.5
     requests>=2.25.1
 

tests/test_sign_assets.py tests/test_signing.pytests/test_sign_assets.py renamed to tests/test_signing.py

@@ -1,13 +1,16 @@
 import os
 import json
+from typing import Any, Dict
 import unittest
 from urllib.parse import parse_qs, urlparse
+import warnings
 
 import requests
 
 import planetary_computer as pc
 from planetary_computer.utils import parse_blob_url
-import pystac
+from pystac import Item
+from pystac_client import ItemCollection, ItemSearch
 
 
 ACCOUNT_NAME = "naipeuwest"
@@ -23,17 +26,26 @@
     "/GRANULE/L2A_T10TET_A018672_20201002T192031/QI_DATA/T10TET_20201002T191229_PVI.tif"
 )
 
+PC_SEARCH_URL = "https://planetarycomputer.microsoft.com/api/stac/v1/search"
 
-def get_sample_item() -> pystac.Item:
+
+def get_sample_item_dict() -> Dict[str, Any]:
     file_path = os.path.abspath(
         os.path.join(os.path.dirname(__file__), "data-files/sample-item.json")
     )
     with open(file_path) as json_file:
-        item_dict = json.load(json_file)
-    return pystac.Item.from_dict(item_dict)
+        return json.load(json_file)
+
+
+def get_sample_item() -> Item:
+    return Item.from_dict(get_sample_item_dict())
+
+
+def get_sample_item_collection() -> ItemCollection:
+    return ItemCollection([get_sample_item()])
 
 
-class TestSignAssests(unittest.TestCase):
+class TestSigning(unittest.TestCase):
     def assertSigned(self, url: str) -> None:
         # Ensure the signed item has an "se" URL parameter added to it,
         # which indicates it has been signed
@@ -57,17 +69,38 @@ def test_unsigned_assets(self) -> None:
         self.assertEqual(EXP_METADATA, item.assets["metadata"].href)
         self.assertEqual(EXP_THUMBNAIL, item.assets["thumbnail"].href)
 
-    def test_signed_assets(self) -> None:
-        unsigned_item = get_sample_item()
-        signed_item = pc.sign_assets(unsigned_item)
-
-        # Ensure the original item wasn't mutated, and all URLs are signed
+    def verify_signed_urls_in_item(self, signed_item: Item) -> None:
         for key in ["image", "metadata", "thumbnail"]:
             signed_url = signed_item.assets[key].href
-            self.assertNotEqual(unsigned_item.assets[key].href, signed_url)
             self.assertSigned(signed_url)
 
+    def test_signed_assets(self) -> None:
+        signed_item = pc.sign(get_sample_item())
+        self.verify_signed_urls_in_item(signed_item)
+
     def test_read_signed_asset(self) -> None:
         signed_href = pc.sign(SENTINEL_THUMBNAIL)
         r = requests.get(signed_href)
         self.assertEqual(r.status_code, 200)
+
+    def test_signed_item_collection(self) -> None:
+        signed_item_collection = pc.sign(get_sample_item_collection())
+        self.assertEqual(len(list(signed_item_collection)), 1)
+        for signed_item in signed_item_collection:
+            self.verify_signed_urls_in_item(signed_item)
+
+    def test_search_and_sign(self) -> None:
+        # Filter out a resource warning coming from within the pystac-client search
+        warnings.simplefilter("ignore", ResourceWarning)
+
+        search = ItemSearch(
+            url=PC_SEARCH_URL,
+            bbox=(-73.21, 43.99, -73.12, 44.05),
+            collections=CONTAINER_NAME,
+            limit=1,
+            max_items=1,
+        )
+        signed_item_collection = pc.sign(search)
+        self.assertEqual(len(list(signed_item_collection)), 1)
+        for signed_item in signed_item_collection:
+            self.verify_signed_urls_in_item(signed_item)