use override for asset uri and maintain az://

Author: liammorrison0

datasets/gbif/gbif_pkg/gbif.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Union
+from typing import List, Union
 
 import pystac
 
@@ -15,28 +15,27 @@ def create_item(
     ) -> Union[List[pystac.Item], WaitTaskResult]:
         from gbif_pkg.gbif_tools import create_item
 
-        storage_options: Optional[Dict[str, str]]
-        asset_extra_fields: Optional[Dict[str, Optional[Dict[str, str]]]]
-
         storage, path = storage_factory.get_storage_for_file(asset_uri)
         if isinstance(storage, BlobStorage):
-            href = f"https://{storage.storage_account_name}.blob.core.windows.net/{storage.container_name}/{path}"
+            az_uri = f"az://{storage.container_name}/{path}".rstrip("/")
+            https_href = f"https://{storage.storage_account_name}.blob.core.windows.net/{storage.container_name}/{path}".rstrip("/")
             storage_options = dict(account_name=storage.storage_account_name)
             asset_extra_fields = {
                 "table:storage_options": storage_options,
                 "msft:partition_info": {
                     "is_partitioned": True
                 }
             }
+            item = create_item(
+                az_uri,
+                storage_options=storage_options,
+                asset_extra_fields=asset_extra_fields,
+                asset_href_override=https_href,
+            )
         else:
-            href = storage.get_url(path)
-            storage_options = dict()
-            asset_extra_fields = None
-
-        href = href.rstrip("/")
-
-        item = create_item(
-            href, storage_options=storage_options, asset_extra_fields=asset_extra_fields
-        )
+            href = storage.get_url(path).rstrip("/")
+            item = create_item(
+                href, storage_options=dict(), asset_extra_fields=None
+            )
 
         return [item]

datasets/gbif/gbif_pkg/gbif_tools.py

@@ -78,7 +78,7 @@
 PARQUET_MEDIA_TYPE = "application/x-parquet"
 
 
-def create_item(asset_href: str, storage_options=None, asset_extra_fields=None) -> pystac.Item:
+def create_item(asset_href: str, storage_options=None, asset_extra_fields=None, asset_href_override=None) -> pystac.Item:
     """
     Create a STAC item for a GBIF Parquet Dataset.
 
@@ -127,6 +127,7 @@ def create_item(asset_href: str, storage_options=None, asset_extra_fields=None)
         proj=False,
         asset_extra_fields=asset_extra_fields,
         count_rows=False,
+        asset_href_override=asset_href_override,
     )
 
     for column in result.properties["table:columns"]:
@@ -223,6 +224,7 @@ def generate(
     proj=True,
     storage_options=None,
     validate=True,
+    asset_href_override=None,
 ) -> T:
     """
     Generate a STAC Item from a Parquet Dataset.
@@ -406,7 +408,7 @@ def generate(
 
     if asset_key:
         asset = pystac.asset.Asset(
-            uri,
+            asset_href_override if asset_href_override is not None else uri,
             title="Dataset root",
             media_type=PARQUET_MEDIA_TYPE,
             roles=["data"],