Batch remaining functions for archive creation

Author: GeigerJ2

src/aiida/tools/archive/create.py

@@ -16,12 +16,11 @@
 import tempfile
 from datetime import datetime
 from pathlib import Path
-from typing import Callable, Iterable, Optional, Sequence, Union
+from typing import Callable, Iterable, Optional, Union
 
 from tabulate import tabulate
 
 from aiida import orm
-from aiida.common.exceptions import LicensingException
 from aiida.common.lang import type_check
 from aiida.common.links import GraphTraversalRules
 from aiida.common.log import AIIDA_LOGGER
@@ -42,17 +41,18 @@
 EXPORT_LOGGER = AIIDA_LOGGER.getChild('export')
 QbType = Callable[[], orm.QueryBuilder]
 
+
 def _batch_query_ids(
     querybuilder: QbType,
     source_entity_type: EntityTypes,
     source_ids: set[int],
     target_entity_type: EntityTypes,
     relationship: str,
     batch_size_limit: int = 10000,
-    query_batch_size: int = 1000
+    query_batch_size: int = 1000,
 ) -> set[int]:
     """Batch query to avoid PostgreSQL parameter limits.
-    
+
     :param querybuilder: QueryBuilder factory function
     :param source_entity_type: Entity type to filter on
     :param source_ids: Set of IDs to filter
@@ -81,7 +81,7 @@ def _batch_query_ids(
     source_id_list = list(source_ids)
 
     for i in range(0, len(source_id_list), batch_size_limit):
-        batch_ids = source_id_list[i:i + batch_size_limit]
+        batch_ids = source_id_list[i : i + batch_size_limit]
 
         qb = querybuilder()
         qb.append(source_orm, filters={'id': {'in': batch_ids}}, tag='source')
@@ -356,12 +356,14 @@ def querybuilder():
             with get_progress_reporter()(desc='Archiving database: ', total=sum(entity_counts.values())) as progress:
                 for etype, ids in entity_ids.items():
                     if etype == EntityTypes.NODE and strip_checkpoints:
+
                         def transform(row):
                             data = row['entity']
                             if data.get('node_type', '').startswith('process.'):
                                 data['attributes'].pop(orm.ProcessNode.CHECKPOINT_KEY, None)
                             return data
                     else:
+
                         def transform(row):
                             return row['entity']
 
@@ -370,14 +372,14 @@ def transform(row):
                         # Batch the IDs to avoid parameter limits
                         ids_list = list(ids)
                         for i in range(0, len(ids_list), 10000):  # Batch of 10k to stay under 32k limit
-                            batch_ids = ids_list[i:i + 10000]
+                            batch_ids = ids_list[i : i + 10000]
                             for nrows, rows in batch_iter(
                                 querybuilder()
                                 .append(
                                     entity_type_to_orm[etype],
                                     filters={'id': {'in': batch_ids}},  # Now only 10k parameters max
                                     tag='entity',
-                                    project=['**']
+                                    project=['**'],
                                 )
                                 .iterdict(batch_size=batch_size),
                                 batch_size,
@@ -569,7 +571,7 @@ def progress_str(name):
         if entity_ids[EntityTypes.GROUP]:
             group_id_list = list(entity_ids[EntityTypes.GROUP])
             for i in range(0, len(group_id_list), 10000):  # Batch to avoid parameter limits
-                batch_ids = group_id_list[i:i + 10000]
+                batch_ids = group_id_list[i : i + 10000]
                 qbuilder = querybuilder()
                 qbuilder.append(orm.Group, filters={'id': {'in': batch_ids}}, project='id', tag='group')
                 qbuilder.append(orm.Node, with_group='group', project='id')
@@ -600,7 +602,7 @@ def progress_str(name):
                     EntityTypes.COMPUTER,
                     'with_node',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
 
@@ -616,7 +618,7 @@ def progress_str(name):
                     EntityTypes.AUTHINFO,
                     'with_computer',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
 
@@ -632,7 +634,7 @@ def progress_str(name):
                     EntityTypes.LOG,
                     'with_node',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
 
@@ -648,7 +650,7 @@ def progress_str(name):
                     EntityTypes.COMMENT,
                     'with_node',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
 
@@ -664,7 +666,7 @@ def progress_str(name):
                     EntityTypes.USER,
                     'with_node',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
         if entity_ids[EntityTypes.GROUP]:
@@ -676,7 +678,7 @@ def progress_str(name):
                     EntityTypes.USER,
                     'with_group',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
         if entity_ids[EntityTypes.COMMENT]:
@@ -688,7 +690,7 @@ def progress_str(name):
                     EntityTypes.USER,
                     'with_comment',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
         if entity_ids[EntityTypes.AUTHINFO]:
@@ -700,7 +702,7 @@ def progress_str(name):
                     EntityTypes.USER,
                     'with_authinfo',
                     batch_size_limit=10000,
-                    query_batch_size=batch_size
+                    query_batch_size=batch_size,
                 )
             )
 
@@ -713,55 +715,81 @@ def _stream_repo_files(
     key_format: str, writer: ArchiveWriterAbstract, node_ids: set[int], backend: StorageBackend, batch_size: int
 ) -> None:
     """Collect all repository object keys from the nodes, then stream the files to the archive."""
-    keys = set(
-        orm.Node.get_collection(backend).iter_repo_keys(filters={'id': {'in': list(node_ids)}}, batch_size=batch_size)
-    )
+
+    # Batch the node IDs to avoid parameter limits when getting repo keys
+    node_ids_list = list(node_ids)
+    batch_size_limit = 10000  # Stay well under 65535 parameter limit
+    all_keys = set()
+
+    for i in range(0, len(node_ids_list), batch_size_limit):
+        batch_ids = node_ids_list[i : i + batch_size_limit]
+        batch_keys = set(
+            orm.Node.get_collection(backend).iter_repo_keys(filters={'id': {'in': batch_ids}}, batch_size=batch_size)
+        )
+        all_keys.update(batch_keys)
 
     repository = backend.get_repository()
     if not repository.key_format == key_format:
         # Here we would have to go back and replace all the keys in the `BackendNode.repository_metadata`s
         raise NotImplementedError(
             f'Backend repository key format incompatible: {repository.key_format!r} != {key_format!r}'
         )
-    with get_progress_reporter()(desc='Archiving files: ', total=len(keys)) as progress:
-        for key, stream in repository.iter_object_streams(keys):  # type: ignore[arg-type]
+    with get_progress_reporter()(desc='Archiving files: ', total=len(all_keys)) as progress:
+        for key, stream in repository.iter_object_streams(all_keys):  # type: ignore[arg-type]
             # to-do should we use assume the key here is correct, or always re-compute and check?
             writer.put_object(stream, key=key)
             progress.update()
 
 
 def _check_unsealed_nodes(querybuilder: QbType, node_ids: set[int], batch_size: int) -> None:
     """Check no process nodes are unsealed, i.e. all processes have completed."""
-    qbuilder = (
-        querybuilder()
-        .append(
-            orm.ProcessNode,
-            filters={
-                'id': {'in': list(node_ids)},
-                'attributes.sealed': {
-                    '!in': [True]  # better operator?
+    if not node_ids:
+        return
+
+    # Batch the node IDs to avoid parameter limits
+    node_ids_list = list(node_ids)
+    batch_size_limit = 10000  # Stay well under 65535 parameter limit
+    all_unsealed_pks = []
+
+    for i in range(0, len(node_ids_list), batch_size_limit):
+        batch_ids = node_ids_list[i : i + batch_size_limit]
+
+        qbuilder = (
+            querybuilder()
+            .append(
+                orm.ProcessNode,
+                filters={
+                    'id': {'in': batch_ids},
+                    'attributes.sealed': {
+                        '!in': [True]  # better operator?
+                    },
                 },
-            },
-            project='id',
+                project='id',
+            )
+            .distinct()
         )
-        .distinct()
-    )
-    unsealed_node_pks = qbuilder.all(batch_size=batch_size, flat=True)
-    if unsealed_node_pks:
+        batch_unsealed_pks = qbuilder.all(batch_size=batch_size, flat=True)
+        all_unsealed_pks.extend(batch_unsealed_pks)
+
+    if all_unsealed_pks:
         raise ExportValidationError(
             'All ProcessNodes must be sealed before they can be exported. '
-            f"Node(s) with PK(s): {', '.join(str(pk) for pk in unsealed_node_pks)} is/are not sealed."
+            f"Node(s) with PK(s): {', '.join(str(pk) for pk in all_unsealed_pks)} is/are not sealed."
         )
 
 
 def _check_node_licenses(
     querybuilder: QbType,
     node_ids: set[int],
-    allowed_licenses: Union[None, Sequence[str], Callable],
-    forbidden_licenses: Union[None, Sequence[str], Callable],
+    allowed_licenses: Optional[Union[list, Callable]],
+    forbidden_licenses: Optional[Union[list, Callable]],
     batch_size: int,
 ) -> None:
     """Check the nodes to be archived for disallowed licences."""
+    from typing import Sequence
+
+    from aiida.common.exceptions import LicensingException
+
     if allowed_licenses is None and forbidden_licenses is None:
         return None
 
@@ -807,24 +835,31 @@ def check_forbidden(lic):
     else:
         raise TypeError('forbidden_licenses not a list or function')
 
-    # create query
-    qbuilder = querybuilder().append(
-        orm.Node,
-        project=['id', 'attributes.source.license'],
-        filters={'id': {'in': list(node_ids)}},
-    )
+    # Batch the node IDs to avoid parameter limits
+    node_ids_list = list(node_ids)
+    batch_size_limit = 10000  # Stay well under 65535 parameter limit
 
-    for node_id, name in qbuilder.iterall(batch_size=batch_size):
-        if name is None:
-            continue
-        if not check_allowed(name):
-            raise LicensingException(
-                f"Node {node_id} is licensed under '{name}' license, which is not in the list of allowed licenses"
-            )
-        if check_forbidden(name):
-            raise LicensingException(
-                f"Node {node_id} is licensed under '{name}' license, which is in the list of forbidden licenses"
-            )
+    for i in range(0, len(node_ids_list), batch_size_limit):
+        batch_ids = node_ids_list[i : i + batch_size_limit]
+
+        # create query for this batch
+        qbuilder = querybuilder().append(
+            orm.Node,
+            project=['id', 'attributes.source.license'],
+            filters={'id': {'in': batch_ids}},
+        )
+
+        for node_id, name in qbuilder.iterall(batch_size=batch_size):
+            if name is None:
+                continue
+            if not check_allowed(name):
+                raise LicensingException(
+                    f"Node {node_id} is licensed under '{name}' license, which is not in the list of allowed licenses"
+                )
+            if check_forbidden(name):
+                raise LicensingException(
+                    f"Node {node_id} is licensed under '{name}' license, which is in the list of forbidden licenses"
+                )
 
 
 def get_init_summary(
@@ -856,4 +891,3 @@ def get_init_summary(
         result += f"\n\n{tabulate(rules_table, headers=['Traversal rules', ''])}"
 
     return result + '\n'
-