From 85eab26e0fd4b2834cfc4630360f00fd96c78fde Mon Sep 17 00:00:00 2001
From: Ariana Barzinpour <ariana.barzinpour@ga.gov.au>
Date: Tue, 3 Sep 2024 02:25:24 +0000
Subject: [PATCH] handle missing _md product attr; update use of DatasetType

---
 docs/stac-vs-odc.rst          |  2 +-
 odc/stac/eo3/_eo3converter.py | 26 +++++++++++++++-----------
 tests/test_eo3converter.py    | 12 ++++++++++++
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/docs/stac-vs-odc.rst b/docs/stac-vs-odc.rst
index 68808d8f..58e3e25b 100644
--- a/docs/stac-vs-odc.rst
+++ b/docs/stac-vs-odc.rst
@@ -13,7 +13,7 @@ similar concepts.
      - ODC
      - Description
    * - :py:class:`~pystac.Collection`
-     - Product or :py:class:`~datacube.model.DatasetType`
+     - :py:class:`~datacube.model.Product`
      - Collection of observations across space and time
    * - :py:class:`~pystac.Item`
      - :py:class:`~datacube.model.Dataset`
diff --git a/odc/stac/eo3/_eo3converter.py b/odc/stac/eo3/_eo3converter.py
index 8ab7e6e0..fd55ec86 100644
--- a/odc/stac/eo3/_eo3converter.py
+++ b/odc/stac/eo3/_eo3converter.py
@@ -20,7 +20,7 @@
 except ImportError:
     from datacube.index.abstract import default_metadata_type_docs  # type: ignore
 
-from datacube.model import Dataset, DatasetType, metadata_from_doc
+from datacube.model import Dataset, Product, metadata_from_doc
 from odc.geo import CRS
 from odc.geo.geobox import GeoBox
 from toolz import dicttoolz
@@ -60,7 +60,7 @@
 )
 
 
-def _to_product(md: RasterCollectionMetadata) -> DatasetType:
+def _to_product(md: RasterCollectionMetadata) -> Product:
     def make_band(
         band_key: BandKey,
         band: RasterBandMetadata,
@@ -95,11 +95,11 @@ def make_band(
             for band_key, band in md.meta.bands.items()
         ],
     }
-    return DatasetType(_eo3, doc)
+    return Product(_eo3, doc)
 
 
 @singledispatch
-def infer_dc_product(x: Any, cfg: Optional[ConversionConfig] = None) -> DatasetType:
+def infer_dc_product(x: Any, cfg: Optional[ConversionConfig] = None) -> Product:
     """Overloaded function."""
     raise TypeError(
         "Invalid type, must be one of: pystac.item.Item, pystac.collection.Collection"
@@ -109,7 +109,7 @@ def infer_dc_product(x: Any, cfg: Optional[ConversionConfig] = None) -> DatasetT
 @infer_dc_product.register(pystac.item.Item)
 def infer_dc_product_from_item(
     item: pystac.item.Item, cfg: Optional[ConversionConfig] = None
-) -> DatasetType:
+) -> Product:
     """
     Infer Datacube product object from a STAC Item.
 
@@ -164,7 +164,7 @@ def _to_dataset(
     item: ParsedItem,
     properties: Dict[str, Any],
     ds_uuid: uuid.UUID,
-    product: DatasetType,
+    product: Product,
 ) -> Dataset:
     # pylint: disable=too-many-locals
 
@@ -227,7 +227,7 @@ def _to_dataset(
 
 
 def _item_to_ds(
-    item: pystac.item.Item, product: DatasetType, cfg: Optional[ConversionConfig] = None
+    item: pystac.item.Item, product: Product, cfg: Optional[ConversionConfig] = None
 ) -> Dataset:
     """
     Construct Dataset object from STAC Item and previously constructed Product.
@@ -251,7 +251,7 @@ def _item_to_ds(
 def stac2ds(
     items: Iterable[pystac.item.Item],
     cfg: Optional[ConversionConfig] = None,
-    product_cache: Optional[Dict[str, DatasetType]] = None,
+    product_cache: Optional[Dict[str, Product]] = None,
 ) -> Iterator[Dataset]:
     """
     STAC :class:`~pystac.item.Item` to :class:`~datacube.model.Dataset` stream converter.
@@ -276,7 +276,7 @@ def stac2ds(
 
     :param product_cache:
        Input/Output parameter, contains mapping from collection name to deduced product definition,
-       i.e. :py:class:`datacube.model.DatasetType` object.
+       i.e. :py:class:`datacube.model.Product` object.
 
     .. rubric: Sample Configuration
 
@@ -312,7 +312,7 @@ def stac2ds(
          warnings: ignore
 
     """
-    products: Dict[str, DatasetType] = {} if product_cache is None else product_cache
+    products: Dict[str, Product] = {} if product_cache is None else product_cache
     for item in items:
         collection_id = _collection_id(item)
         product = products.get(collection_id)
@@ -321,6 +321,10 @@ def stac2ds(
         if product is None:
             product = infer_dc_product(item, cfg)
             products[collection_id] = product
+        # if product_cache was provided, the products within may not have the custom _md attr
+        if not hasattr(product, "_md"):  # pylint: disable=protected-access
+            md = extract_collection_metadata(item, cfg)
+            setattr(product, "_md", md)  # pylint: disable=protected-access
 
         yield _item_to_ds(item, product, cfg)
 
@@ -328,7 +332,7 @@ def stac2ds(
 @infer_dc_product.register(pystac.collection.Collection)
 def infer_dc_product_from_collection(
     collection: pystac.collection.Collection, cfg: Optional[ConversionConfig] = None
-) -> DatasetType:
+) -> Product:
     """
     Construct Datacube Product definition from STAC Collection.
 
diff --git a/tests/test_eo3converter.py b/tests/test_eo3converter.py
index 0d655466..2cb7e996 100644
--- a/tests/test_eo3converter.py
+++ b/tests/test_eo3converter.py
@@ -234,3 +234,15 @@ def test_old_imports():
 
     with pytest.raises(AttributeError):
         _ = odc.stac.no_such_thing
+
+
+def test_product_cache(sentinel_stac_ms: pystac.item.Item):
+    item = sentinel_stac_ms
+    # simulate a product that was not created via infer_dc_product
+    # (and therefore did not have the _md attr set)
+    product = infer_dc_product(item, STAC_CFG)
+    delattr(product, "_md")
+    
+    # make sure it doesn't error when product_cache is provided
+    (ds,) = stac2ds([item], STAC_CFG, {product.name: product})
+    assert ds.id