FIX: prefer_streaming=True now returns all file links (#92)

* properly return a list of links, selected by the fastest host that returns a response

* tests now require actual OPENDAP links to get partitioned

Author: nocollier

intake_esgf/base.py

@@ -89,6 +89,27 @@ def get_globus_endpoints(info: dict) -> list[str]:
     return globus_endpoints
 
 
+def select_streaming_link(links: list[str], df_rate: pd.DataFrame) -> str:
+    """
+    Select a streaming link from a list of options based on fastest download speeds.
+    """
+    # There are some .nc and .nc.html links in these records, only the .nc links
+    # can be read by xarray.
+    links = [link for link in links if link.endswith(".nc")]
+    # Sort the links by the fastest hostname we have seen when transferring.
+    links = sorted(
+        links, key=partial(sort_download_links, df_rate=df_rate), reverse=True
+    )
+    # Return the first of these links whose html page returns a valid response.
+    # This is particular to OPENDAP and will need rethought for other virtual
+    # methods.
+    for link in links:
+        resp = requests.get(link + ".html", stream=True, timeout=10)
+        if resp.status_code == 200:
+            return link
+    raise ValueError(f"None of these links appears functional {links}")
+
+
 def partition_infos(
     infos: list[dict], prefer_streaming: bool, prefer_globus: bool
 ) -> tuple[dict, dict]:
@@ -132,6 +153,9 @@ def partition_infos(
     active_endpoints = set()
 
     # Partition and setup all the file infos based on a priority
+    df_rate = get_download_rate_dataframe(
+        Path(intake_esgf.conf["download_db"]).expanduser()
+    )
     for i, info in enumerate(infos):
         key = info["key"]
 
@@ -151,16 +175,23 @@ def partition_infos(
 
         # 2) does the user prefer to stream data?
         if prefer_streaming:
-            # how do we choose a link?
+            # What possible links are there to stream from?
             preferred_sources = ["VirtualZarr", "OPENDAP"]  # move to configure
             links = [
                 link
                 for src in (set(preferred_sources) & set(info))
                 for link in info[src]
             ]
             if links:
-                # for now just use first link, we need to do better
-                ds[key] = [links[0]]
+                # Try to find a streaming link from the fastest servers we have
+                # seen while downloading. If this fails, we revert to https.
+                try:
+                    link = select_streaming_link(links, df_rate)
+                except ValueError:
+                    break
+                if key not in ds:
+                    ds[key] = []
+                ds[key].append(link)
                 infos_stream.append(info)  # maybe not needed
                 continue
 

intake_esgf/tests/test_basic.py

@@ -151,17 +151,18 @@ def test_partition_infos():
 
 
 def test_partition_infos_stream():
+    VALID_OPENDAP_LINK = "https://esgf-data1.llnl.gov/thredds/dodsC/css03_data/CMIP6/CMIP/AS-RCEC/TaiESM1/historical/r1i1p1f1/day/tasmax/gn/v20210517/tasmax_day_TaiESM1_historical_r1i1p1f1_gn_20100101-20141231.nc"
     infos = [
         {
             "key": "dataset1",
             "path": Path("file1"),
-            "VirtualZarr": ["link1", "link2"],
+            "VirtualZarr": [VALID_OPENDAP_LINK, VALID_OPENDAP_LINK],
         },
         {
             "key": "dataset2",
             "path": Path("file1"),
             "HTTPServer": ["link1", "link2"],
-            "OPENDAP": ["link1", "link2"],
+            "OPENDAP": [VALID_OPENDAP_LINK, VALID_OPENDAP_LINK],
         },
     ]
     infos_, ds = partition_infos(infos, False, False)