Repository: opendatacube/datacube-core
Issue: Doc2Dataset fails if lineage sources are specified as UUIDs rather than strings

Converting an input doc that contains lineage references to a Dataset (using `Doc2Dataset(index, fail_on_missing_lineage=True, verify_lineage=False, skip_lineage=False)`) fails if the lineage references in the input doc are of type `UUID` instead of `str`.

In [dataset_resolver, hl.py around line 174](https://github.com/opendatacube/datacube-core/blob/c5486c854dc4fe3e1bef670a5ce01010d0caa938/datacube/index/hl.py#L174), the identifiers of the resolved lineage sources are converted to strings. Then, missing lineage items are determined by subtracting the resolved identifiers from the input set of lineage sources. But if the input references are of type `UUID`, this simple set subtraction is incorrect, and as a result the lineage items will be considered missing: 
```python
        db_dss = {str(ds.id): ds for ds in index.datasets.bulk_get(all_uuid)}

        lineage_uuids = set(filter(lambda x: x != main_uuid, all_uuid))
        missing_lineage = lineage_uuids - set(db_dss)
```
(After that, also generating the error message in line 180 fails, because it expects missing_lineage items to be strings).

In practice, it is very likely that lineage source identifiers will be of type `UUID`, given that the `Dataset` class uses `UUID`s for identifiers, as do external interfaces like e.g. `eodatasets3.DatasetPrepare.note_source_datasets()`.
Before doing the subtraction, both lineage_uuids and db_dss sets should be normalized to contain a single datatype.


Changed files:
datacube/model/utils.py, datacube/index/eo3.py, datacube/index/hl.py, integration_tests/test_dataset_add.py, tests/test_utils_docs.py, integration_tests/index/test_memory_index.py, datacube/utils/documents.py

diff --git a/datacube/index/eo3.py b/datacube/index/eo3.py
index 12aea720..74bb06c6 100644
--- a/datacube/index/eo3.py
+++ b/datacube/index/eo3.py
@@ -9,7 +9,8 @@
 """
 from affine import Affine
 from functools import reduce
-from typing import Dict, Any, Iterable, Optional, Tuple
+from typing import Dict, Any, Iterable, Optional, Tuple, Union
+from uuid import UUID
 
 from datacube.utils.geometry import (
     SomeCRS,
@@ -217,6 +218,11 @@ def prep_eo3(doc: Dict[str, Any],
         if not is_doc_eo3(doc):
             return doc
 
+    def stringify(u: Optional[Union[str, UUID]]) -> Optional[str]:
+        return u if isinstance(u, str) else str(u) if u else None
+
+    doc['id'] = stringify(doc.get('id', None))
+
     doc = add_eo3_parts(doc, resolution=resolution)
     lineage = doc.pop('lineage', {})
 
@@ -228,11 +234,11 @@ def prep_eo3(doc: Dict[str, Any],
         if isinstance(uuids, dict) or isinstance(uuids[0], dict):
             raise ValueError("Embedded lineage not supported for eo3 metadata types")
         if len(uuids) == 1:
-            return {name: {'id': uuids[0]}}
+            return {name: {'id': stringify(uuids[0])}}
 
         out = {}
         for idx, uuid in enumerate(uuids, start=1):
-            out[name+str(idx)] = {'id': uuid}
+            out[name+str(idx)] = {'id': stringify(uuid)}
         return out
 
     sources = {}
diff --git a/datacube/index/hl.py b/datacube/index/hl.py
index 6a86ec60..6d09ab4f 100644
--- a/datacube/index/hl.py
+++ b/datacube/index/hl.py
@@ -171,13 +171,14 @@ def dataset_resolver(index: AbstractIndex,
 
         ds_by_uuid = toolz.valmap(toolz.first, flatten_datasets(main_ds))
         all_uuid = list(ds_by_uuid)
-        db_dss = {str(ds.id): ds for ds in index.datasets.bulk_get(all_uuid)}
+        db_dss = {ds.id: ds for ds in index.datasets.bulk_get(all_uuid)}
 
         lineage_uuids = set(filter(lambda x: x != main_uuid, all_uuid))
         missing_lineage = lineage_uuids - set(db_dss)
 
         if missing_lineage and fail_on_missing_lineage:
-            return None, "Following lineage datasets are missing from DB: %s" % (','.join(missing_lineage))
+            return None, "Following lineage datasets are missing from DB: %s" % (
+                ','.join(str(m) for m in missing_lineage))
 
         if not is_doc_eo3(main_ds.doc):
             if is_doc_geo(main_ds.doc, check_eo3=False):
@@ -206,7 +207,7 @@ def dataset_resolver(index: AbstractIndex,
             return v
 
         def resolve_ds(ds: SimpleDocNav,
-                       sources: Optional[Mapping[str, Dataset]],
+                       sources: Optional[Mapping[UUID, Dataset]],
                        cache: MutableMapping[UUID, Dataset]) -> Dataset:
             cached = cache.get(ds.id)
             if cached is not None:
diff --git a/datacube/model/utils.py b/datacube/model/utils.py
index 75409427..90746d51 100644
--- a/datacube/model/utils.py
+++ b/datacube/model/utils.py
@@ -353,7 +353,7 @@ def remap_lineage_doc(root, mk_node, **kwargs):
     try:
         return visit(root)
     except BadMatch as e:
-        if root.id not in str(e):
+        if str(root.id) not in str(e):
             raise BadMatch(f"Error loading lineage dataset: {e}") from None
         else:
             raise
diff --git a/datacube/utils/documents.py b/datacube/utils/documents.py
index 635658e6..34dcbdb1 100644
--- a/datacube/utils/documents.py
+++ b/datacube/utils/documents.py
@@ -18,6 +18,7 @@ from urllib.parse import urlparse
 from urllib.request import urlopen
 from typing import Dict, Any, Mapping
 from copy import deepcopy
+from uuid import UUID
 
 import numpy
 import toolz  # type: ignore[import]
@@ -372,6 +373,7 @@ class SimpleDocNav(object):
         self._doc_without = None
         self._sources_path = ('lineage', 'source_datasets')
         self._sources = None
+        self._doc_uuid = None
 
     @property
     def doc(self):
@@ -386,7 +388,11 @@ class SimpleDocNav(object):
 
     @property
     def id(self):
-        return self._doc.get('id', None)
+        if not self._doc_uuid:
+            doc_id = self._doc.get('id', None)
+            if doc_id:
+                self._doc_uuid = doc_id if isinstance(doc_id, UUID) else UUID(doc_id)
+        return self._doc_uuid
 
     @property
     def sources(self):
diff --git a/integration_tests/index/test_memory_index.py b/integration_tests/index/test_memory_index.py
index 640a51d5..97b276c1 100644
--- a/integration_tests/index/test_memory_index.py
+++ b/integration_tests/index/test_memory_index.py
@@ -5,7 +5,6 @@
 import datetime
 
 import pytest
-from uuid import UUID
 from datacube.testutils import gen_dataset_test_dag
 
 from datacube.utils import InvalidDocException, read_documents, SimpleDocNav
@@ -549,7 +548,7 @@ def test_memory_dataset_add(dataset_add_configs, mem_index_fresh):
     for id_ in ds_ids:
         assert id_ in loc_ids
     ds_ = SimpleDocNav(gen_dataset_test_dag(1, force_tree=True))
-    assert UUID(ds_.id) in ds_ids
+    assert ds_.id in ds_ids
     ds_from_idx = idx.datasets.get(ds_.id, include_sources=True)
-    assert str(ds_from_idx.sources['ab'].id) == ds_.sources['ab'].id
-    assert str(ds_from_idx.sources['ac'].sources["cd"].id) == ds_.sources['ac'].sources['cd'].id
+    assert ds_from_idx.sources['ab'].id == ds_.sources['ab'].id
+    assert ds_from_idx.sources['ac'].sources["cd"].id == ds_.sources['ac'].sources['cd'].id
diff --git a/integration_tests/test_dataset_add.py b/integration_tests/test_dataset_add.py
index b235923b..5d8e8744 100644
--- a/integration_tests/test_dataset_add.py
+++ b/integration_tests/test_dataset_add.py
@@ -25,7 +25,7 @@ def check_skip_lineage_test(clirunner, index):
 
     ds_ = index.datasets.get(ds.id, include_sources=True)
     assert ds_ is not None
-    assert str(ds_.id) == ds.id
+    assert ds_.id == ds.id
     assert ds_.sources == {}
 
     assert index.datasets.get(ds.sources['ab'].id) is None
@@ -234,31 +234,31 @@ def test_dataset_add(dataset_add_configs, index_empty, clirunner):
     doc2ds = Doc2Dataset(index)
     _ds, _err = doc2ds(ds.doc, 'file:///something')
     assert _err is None
-    assert str(_ds.id) == ds.id
+    assert _ds.id == ds.id
     assert _ds.metadata_doc == ds.doc
 
     # Check dataset search
     r = clirunner(['dataset', 'search'], expect_success=True)
-    assert ds.id in r.output
-    assert ds_bad1.id not in r.output
-    assert ds.sources['ab'].id in r.output
-    assert ds.sources['ac'].sources['cd'].id in r.output
+    assert str(ds.id) in r.output
+    assert str(ds_bad1.id) not in r.output
+    assert str(ds.sources['ab'].id) in r.output
+    assert str(ds.sources['ac'].sources['cd'].id) in r.output
 
-    r = clirunner(['dataset', 'info', '-f', 'csv', ds.id])
-    assert ds.id in r.output
+    r = clirunner(['dataset', 'info', '-f', 'csv', str(ds.id)])
+    assert str(ds.id) in r.output
 
-    r = clirunner(['dataset', 'info', '-f', 'yaml', '--show-sources', ds.id])
-    assert ds.sources['ae'].id in r.output
+    r = clirunner(['dataset', 'info', '-f', 'yaml', '--show-sources', str(ds.id)])
+    assert str(ds.sources['ae'].id) in r.output
 
-    r = clirunner(['dataset', 'info', '-f', 'yaml', '--show-derived', ds.sources['ae'].id])
-    assert ds.id in r.output
+    r = clirunner(['dataset', 'info', '-f', 'yaml', '--show-derived', str(ds.sources['ae'].id)])
+    assert str(ds.id) in r.output
 
     ds_ = SimpleDocNav(gen_dataset_test_dag(1, force_tree=True))
     assert ds_.id == ds.id
 
     x = index.datasets.get(ds.id, include_sources=True)
-    assert str(x.sources['ab'].id) == ds.sources['ab'].id
-    assert str(x.sources['ac'].sources['cd'].id) == ds.sources['ac'].sources['cd'].id
+    assert x.sources['ab'].id == ds.sources['ab'].id
+    assert x.sources['ac'].sources['cd'].id == ds.sources['ac'].sources['cd'].id
 
     check_skip_lineage_test(clirunner, index)
     check_no_product_match(clirunner, index)
@@ -453,11 +453,11 @@ measurements:
     print(r.output)
 
     r = clirunner(['dataset', 'search', '-f', 'csv'])
-    assert ds1.id not in r.output
-    assert ds2.id not in r.output
-    assert ds3.id not in r.output
-    assert ds4.id in r.output
-    assert ds5.id in r.output
+    assert str(ds1.id) not in r.output
+    assert str(ds2.id) not in r.output
+    assert str(ds3.id) not in r.output
+    assert str(ds4.id) in r.output
+    assert str(ds5.id) in r.output
 
 
 def dataset_archive_prep(dataset_add_configs, index_empty, clirunner):
@@ -483,8 +483,8 @@ def test_dataset_archive_dry_run(dataset_add_configs, index_empty, clirunner):
     non_existent_uuid = '00000000-1036-5607-a62f-fde5e3fec985'
 
     # Single valid UUID is detected and not archived
-    single_valid_uuid = clirunner(['dataset', 'archive', '--dry-run', ds.id])
-    assert ds.id in single_valid_uuid.output
+    single_valid_uuid = clirunner(['dataset', 'archive', '--dry-run', str(ds.id)])
+    assert str(ds.id) in single_valid_uuid.output
     assert index.datasets.has(ds.id) is True
 
     # Single invalid UUID is detected
@@ -498,7 +498,7 @@ def test_dataset_archive_dry_run(dataset_add_configs, index_empty, clirunner):
     valid_and_invalid_uuid = clirunner(['dataset',
                                         'archive',
                                         '--dry-run',
-                                        ds.id,
+                                        str(ds.id),
                                         non_existent_uuid],
                                        expect_success=False)
     assert non_existent_uuid in valid_and_invalid_uuid.output
@@ -509,7 +509,7 @@ def test_dataset_archive_dry_run(dataset_add_configs, index_empty, clirunner):
                                         'archive',
                                         '--dry-run',
                                         '--archive-derived',
-                                        ds.id,
+                                        str(ds.id),
                                         non_existent_uuid
                                         ],
                                        expect_success=False)
@@ -519,9 +519,11 @@ def test_dataset_archive_dry_run(dataset_add_configs, index_empty, clirunner):
 
     # Multiple Valid UUIDs
     # Not archived in the database and are shown in output
-    multiple_valid_uuid = clirunner(['dataset', 'archive', '--dry-run', ds.sources['ae'].id, ds.sources['ab'].id])
-    assert ds.sources['ae'].id in multiple_valid_uuid.output
-    assert ds.sources['ab'].id in multiple_valid_uuid.output
+    multiple_valid_uuid = clirunner([
+        'dataset', 'archive', '--dry-run',
+        str(ds.sources['ae'].id), str(ds.sources['ab'].id)])
+    assert str(ds.sources['ae'].id) in multiple_valid_uuid.output
+    assert str(ds.sources['ab'].id) in multiple_valid_uuid.output
     assert index.datasets.has(ds.sources['ae'].id) is True
     assert index.datasets.has(ds.sources['ab'].id) is True
 
@@ -529,11 +531,11 @@ def test_dataset_archive_dry_run(dataset_add_configs, index_empty, clirunner):
                                  'archive',
                                  '--dry-run',
                                  '--archive-derived',
-                                 ds.sources['ae'].id,
-                                 ds.sources['ab'].id
+                                 str(ds.sources['ae'].id),
+                                 str(ds.sources['ab'].id)
                                  ])
-    assert ds.id in archive_derived.output
-    assert ds.sources['ae'].id in archive_derived.output
+    assert str(ds.id) in archive_derived.output
+    assert str(ds.sources['ae'].id) in archive_derived.output
     assert index.datasets.has(ds.id) is True
 
 
@@ -545,15 +547,15 @@ def test_dataset_archive_restore_invalid(dataset_add_configs, index_empty, cliru
     non_existent_uuid = '00000000-1036-5607-a62f-fde5e3fec985'
 
     # With non-existent uuid, operations should halt.
-    r = clirunner(['dataset', 'archive', ds.id, non_existent_uuid], expect_success=False)
-    r = clirunner(['dataset', 'info', ds.id])
+    r = clirunner(['dataset', 'archive', str(ds.id), non_existent_uuid], expect_success=False)
+    r = clirunner(['dataset', 'info', str(ds.id)])
     assert 'status: archived' not in r.output
     assert index.datasets.has(ds.id) is True
 
     # With non-existent uuid, operations should halt.
     d_id = ds.sources['ac'].sources['cd'].id
-    r = clirunner(['dataset', 'archive', '--archive-derived', d_id, non_existent_uuid], expect_success=False)
-    r = clirunner(['dataset', 'info', ds.id, ds.sources['ab'].id, ds.sources['ac'].id])
+    r = clirunner(['dataset', 'archive', '--archive-derived', str(d_id), non_existent_uuid], expect_success=False)
+    r = clirunner(['dataset', 'info', str(ds.id), str(ds.sources['ab'].id), str(ds.sources['ac'].id)])
     assert 'status: active' in r.output
     assert 'status: archived' not in r.output
     assert index.datasets.has(ds.id) is True
@@ -565,30 +567,30 @@ def test_dataset_archive_restore(dataset_add_configs, index_empty, clirunner):
     p, index, ds = dataset_archive_prep(dataset_add_configs, index_empty, clirunner)
 
     # Run for real
-    r = clirunner(['dataset', 'archive', ds.id])
-    r = clirunner(['dataset', 'info', ds.id])
+    r = clirunner(['dataset', 'archive', str(ds.id)])
+    r = clirunner(['dataset', 'info', str(ds.id)])
     assert 'status: archived' in r.output
 
     # restore dry run
-    r = clirunner(['dataset', 'restore', '--dry-run', ds.id])
-    r = clirunner(['dataset', 'info', ds.id])
+    r = clirunner(['dataset', 'restore', '--dry-run', str(ds.id)])
+    r = clirunner(['dataset', 'info', str(ds.id)])
     assert 'status: archived' in r.output
 
     # restore for real
-    r = clirunner(['dataset', 'restore', ds.id])
-    r = clirunner(['dataset', 'info', ds.id])
+    r = clirunner(['dataset', 'restore', str(ds.id)])
+    r = clirunner(['dataset', 'info', str(ds.id)])
     assert 'status: active' in r.output
 
     # archive derived
     d_id = ds.sources['ac'].sources['cd'].id
-    r = clirunner(['dataset', 'archive', '--archive-derived', d_id])
-    r = clirunner(['dataset', 'info', ds.id, ds.sources['ab'].id, ds.sources['ac'].id])
+    r = clirunner(['dataset', 'archive', '--archive-derived', str(d_id)])
+    r = clirunner(['dataset', 'info', str(ds.id), str(ds.sources['ab'].id), str(ds.sources['ac'].id)])
     assert 'status: active' not in r.output
     assert 'status: archived' in r.output
 
     # restore derived
-    r = clirunner(['dataset', 'restore', '--restore-derived', d_id])
-    r = clirunner(['dataset', 'info', ds.id, ds.sources['ab'].id, ds.sources['ac'].id])
+    r = clirunner(['dataset', 'restore', '--restore-derived', str(d_id)])
+    r = clirunner(['dataset', 'info', str(ds.id), str(ds.sources['ab'].id), str(ds.sources['ac'].id)])
     assert 'status: active' in r.output
     assert 'status: archived' not in r.output
 
diff --git a/tests/test_utils_docs.py b/tests/test_utils_docs.py
index 74a523b9..a09db82d 100644
--- a/tests/test_utils_docs.py
+++ b/tests/test_utils_docs.py
@@ -12,6 +12,7 @@ from pathlib import Path
 from collections import OrderedDict
 from types import SimpleNamespace
 from typing import Tuple, Iterable
+from uuid import UUID, uuid4
 
 import numpy as np
 import pytest
@@ -237,7 +238,7 @@ def test_dataset_maker():
 
     assert a.id != b.id
     assert a.doc['creation_dt'] == b.doc['creation_dt']
-    assert isinstance(a.id, str)
+    assert isinstance(a.id, UUID)
     assert a.sources == {}
 
     a1, a2 = [dataset_maker(i)('A', product_type='eo') for i in (0, 1)]
@@ -321,8 +322,11 @@ def test_simple_doc_nav():
       +--> E
     """
 
+    nu_map = {n: uuid4() for n in ['A', 'B', 'C', 'D', 'E']}
+    un_map = {u: n for n, u in nu_map.items()}
+
     def node(name, **kwargs):
-        return dict(id=name, lineage=dict(source_datasets=kwargs))
+        return dict(id=nu_map[name], lineage=dict(source_datasets=kwargs))
 
     A, _, C, _, _ = make_graph_abcde(node)  # noqa: N806
     rdr = SimpleDocNav(A)
@@ -336,7 +340,7 @@ def test_simple_doc_nav():
     assert isinstance(rdr.sources_path, tuple)
 
     def visitor(node, name=None, depth=0, out=None):
-        s = '{}:{}:{:d}'.format(node.id, name if name else '..', depth)
+        s = '{}:{}:{:d}'.format(un_map[node.id], name if name else '..', depth)
         out.append(s)
 
     expect_preorder = '''
@@ -367,17 +371,17 @@ A:..:0
 
     fv = flatten_datasets(rdr)
 
-    assert len(fv['A']) == 1
-    assert len(fv['C']) == 2
-    assert len(fv['E']) == 1
-    assert set(fv.keys()) == set('ABCDE')
+    assert len(fv[nu_map['A']]) == 1
+    assert len(fv[nu_map['C']]) == 2
+    assert len(fv[nu_map['E']]) == 1
+    assert set(fv.keys()) == set(un_map.keys())
 
     fv, dg = flatten_datasets(rdr, with_depth_grouping=True)
 
-    assert len(fv['A']) == 1
-    assert len(fv['C']) == 2
-    assert len(fv['E']) == 1
-    assert set(fv.keys()) == set('ABCDE')
+    assert len(fv[nu_map['A']]) == 1
+    assert len(fv[nu_map['C']]) == 2
+    assert len(fv[nu_map['E']]) == 1
+    assert set(fv.keys()) == set(un_map.keys())
     assert isinstance(dg, list)
     assert len(dg) == 4
     assert [len(dss) for dss in dg] == [1, 3, 2, 1]
@@ -385,10 +389,9 @@ A:..:0
     def to_set(xx):
         return set(x.id for x in xx)
 
-    assert [set(s) for s in ('A',
-                             'BCE',
-                             'CD',
-                             'D')] == [to_set(xx) for xx in dg]
+    assert [set(nu_map[n] for n in s)
+            for s in ('A', 'BCE', 'CD', 'D')
+            ] == [to_set(xx) for xx in dg]
 
     with pytest.raises(ValueError):
         SimpleDocNav([])