Skip to content

Commit df574cb

Browse files
committed
Include usage of filter_size for _import_logs
Building the query for `backend_uuid_id` does not use the `filter_size` and thus fails in case sqlite is used.
1 parent e768b70 commit df574cb

File tree

2 files changed

+54
-5
lines changed

2 files changed

+54
-5
lines changed

src/aiida/tools/archive/imports.py

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -544,12 +544,31 @@ def _import_logs(
544544

545545
# get matching uuids from the backend
546546
backend_uuid_id: Dict[str, int] = {}
547+
548+
# TODO put somewhere else
549+
def chunk_dict(d, chunk_size):
550+
from itertools import islice
551+
552+
it = iter(d.items())
553+
while True:
554+
chunk = dict(islice(it, chunk_size))
555+
if not chunk:
556+
break
557+
yield chunk
558+
547559
if input_id_uuid:
548-
backend_uuid_id = dict(
549-
orm.QueryBuilder(backend=backend_to)
550-
.append(orm.Log, filters={'uuid': {'in': list(input_id_uuid.values())}}, project=['uuid', 'id'])
551-
.all(batch_size=query_params.batch_size)
552-
)
560+
backend_uuid_id = {
561+
key: value
562+
for lis in [
563+
dict(
564+
orm.QueryBuilder(backend=backend_to)
565+
.append(orm.Log, filters={'uuid': {'in': chunk}}, project=['uuid', 'id'])
566+
.all(batch_size=query_params.batch_size)
567+
)
568+
for _, chunk in batch_iter(set(input_id_uuid.values()), query_params.filter_size)
569+
]
570+
for key, value in lis.items()
571+
}
553572

554573
new_logs = len(input_id_uuid) - len(backend_uuid_id)
555574
existing_logs = len(backend_uuid_id)

tests/tools/archive/orm/test_logs.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,3 +378,33 @@ def test_reimport_of_logs_for_single_node(tmp_path, aiida_profile_clean):
378378
log_message = str(log[1])
379379
assert log_uuid in total_log_uuids
380380
assert log_message in log_msgs
381+
382+
383+
def test_filter_size_logs(tmp_path, aiida_profile_clean):
384+
"""Tests"""
385+
386+
# Create Node and initial log message and save UUIDs prior to export
387+
node = orm.CalculationNode().store()
388+
node.seal()
389+
# To test the query parameter limit of Linux we need a limit above 1000,
390+
# above 2000 for macos
391+
nb_logs = 2002
392+
for _ in range(nb_logs):
393+
node.logger.critical('some')
394+
log_uuid_existing = orm.QueryBuilder().append(orm.Log, project=['uuid']).all()
395+
log_uuid_existing = str(log_uuid_existing[0][0])
396+
397+
# Export DB
398+
export_file_existing = tmp_path.joinpath('export.aiida')
399+
create_archive([node], filename=export_file_existing)
400+
401+
# Clean database and reimport DB
402+
aiida_profile_clean.reset_storage()
403+
import_archive(export_file_existing)
404+
405+
# Check correct import
406+
builder = orm.QueryBuilder().append(orm.Node, tag='node', project=['uuid'])
407+
builder.append(orm.Log, with_node='node', project=['uuid', 'message'])
408+
builder = builder.all()
409+
410+
assert len(builder) == nb_logs

0 commit comments

Comments
 (0)