From 4de8c61d5c3229721250da593b9f87e4e1c8c6cf Mon Sep 17 00:00:00 2001
From: Stefaan Lippens <stefaan.lippens@vito.be>
Date: Mon, 25 Nov 2024 17:31:21 +0100
Subject: [PATCH] Issue #666 load_stac: fix band filtering in cube metadata

---
 openeo/rest/datacube.py       |  4 ++++
 tests/rest/test_connection.py | 13 +++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/openeo/rest/datacube.py b/openeo/rest/datacube.py
index c91fb722a..f0da58849 100644
--- a/openeo/rest/datacube.py
+++ b/openeo/rest/datacube.py
@@ -203,6 +203,7 @@ def load_collection(
                 metadata = None
             if metadata:
                 bands = [b if isinstance(b, str) else metadata.band_dimension.band_name(b) for b in bands]
+                # TODO: also apply spatial/temporal filters to metadata?
                 metadata = metadata.filter_bands(bands)
             arguments['bands'] = bands
 
@@ -385,6 +386,9 @@ def load_stac(
         graph = PGNode("load_stac", arguments=arguments)
         try:
             metadata = metadata_from_stac(url)
+            if bands:
+                # TODO: also apply spatial/temporal filters to metadata?
+                metadata = metadata.filter_bands(band_names=bands)
         except Exception:
             log.warning(f"Failed to extract cube metadata from STAC URL {url}", exc_info=True)
             metadata = None
diff --git a/tests/rest/test_connection.py b/tests/rest/test_connection.py
index f392d0321..5f98890e1 100644
--- a/tests/rest/test_connection.py
+++ b/tests/rest/test_connection.py
@@ -2651,6 +2651,19 @@ def test_load_stac_no_cube_extension_temporal_dimension(self, con120, tmp_path,
         cube = con120.load_stac(str(stac_path))
         assert cube.metadata.temporal_dimension == TemporalDimension(name="t", extent=dim_extent)
 
+    def test_load_stac_band_filtering(self, con120, tmp_path):
+        stac_path = tmp_path / "stac.json"
+        stac_data = StacDummyBuilder.collection(
+            summaries={"eo:bands": [{"name": "B01"}, {"name": "B02"}, {"name": "B03"}]}
+        )
+        stac_path.write_text(json.dumps(stac_data))
+
+        cube = con120.load_stac(str(stac_path))
+        assert cube.metadata.band_names == ["B01", "B02", "B03"]
+
+        cube = con120.load_stac(str(stac_path), bands=["B03", "B02"])
+        assert cube.metadata.band_names == ["B03", "B02"]
+
 
 @pytest.mark.parametrize(
     "data",