Handle dangling nodes in links/groups for legacy_to_main

aiida/storage/sqlite_zip/migrations/legacy_to_main.py

@@ -143,7 +143,7 @@ def in_archive_context(_inpath):
     return working / DB_FILENAME
 
 
-def _json_to_sqlite(
+def _json_to_sqlite(  # pylint: disable=too-many-branches,too-many-locals
     outpath: Path, data: dict, node_repos: Dict[str, List[Tuple[str, Optional[str]]]], batch_size: int = 100
 ) -> None:
     """Convert a JSON archive format to SQLite."""
@@ -188,9 +188,17 @@ def _json_to_sqlite(
         if data['links_uuid']:
 
             def _transform_link(link_row):
+                try:
+                    input_id = node_uuid_map[link_row['input']]
+                except KeyError:
+                    raise StorageMigrationError(f'Database contains link with unknown input node: {link_row}')
+                try:
+                    output_id = node_uuid_map[link_row['output']]
+                except KeyError:
+                    raise StorageMigrationError(f'Database contains link with unknown output node: {link_row}')
                 return {
-                    'input_id': node_uuid_map[link_row['input']],
-                    'output_id': node_uuid_map[link_row['output']],
+                    'input_id': input_id,
+                    'output_id': output_id,
                     'label': link_row['label'],
                     'type': link_row['type']
                 }
@@ -207,16 +215,20 @@ def _transform_link(link_row):
                 uuid: pk for uuid, pk in connection.execute(select(v1_schema.DbGroup.uuid, v1_schema.DbGroup.id))  # pylint: disable=unnecessary-comprehension
             }
             length = sum(len(uuids) for uuids in data['groups_uuid'].values())
+            unknown_nodes = {}
             with get_progress_reporter()(desc='Adding Group-Nodes', total=length) as progress:
                 for group_uuid, node_uuids in data['groups_uuid'].items():
                     group_id = group_uuid_map[group_uuid]
-                    connection.execute(
-                        insert(v1_schema.DbGroupNodes.__table__), [{
-                            'dbnode_id': node_uuid_map[uuid],
-                            'dbgroup_id': group_id
-                        } for uuid in node_uuids]
-                    )
+                    rows = []
+                    for uuid in node_uuids:
+                        if uuid in node_uuid_map:
+                            rows.append({'dbnode_id': node_uuid_map[uuid], 'dbgroup_id': group_id})
+                        else:
+                            unknown_nodes.setdefault(group_uuid, set()).add(uuid)
+                    connection.execute(insert(v1_schema.DbGroupNodes.__table__), rows)
                     progress.update(len(node_uuids))
+            if unknown_nodes:
+                MIGRATE_LOGGER.warning(f'Dropped unknown nodes in groups: {unknown_nodes}')
 
 
 def _convert_datetime(key, value):

tests/tools/archive/migration/test_legacy_to_main.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Test archive file migration from legacy format (JSON) to main format (SQLite)."""
+import pytest
+
+from aiida.common.exceptions import StorageMigrationError
+from aiida.storage.sqlite_zip.migrator import migrate
+from tests.utils.archives import get_archive_file
+
+
+def test_dangling_links(tmp_path):
+    """Test that links with node UUIDs that are not in the archive are correctly handled."""
+    filepath_archive = get_archive_file('0.10_dangling_link.aiida', 'export/migrate')
+    with pytest.raises(StorageMigrationError, match='Database contains link with unknown input node'):
+        migrate(filepath_archive, tmp_path / 'archive.aiida', 'main_0001')
+
+
+def test_missing_nodes_in_groups(tmp_path, aiida_caplog):
+    """Test that groups with listed node UUIDs that are not in the archive are correctly handled."""
+    filepath_archive = get_archive_file('0.10_unknown_nodes_in_group.aiida', 'export/migrate')
+    migrate(filepath_archive, tmp_path / 'archive.aiida', 'main_0001')
+    assert 'Dropped unknown nodes in groups' in aiida_caplog.text, aiida_caplog.text