Repository: opendatacube/datacube-core

Issue:
Add prefix search option for `get_datasets_for_location`
@jeremyh pointed out on slack that `get_datasets_for_location` will be affected by the recent changes to better support stacked datasets (see #415)

Previously we assumed that `dataset.location` points to a file, so reverse query from file to a set of stacked datasets worked by just finding all datasets that have locations with exact match.

https://github.com/opendatacube/datacube-core/blob/5f7211d2eb26ac10d6dfb754418d4bd3079cdf7a/datacube/drivers/postgres/_api.py#L252-L263

Now however every dataset will get a unique location uri, with fragment part (e.g. `file:///some/path/file.nc#part=0`, same applies for multi-doc yaml documents too). So to get the original behaviour back we need a version of that same query but with string prefix match instead of equality check. Any tools that uses this will need to be updated to use prefix search option.

Bonus: this would allow querying a folder as well.




Discussion:

Resolving patch:
diff --git a/datacube/drivers/postgres/_api.py b/datacube/drivers/postgres/_api.py
--- a/datacube/drivers/postgres/_api.py
+++ b/datacube/drivers/postgres/_api.py
@@ -288,16 +288,26 @@ def contains_dataset(self, dataset_id):
             ).fetchone()
         )
 
-    def get_datasets_for_location(self, uri):
+    def get_datasets_for_location(self, uri, mode=None):
         scheme, body = _split_uri(uri)
+
+        if mode is None:
+            mode = 'exact' if body.count('#') > 0 else 'prefix'
+
+        if mode == 'exact':
+            body_query = DATASET_LOCATION.c.uri_body == body
+        elif mode == 'prefix':
+            body_query = DATASET_LOCATION.c.uri_body.startswith(body)
+        else:
+            raise ValueError('Unsupported query mode {}'.format(mode))
+
         return self._connection.execute(
             select(
                 _DATASET_SELECT_FIELDS
             ).select_from(
                 DATASET_LOCATION.join(DATASET)
             ).where(
-                and_(DATASET_LOCATION.c.uri_scheme == scheme,
-                     DATASET_LOCATION.c.uri_body == body)
+                and_(DATASET_LOCATION.c.uri_scheme == scheme, body_query)
             )
         ).fetchall()
 
diff --git a/datacube/index/_datasets.py b/datacube/index/_datasets.py
--- a/datacube/index/_datasets.py
+++ b/datacube/index/_datasets.py
@@ -370,9 +370,9 @@ def add_location(self, id_, uri):
             except DuplicateRecordError:
                 return False
 
-    def get_datasets_for_location(self, uri):
+    def get_datasets_for_location(self, uri, mode=None):
         with self._db.connect() as connection:
-            return (self._make(row) for row in connection.get_datasets_for_location(uri))
+            return (self._make(row) for row in connection.get_datasets_for_location(uri, mode=mode))
 
     def remove_location(self, id_, uri):
         """


Test patch:
diff --git a/integration_tests/index/test_index_data.py b/integration_tests/index/test_index_data.py
--- a/integration_tests/index/test_index_data.py
+++ b/integration_tests/index/test_index_data.py
@@ -352,8 +352,14 @@ def test_index_dataset_with_location(index, default_metadata_type):
     second_ds_doc = copy.deepcopy(_telemetry_dataset)
     second_ds_doc['id'] = '366f32d8-e1f8-11e6-94b4-185e0f80a5c0'
     index.datasets.add(Dataset(type_, second_ds_doc, uris=[second_file.as_uri()], sources={}))
-    dataset_ids = [d.id for d in index.datasets.get_datasets_for_location(first_file.as_uri())]
-    assert dataset_ids == [dataset.id]
+    for mode in ('exact', 'prefix', None):
+        dataset_ids = [d.id for d in index.datasets.get_datasets_for_location(first_file.as_uri(), mode=mode)]
+        assert dataset_ids == [dataset.id]
+
+    assert list(index.datasets.get_datasets_for_location(first_file.as_uri() + "#part=100")) == []
+
+    with pytest.raises(ValueError):
+        list(index.datasets.get_datasets_for_location(first_file.as_uri(), mode="nosuchmode"))
 
 
 def utc_now():