Repository: opendatacube/datacube-core

Issue:
`IndexError` when accessing `.geobox` on the result of `xarray.open_rasterio` for epsgs:4326 files
### Expected behaviour

Get proper `GeoBox` object.

### Actual behaviour

`KeyError` is raised

```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/.envs/odc/lib/python3.8/site-packages/xarray/core/dataarray.py in _getitem_coord(self, key)
    616         try:
--> 617             var = self._coords[key]
    618         except KeyError:

KeyError: 'longitude'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
<ipython-input-12-e899bd2aebfe> in <module>
----> 1 xx.geobox

~/wk/datacube-core/datacube/utils/xarray_geoextensions.py in _xarray_geobox(obj)
    148 
    149     dims = crs.dimensions
--> 150     return geometry.GeoBox(obj[dims[1]].size, obj[dims[0]].size, obj.affine, crs)
```

Problematic code is here:
https://github.com/opendatacube/datacube-core/blob/635c8bbeedd522d6313abf41f3dadf2c743e1a3f/datacube/utils/xarray_geoextensions.py#L149-L150

It assumes that "YX" dimensions as named by CRS are present on the xarray object as well, which is not the case when xarray is no constructed by datacube code but by other libraries, particularly when CRS is geographic.





Discussion:

Resolving patch:
diff --git a/datacube/utils/dask.py b/datacube/utils/dask.py
--- a/datacube/utils/dask.py
+++ b/datacube/utils/dask.py
@@ -85,6 +85,14 @@ def start_local_dask(n_workers: int = 1,
 
     NOTE: if `memory_limit` is supplied, it will be parsed and divided equally between workers.
     """
+
+    # if dashboard.link set to default value and running behind hub, make dashboard link go via proxy
+    if dask.config.get("distributed.dashboard.link") == '{scheme}://{host}:{port}/status':
+        jup_prefix = os.environ.get('JUPYTERHUB_SERVICE_PREFIX')
+        if jup_prefix is not None:
+            jup_prefix = jup_prefix.rstrip('/')
+            dask.config.set({"distributed.dashboard.link": f"{jup_prefix}/proxy/{{port}}/status"})
+
     memory_limit = compute_memory_per_worker(n_workers=n_workers,
                                              memory_limit=memory_limit,
                                              mem_safety_margin=mem_safety_margin)
diff --git a/datacube/utils/math.py b/datacube/utils/math.py
--- a/datacube/utils/math.py
+++ b/datacube/utils/math.py
@@ -1,4 +1,4 @@
-from typing import Tuple, Optional, Any, cast
+from typing import Tuple, Union, Optional, Any, cast
 from math import ceil, fmod
 
 import numpy
@@ -35,7 +35,8 @@ def unsqueeze_dataset(ds: xr.Dataset, dim: str, coord: int = 0, pos: int = 0) ->
     return ds
 
 
-def spatial_dims(xx: xr.DataArray, relaxed: bool = False) -> Optional[Tuple[str, str]]:
+def spatial_dims(xx: Union[xr.DataArray, xr.Dataset],
+                 relaxed: bool = False) -> Optional[Tuple[str, str]]:
     """ Find spatial dimensions of `xx`.
 
         Checks for presence of dimensions named:
@@ -64,19 +65,12 @@ def spatial_dims(xx: xr.DataArray, relaxed: bool = False) -> Optional[Tuple[str,
     return None
 
 
-def clamp(x, l, u):
+def clamp(x, lo, up):
     """
-    clamp x to be l <= x <= u
-
-    >>> clamp(5, 1, 10)
-    5
-    >>> clamp(-1, 1, 10)
-    1
-    >>> clamp(12, 1, 10)
-    10
+    clamp x to be lo <= x <= up
     """
-    assert l <= u
-    return l if x < l else u if x > u else x
+    assert lo <= up
+    return lo if x < lo else up if x > up else x
 
 
 def is_almost_int(x: float, tol: float):
diff --git a/datacube/utils/xarray_geoextensions.py b/datacube/utils/xarray_geoextensions.py
--- a/datacube/utils/xarray_geoextensions.py
+++ b/datacube/utils/xarray_geoextensions.py
@@ -25,7 +25,7 @@ def _norm_crs(crs):
         raise ValueError('Can not interpret {} as CRS'.format(type(crs)))
 
 
-def _get_crs_from_attrs(obj):
+def _get_crs_from_attrs(obj, sdims):
     """ Looks for attribute named `crs` containing CRS string
         1. Checks spatials coords attrs
         2. Checks data variable attrs
@@ -45,16 +45,12 @@ def _get_crs_from_attrs(obj):
     else:
         data_array = obj
 
-    sdims = spatial_dims(data_array, relaxed=True)
-    if sdims is not None:
-        crs_set = set(data_array[d].attrs.get('crs', None) for d in sdims)
-        crs = None
-        if len(crs_set) > 1:
-            raise ValueError('Spatial dimensions have different crs.')
-        elif len(crs_set) == 1:
-            crs = crs_set.pop()
-    else:
-        crs = None
+    crs_set = set(data_array[d].attrs.get('crs', None) for d in sdims)
+    crs = None
+    if len(crs_set) > 1:
+        raise ValueError('Spatial dimensions have different crs.')
+    elif len(crs_set) == 1:
+        crs = crs_set.pop()
 
     if crs is None:
         # fall back option
@@ -109,15 +105,20 @@ def _get_crs_from_coord(obj, mode='strict'):
         raise ValueError(f"Mode needs to be: strict|any|all got {mode}")
 
 
-def _xarray_affine(obj):
+def _xarray_affine_impl(obj):
     sdims = spatial_dims(obj, relaxed=True)
     if sdims is None:
-        return None
+        return None, None
 
     yy, xx = (obj[dim] for dim in sdims)
     fallback_res = (coord.attrs.get('resolution', None) for coord in (xx, yy))
 
-    return affine_from_axis(xx.values, yy.values, fallback_res)
+    return affine_from_axis(xx.values, yy.values, fallback_res), sdims
+
+
+def _xarray_affine(obj):
+    transform, _ = _xarray_affine_impl(obj)
+    return transform
 
 
 def _xarray_extent(obj):
@@ -126,6 +127,10 @@ def _xarray_extent(obj):
 
 
 def _xarray_geobox(obj):
+    transform, sdims = _xarray_affine_impl(obj)
+    if sdims is None:
+        return None
+
     crs = None
     try:
         crs = _get_crs_from_coord(obj)
@@ -134,7 +139,7 @@ def _xarray_geobox(obj):
 
     if crs is None:
         try:
-            crs = _get_crs_from_attrs(obj)
+            crs = _get_crs_from_attrs(obj, sdims)
         except ValueError:
             pass
 
@@ -146,13 +151,14 @@ def _xarray_geobox(obj):
     except ValueError:
         return None
 
-    dims = crs.dimensions
-    return geometry.GeoBox(obj[dims[1]].size, obj[dims[0]].size, obj.affine, crs)
+    h, w = (obj.coords[dim].size for dim in sdims)
+
+    return geometry.GeoBox(w, h, transform, crs)
 
 
-xarray.Dataset.geobox = property(_xarray_geobox)   # type: ignore
-xarray.Dataset.affine = property(_xarray_affine)   # type: ignore
-xarray.Dataset.extent = property(_xarray_extent)   # type: ignore
-xarray.DataArray.geobox = property(_xarray_geobox) # type: ignore
-xarray.DataArray.affine = property(_xarray_affine) # type: ignore
-xarray.DataArray.extent = property(_xarray_extent) # type: ignore
+xarray.Dataset.geobox = property(_xarray_geobox)    # type: ignore
+xarray.Dataset.affine = property(_xarray_affine)    # type: ignore
+xarray.Dataset.extent = property(_xarray_extent)    # type: ignore
+xarray.DataArray.geobox = property(_xarray_geobox)  # type: ignore
+xarray.DataArray.affine = property(_xarray_affine)  # type: ignore
+xarray.DataArray.extent = property(_xarray_extent)  # type: ignore


Test patch:
diff --git a/tests/test_utils_dask.py b/tests/test_utils_dask.py
--- a/tests/test_utils_dask.py
+++ b/tests/test_utils_dask.py
@@ -38,6 +38,12 @@ def test_compute_tasks():
     del client
 
 
+def test_start_local_dask_dashboard_link(monkeypatch):
+    monkeypatch.setenv('JUPYTERHUB_SERVICE_PREFIX', 'user/test/')
+    client = start_local_dask()
+    assert client.dashboard_link.startswith('user/test/proxy/')
+
+
 def test_partition_map():
     tasks = partition_map(10, str, range(101))
     tt = [t for t in tasks]
diff --git a/tests/test_xarray_extension.py b/tests/test_xarray_extension.py
--- a/tests/test_xarray_extension.py
+++ b/tests/test_xarray_extension.py
@@ -128,17 +128,17 @@ def test_crs_from_attrs():
     xx_3857 = mk_sample_xr_dataset(crs=epsg3857)
     xx_4326 = mk_sample_xr_dataset(crs=epsg4326)
 
-    assert _get_crs_from_attrs(xx_none) is None
-    assert _get_crs_from_attrs(xx_none.band) is None
-    assert _get_crs_from_attrs(xx_3857.band) == epsg3857
-    assert _get_crs_from_attrs(xx_3857) == epsg3857
-    assert _get_crs_from_attrs(xx_4326.band) == epsg4326
-    assert _get_crs_from_attrs(xx_4326) == epsg4326
+    assert _get_crs_from_attrs(xx_none, ('y', 'x')) is None
+    assert _get_crs_from_attrs(xx_none.band, ('y', 'x')) is None
+    assert _get_crs_from_attrs(xx_3857.band, ('y', 'x')) == epsg3857
+    assert _get_crs_from_attrs(xx_3857, ('y', 'x')) == epsg3857
+    assert _get_crs_from_attrs(xx_4326.band, ('latitude', 'longitude')) == epsg4326
+    assert _get_crs_from_attrs(xx_4326, ('latitude', 'longitude')) == epsg4326
 
-    assert _get_crs_from_attrs(xr.Dataset()) is None
+    assert _get_crs_from_attrs(xr.Dataset(), None) is None
 
     # check inconsistent CRSs
     xx = xx_3857.copy()
     xx.x.attrs['crs'] = xx_4326.attrs['crs']
     with pytest.raises(ValueError):
-        _get_crs_from_attrs(xx)
+        _get_crs_from_attrs(xx, ('y', 'x'))