Include usage of filter_size for importing logs, user, computer

Building the query for the matching nodes from the backend can for
surpass the sqlite limit of allowed query parameters and thus result in
failure when importing the archive. In this PR we impement the usage of
the `filter_size` which limits the number of parameters in the query for
importing logs, user and computers.

Author: agoscinski

src/aiida/tools/archive/imports.py

@@ -262,11 +262,18 @@ def _import_users(
     # get matching emails from the backend
     output_email_id: Dict[str, int] = {}
     if input_id_email:
-        output_email_id = dict(
-            orm.QueryBuilder(backend=backend_to)
-            .append(orm.User, filters={'email': {'in': list(input_id_email.values())}}, project=['email', 'id'])
-            .all(batch_size=query_params.batch_size)
-        )
+        output_email_id = {
+            key: value
+            for query_results in [
+                dict(
+                    orm.QueryBuilder(backend=backend_to)
+                    .append(orm.User, filters={'email': {'in': chunk}}, project=['email', 'id'])
+                    .all(batch_size=query_params.batch_size)
+                )
+                for _, chunk in batch_iter(set(input_id_email.values()), query_params.filter_size)
+            ]
+            for key, value in query_results.items()
+        }
 
     new_users = len(input_id_email) - len(output_email_id)
     existing_users = len(output_email_id)
@@ -300,11 +307,18 @@ def _import_computers(
     # get matching uuids from the backend
     backend_uuid_id: Dict[str, int] = {}
     if input_id_uuid:
-        backend_uuid_id = dict(
-            orm.QueryBuilder(backend=backend_to)
-            .append(orm.Computer, filters={'uuid': {'in': list(input_id_uuid.values())}}, project=['uuid', 'id'])
-            .all(batch_size=query_params.batch_size)
-        )
+        backend_uuid_id = {
+            key: value
+            for query_results in [
+                dict(
+                    orm.QueryBuilder(backend=backend_to)
+                    .append(orm.Computer, filters={'uuid': {'in': chunk}}, project=['uuid', 'id'])
+                    .all(batch_size=query_params.batch_size)
+                )
+                for _, chunk in batch_iter(set(input_id_uuid.values()), query_params.filter_size)
+            ]
+            for key, value in query_results.items()
+        }
 
     new_computers = len(input_id_uuid) - len(backend_uuid_id)
     existing_computers = len(backend_uuid_id)
@@ -460,17 +474,20 @@ def _import_nodes(
 
     # get matching uuids from the backend
     backend_uuid_id: Dict[str, int] = {}
-    input_id_uuid_uuids = list(input_id_uuid.values())
 
     if input_id_uuid:
-        for _, batch in batch_iter(input_id_uuid_uuids, query_params.filter_size):
-            backend_uuid_id.update(
+        backend_uuid_id = {
+            key: value
+            for query_results in [
                 dict(
                     orm.QueryBuilder(backend=backend_to)
-                    .append(orm.Node, filters={'uuid': {'in': batch}}, project=['uuid', 'id'])
+                    .append(orm.Node, filters={'uuid': {'in': chunk}}, project=['uuid', 'id'])
                     .all(batch_size=query_params.batch_size)
                 )
-            )
+                for _, chunk in batch_iter(set(input_id_uuid.values()), query_params.filter_size)
+            ]
+            for key, value in query_results.items()
+        }
 
     new_nodes = len(input_id_uuid) - len(backend_uuid_id)
 
@@ -544,12 +561,20 @@ def _import_logs(
 
     # get matching uuids from the backend
     backend_uuid_id: Dict[str, int] = {}
+
     if input_id_uuid:
-        backend_uuid_id = dict(
-            orm.QueryBuilder(backend=backend_to)
-            .append(orm.Log, filters={'uuid': {'in': list(input_id_uuid.values())}}, project=['uuid', 'id'])
-            .all(batch_size=query_params.batch_size)
-        )
+        backend_uuid_id = {
+            key: value
+            for query_results in [
+                dict(
+                    orm.QueryBuilder(backend=backend_to)
+                    .append(orm.Log, filters={'uuid': {'in': chunk}}, project=['uuid', 'id'])
+                    .all(batch_size=query_params.batch_size)
+                )
+                for _, chunk in batch_iter(set(input_id_uuid.values()), query_params.filter_size)
+            ]
+            for key, value in query_results.items()
+        }
 
     new_logs = len(input_id_uuid) - len(backend_uuid_id)
     existing_logs = len(backend_uuid_id)

tests/tools/archive/orm/test_computers.py

@@ -328,3 +328,36 @@ def test_import_of_django_sqla_export_file(aiida_localhost, backend):
     res = builder.dict()[0]
 
     assert res['comp']['metadata'] == comp1_metadata
+
+
+def test_filter_size(tmp_path, aiida_profile_clean):
+    """Tests if the query still works when the number of computer is beyond the `filter_size limit."""
+    nb_nodes = 5
+    nodes = []
+    for i in range(nb_nodes):
+        node = orm.CalcJobNode()
+        node.computer = orm.Computer(
+            label=f'{i}', hostname='localhost', transport_type='core.local', scheduler_type='core.direct'
+        ).store()
+        node.set_option('resources', {'num_machines': 1, 'num_mpiprocs_per_machine': 1})
+        node.label = f'{i}'
+        node.store()
+        node.seal()
+        nodes.append(node)
+
+    builder = orm.QueryBuilder().append(orm.Computer, project=['uuid', 'label'])
+    builder = builder.all()
+
+    # Export DB
+    export_file_existing = tmp_path.joinpath('export.aiida')
+    create_archive(nodes, filename=export_file_existing)
+
+    # Clean database and reimport DB
+    aiida_profile_clean.reset_storage()
+    import_archive(export_file_existing, filter_size=2)
+
+    # Check correct import
+    builder = orm.QueryBuilder().append(orm.Computer, project=['uuid', 'label'])
+    builder = builder.all()
+
+    assert len(builder) == nb_nodes

tests/tools/archive/orm/test_logs.py

@@ -378,3 +378,28 @@ def test_reimport_of_logs_for_single_node(tmp_path, aiida_profile_clean):
         log_message = str(log[1])
         assert log_uuid in total_log_uuids
         assert log_message in log_msgs
+
+
+def test_filter_size(tmp_path, aiida_profile_clean):
+    """Tests if the query still works when the number of logs is beyond the `filter_size limit."""
+    node = orm.CalculationNode().store()
+    node.seal()
+
+    nb_nodes = 5
+    for _ in range(nb_nodes):
+        node.logger.critical('some')
+
+    # Export DB
+    export_file_existing = tmp_path.joinpath('export.aiida')
+    create_archive([node], filename=export_file_existing)
+
+    # Clean database and reimport DB
+    aiida_profile_clean.reset_storage()
+    import_archive(export_file_existing, filter_size=2)
+
+    # Check correct import
+    builder = orm.QueryBuilder().append(orm.Node, tag='node', project=['uuid'])
+    builder.append(orm.Log, with_node='node', project=['uuid'])
+    builder = builder.all()
+
+    assert len(builder) == nb_nodes

tests/tools/archive/orm/test_users.py

@@ -161,3 +161,30 @@ def test_non_default_user_nodes(aiida_profile_clean, tmp_path, aiida_localhost_f
         assert orm.load_node(uuid).user.email == new_email
     for uuid in uuids2:
         assert orm.load_node(uuid).user.email == manager.get_profile().default_user_email
+
+
+def test_filter_size(tmp_path, aiida_profile_clean):
+    """Tests if the query still works when the number of users is beyond the `filter_size limit."""
+    nb_nodes = 5
+    nodes = []
+    # We need to attach a node to the user otherwise it is not exported
+    for i in range(nb_nodes):
+        node = orm.Int(5, user=orm.User(email=f'{i}').store())
+        node.store()
+        nodes.append(node)
+
+    # Export DB
+    export_file_existing = tmp_path.joinpath('export.aiida')
+    create_archive(nodes, filename=export_file_existing)
+
+    # Clean database and reimport DB
+    aiida_profile_clean.reset_storage()
+    import_archive(export_file_existing, filter_size=2)
+
+    # Check correct import
+    builder = orm.QueryBuilder().append(orm.User, project=['id'])
+    builder = builder.all()
+
+    # We need to add one because default profile is added by reset_storage
+    # automatically
+    assert len(builder) == nb_nodes + 1