MV search redux (#378)

* Basic SQLAlchemy interface to space_time_view

* Search returns unioned extent.

* Search returns ODC Datasets.

* Use MVs for dataset search.

* Delete obsolete file

* only return bits of unioned extent inside out query.

* Interact with Postgres with EPSG:4326

* Use Count and Extent modes in GetMap.

* Pylint is just doing its job.

* Unit test MVSelectOpts.

* Make unit tests more fun.

* Test mv_index.py

* Remove legacy dataset search code.  Instant coverage boost!

* Grant select permission to public.

* Handle null extents.

* More readable refactor of previous bug fix.

* I guess I was right the first time.

* Remove support for blocking materialised view updates.

* Oops - unfinished refactor.

* Update database documentation.

* Use MVSelectOpts.COUNT in WCS GetCoverage requests.

* Embarassing typo in last commit.

Co-authored-by: phaesler <paul.haesler@data61.csiro.au>
Co-authored-by: Pin Jin <pindge.j@gmail.com>
Co-authored-by: Tisham Dhar <tisham.dhar@ga.gov.au>

Author: 4 people

datacube_ows/band_mapper.py

@@ -24,7 +24,7 @@
     GetFeatureInfoParameters, solar_correct_data, collapse_datasets_to_times
 from datacube_ows.ogc_utils import local_solar_date_range, dataset_center_time, ConfigException, tz_for_geometry, \
     solar_date, year_date_range, month_date_range
-
+from datacube_ows.mv_index import MVSelectOpts, mv_search_datasets
 from datacube_ows.utils import log_call, group_by_statistical
 
 import logging
@@ -62,51 +62,37 @@ def needed_bands(self):
         return self._needed_bands
 
     @log_call
-    def datasets(self, index, mask=False, all_time=False, point=None):
+    def n_datasets(self, index, all_time=False, point=None):
+        return self.datasets(index,
+                             all_time=all_time, point=point,
+                             mode=MVSelectOpts.COUNT)
+    @log_call
+    def datasets(self, index,
+                 mask=False, all_time=False, point=None,
+                 mode=MVSelectOpts.DATASETS):
         # Return datasets as a time-grouped xarray DataArray. (or None if no datasets)
         # No PQ product, so no PQ datasets.
         if not self._product.pq_name and mask:
             return None
 
-        if self._product.multi_product:
-            prod_name = self._product.pq_names if mask and self._product.pq_name else self._product.product_names
-            query_args = {
-                "geopolygon": self._geobox.extent
-            }
+        if point:
+            geom = point
         else:
-            prod_name = self._product.pq_name if mask and self._product.pq_name else self._product.product_name
-            query_args = {
-                "product": prod_name,
-                "geopolygon": self._geobox.extent
-            }
-        if all_time or (mask and self._product.pq_ignore_time):
-            all_datasets = self._dataset_query(index, prod_name, query_args)
+            geom = self._geobox.extent
+
+        if all_time:
+            times = None
         else:
-            all_datasets = []
-            for th in self._times:
-                query_args["time"] = th
-                all_datasets.extend(self._dataset_query(index, prod_name, query_args))
-        return datacube.Datacube.group_datasets(all_datasets, self.group_by)
-
-    def _dataset_query(self, index, prod_name, query_args):
-        # ODC Dataset Query
-        if self._product.multi_product:
-            queries = []
-            for pn in prod_name:
-                query_args["product"] = pn
-                queries.append(datacube.api.query.Query(**query_args))
-            _LOG.debug("query start %s", datetime.now().time())
-            datasets = []
-            for q in queries:
-                datasets.extend(index.datasets.search_eager(**q.search_terms))
-            _LOG.debug("query stop %s", datetime.now().time())
+            times = self._times
+        result = mv_search_datasets(index, mode,
+                                  layer=self._product,
+                                  times=times,
+                                  geom=geom,
+                                  mask=mask)
+        if mode == MVSelectOpts.DATASETS:
+            return datacube.Datacube.group_datasets(result, self.group_by)
         else:
-            query = datacube.api.query.Query(**query_args)
-            _LOG.debug("query start %s", datetime.now().time())
-            datasets = index.datasets.search_eager(**query.search_terms)
-            _LOG.debug("query stop %s", datetime.now().time())
-
-        return datasets
+            return result
 
     @log_call
     def data(self, datasets, mask=False, manual_merge=False, skip_corrections=False, **kwargs):
@@ -236,38 +222,26 @@ def get_map(args):
             raise WMSException("Database connectivity failure")
         # Tiling.
         stacker = DataStacker(params.product, params.geobox, params.times, params.resampling, style=params.style)
-        datasets = stacker.datasets(dc.index)
-        n_datasets = datasets_in_xarray(datasets)
         zoomed_out = params.zf < params.product.min_zoom
-        too_many_datasets = (params.product.max_datasets_wms > 0
-                             and n_datasets > params.product.max_datasets_wms
-        )
-        if n_datasets == 0:
+        too_many_datasets = False
+        if not zoomed_out:
+            n_datasets = stacker.datasets(dc.index, mode=MVSelectOpts.COUNT)
+            too_many_datasets = (params.product.max_datasets_wms > 0
+                                 and n_datasets > params.product.max_datasets_wms
+            )
+        if too_many_datasets or zoomed_out:
+            extent  = stacker.datasets(dc.index, mode=MVSelectOpts.EXTENT)
+            if extent is None:
+                body = _write_empty(params.geobox)
+            else:
+                body = _write_polygon(
+                    params.geobox,
+                    extent,
+                    params.product.zoom_fill)
+        elif n_datasets == 0:
             body = _write_empty(params.geobox)
-        elif too_many_datasets:
-            body = _write_polygon(
-                params.geobox,
-                params.geobox.extent,
-                params.product.zoom_fill)
-        elif zoomed_out:
-            # Zoomed out to far to properly render data.
-            # Construct a polygon which is the union of the extents of the matching datasets.
-            extent = None
-            extent_crs = None
-            for dt in datasets.time.values:
-                tds = datasets.sel(time=dt)
-                for ds in tds.values.item():
-                    if extent:
-                        new_extent = bbox_to_geom(ds.extent.boundingbox, ds.extent.crs)
-                        if new_extent.crs != extent_crs:
-                            new_extent = new_extent.to_crs(extent_crs)
-                        extent = extent.union(new_extent)
-                    else:
-                        extent = bbox_to_geom(ds.extent.boundingbox, ds.extent.crs)
-                        extent_crs = extent.crs
-            extent = extent.to_crs(params.crs)
-            body = _write_polygon(params.geobox, extent, params.product.zoom_fill)
         else:
+            datasets = stacker.datasets(dc.index)
             _LOG.debug("load start %s %s", datetime.now().time(), args["requestid"])
             data = stacker.data(datasets,
                                 manual_merge=params.product.data_manual_merge,

datacube_ows/data.py

@@ -0,0 +1,123 @@
+from enum import Enum
+import json
+
+from geoalchemy2 import Geometry
+from datacube.utils.geometry import Geometry as ODCGeom
+from sqlalchemy import MetaData, Table, select, or_, Column, SMALLINT, text
+from psycopg2.extras import DateTimeTZRange
+from sqlalchemy.dialects.postgresql import UUID, TSTZRANGE
+from sqlalchemy.sql.functions import count, func
+
+def get_sqlalc_engine(index):
+    # pylint: disable=protected-access
+    return index._db._engine
+
+def get_st_view(meta):
+    return Table('space_time_view', meta,
+             Column('id', UUID()),
+             Column('dataset_type_ref', SMALLINT()),
+             Column('spatial_extent', Geometry(from_text='ST_GeomFromGeoJSON', name='geometry')),
+             Column('temporal_extent', TSTZRANGE())
+                 )
+_meta = MetaData()
+st_view = get_st_view(_meta)
+
+class MVSelectOpts(Enum):
+    """
+    Enum for mv_search_datasets sel parameter.
+
+    ALL: return all columns, select *, as result set
+    IDS: return list of database_ids only.
+    DATASETS: return list of ODC dataset objects
+    COUNT: return a count of matching datasets
+    EXTENT: return full extent of query as a Geometry
+    """
+    ALL = 0
+    IDS = 1
+    COUNT = 2
+    EXTENT = 3
+    DATASETS = 4
+
+    def sel(self, stv):
+        if self == self.ALL:
+            return [stv]
+        if self == self.IDS or self == self.DATASETS:
+            return [stv.c.id]
+        if self == self.COUNT:
+            return [count(stv.c.id)]
+        if self == self.EXTENT:
+            return [text("ST_AsGeoJSON(ST_Union(spatial_extent))")]
+        assert False
+
+def mv_search_datasets(index,
+                       sel=MVSelectOpts.IDS,
+                       times=None,
+                       layer=None,
+                       geom=None,
+                       mask=False):
+    """
+    Perform a dataset query via the space_time_view
+
+    :param layer: A ows_configuration.OWSNamedLayer object (single or multiproduct)
+    :param index: A datacube index (required)
+
+    :param sel: Selection mode - a MVSelectOpts enum. Defaults to IDS.
+    :param times: A list of pairs of datetimes (with time zone)
+    :param geom: A datacube.utils.geometry.Geometry object
+    :param mask: Bool, if true use the flags product of layer
+
+    :return: See MVSelectOpts doc
+    """
+    engine = get_sqlalc_engine(index)
+    stv = st_view
+    if layer is None:
+        raise Exception("Must filter by product/layer")
+    if mask:
+        prod_ids = [p.id for p in layer.pq_products]
+    else:
+        prod_ids = [p.id for p in layer.products]
+
+    s = select(sel.sel(stv)).where(stv.c.dataset_type_ref.in_(prod_ids))
+    if times is not None:
+        s = s.where(
+            or_(
+                *[
+                    stv.c.temporal_extent.op("&&")(DateTimeTZRange(*t))
+                    for t in times
+                ]
+            )
+        )
+    orig_crs = None
+    if geom is not None:
+        orig_crs = geom.crs
+        if str(geom.crs) != "EPSG:4326":
+            geom = geom.to_crs("EPSG:4326")
+        geom_js = json.dumps(geom.json)
+        s = s.where(stv.c.spatial_extent.intersects(geom_js))
+    # print(s) # Print SQL Statement
+    conn = engine.connect()
+    if sel == MVSelectOpts.ALL:
+        return conn.execute(s)
+    if sel == MVSelectOpts.IDS:
+        return [r[0] for r in conn.execute(s)]
+    if sel in (MVSelectOpts.COUNT, MVSelectOpts.EXTENT):
+        for r in conn.execute(s):
+            if sel == MVSelectOpts.COUNT:
+                return r[0]
+            if sel == MVSelectOpts.EXTENT:
+                geojson = r[0]
+                if geojson is None:
+                    return None
+                uniongeom = ODCGeom(json.loads(geojson), crs="EPSG:4326")
+                if geom:
+                    intersect = uniongeom.intersection(geom)
+                    if orig_crs and orig_crs != "EPSG:4326":
+                        intersect = intersect.to_crs(orig_crs)
+                else:
+                    intersect = uniongeom
+                return intersect
+    if sel == MVSelectOpts.DATASETS:
+        return index.datasets.bulk_get(
+                 [r[0] for r in conn.execute(s)]
+        )
+    assert False

datacube_ows/mv_index.py

@@ -1,3 +1,3 @@
--- Refreshing TIME materialized view
+-- Refreshing TIME materialized view (Blocking)
 
 REFRESH MATERIALIZED VIEW time_view

datacube_ows/sql/extent_views/refresh/002_refresh_time.sql

@@ -1,3 +1,3 @@
--- Refreshing SPACE materialized view
+-- Refreshing SPACE materialized view (blocking)
 
 REFRESH MATERIALIZED VIEW space_view

datacube_ows/sql/extent_views/refresh/003_refresh_space.sql

@@ -0,0 +1,3 @@
+-- Refreshing combined SPACE-TIME materialized view (concurrently)
+
+REFRESH MATERIALIZED VIEW CONCURRENTLY space_time_view

datacube_ows/sql/extent_views/refresh/004_refresh_spacetime.sql

@@ -14,7 +14,7 @@
 
 @click.command()
 @click.option("--views", is_flag=True, default=False, help="Refresh the ODC spatio-temporal materialised views.")
-@click.option("--blocking/--no-blocking", is_flag=True, default=False, help="In conjunction with --views, controls whether the materialised view refresh is blocking or concurrent. (defaults to concurrent)")
+@click.option("--blocking/--no-blocking", is_flag=True, default=False, help="Obsolete")
 @click.option("--schema", is_flag=True, default=False, help="Create or update the OWS database schema, including the spatio-temporal materialised views.")
 @click.option("--role", default=None, help="Role to grant database permissions to")
 @click.option("--summary", is_flag=True, default=False, help="Treat any named ODC products with no corresponding configured OWS Layer as summary products" )
@@ -107,8 +107,11 @@ def main(layers, blocking,
         print("Done")
         return 0
     elif views:
+        if blocking:
+            print("WARNING: The 'blocking' flag is "
+                  "no longer supported and will be ignored.")
         print("Refreshing materialised views...")
-        refresh_views(dc, blocking)
+        refresh_views(dc)
         print("Done")
         return 0
 
@@ -137,14 +140,8 @@ def create_views(dc):
     run_sql(dc, "extent_views/create", database=dbname)
 
 
-def refresh_views(dc, blocking):
-    if blocking:
-        blocking = ""
-        print("N.B. Blocking refresh. Spacetime view will be locked until completion.")
-    else:
-        blocking = "CONCURRENTLY"
-        print("N.B. Concurrent refresh. Refresh will not take place immediately.")
-    run_sql(dc, "extent_views/refresh", blocking=blocking)
+def refresh_views(dc):
+    run_sql(dc, "extent_views/refresh")
 
 
 def create_schema(dc, role):

datacube_ows/sql/extent_views/refresh/004_refresh_spacetime_requires_blocking.sql

@@ -16,6 +16,7 @@
 from datacube_ows.ogc_exceptions import WCS1Exception
 from datacube_ows.ogc_utils import ProductLayerException
 from datacube_ows.ows_configuration import get_config
+from datacube_ows.mv_index import MVSelectOpts
 
 class WCS1GetCoverageRequest():
     version = Version(1,0,0)
@@ -291,9 +292,9 @@ def get_coverage_data(req):
                               req.geobox,
                               req.times,
                               bands=req.bands)
-        datasets = stacker.datasets(dc.index)
-        if not datasets:
-            # TODO: Return an empty coverage file with full metadata?
+        n_datasets = stacker.datasets(dc.index, mode=MVSelectOpts.COUNT)
+        if n_datasets == 0:
+            # Return an empty coverage file with full metadata?
             cfg = get_config()
             x_range = (req.minx, req.maxx)
             y_range = (req.miny, req.maxy)
@@ -336,12 +337,11 @@ def get_coverage_data(req):
 
             return data
 
-        n_datasets = datasets_in_xarray(datasets)
         if req.product.max_datasets_wcs > 0 and n_datasets > req.product.max_datasets_wcs:
             raise WCS1Exception("This request processes too much data to be served in a reasonable amount of time."
                                 "Please reduce the bounds of your request and try again."
                                 "(max: %d, this request requires: %d)" % (req.product.max_datasets_wcs, n_datasets))
-
+        datasets = stacker.datasets(index=dc.index)
         stacker = DataStacker(req.product,
                               req.geobox,
                               req.times,