Added msgpacked function in uploads

nomad/cli/admin/uploads.py

@@ -19,9 +19,11 @@ from mongoengine import Q
 from pymongo import UpdateOne
 import threading
 import elasticsearch_dsl as es
+import json
 
 from nomad import processing as proc, config, infrastructure, utils, search, files, datamodel
 from .admin import admin
+from nomad.archive import write_archive
 
 
 @admin.group(help='Upload related commands')
@@ -231,6 +233,30 @@ def rm(ctx, uploads, skip_es, skip_mongo, skip_files):
         delete_upload(upload, skip_es=skip_es, skip_mongo=skip_mongo, skip_files=skip_files)
 
 
+@uploads.command(help='Create msgpack file for upload')
+@click.argument('UPLOADS', nargs=-1)
+@click.pass_context
+def msgpacked(ctx, uploads):
+    _, uploads = query_uploads(ctx, uploads)
+
+    for upload in uploads:
+        upload_files = files.UploadFiles.get(upload_id=upload.upload_id)
+
+        if isinstance(upload_files, files.PublicUploadFiles):
+            def iterator(zf, names):
+                for name in names:
+                    calc_id = name.strip('.json')
+                    with zf.open(name) as f:
+                        yield (calc_id, json.load(f))
+
+            for access in ['public', 'restricted']:
+                zf = upload_files.open_zip_file('archive', access, upload_files._archive_ext)
+                archive_name = zf.filename.replace('.zip', '.msg')
+                names = [name for name in zf.namelist() if name.endswith('json')]
+                print('writing msgpack archive %s' % archive_name)
+                write_archive(archive_name, len(names), iterator(zf, names))
+
+
 def __run_processing(
         ctx, uploads, parallel: int, process: Callable[[proc.Upload], None], label: str):
     _, uploads = query_uploads(ctx, uploads)

tests/test_cli.py

@@ -134,6 +134,15 @@ class TestAdminUploads:
         assert Upload.objects(upload_id=upload_id).first() is None
         assert Calc.objects(upload_id=upload_id).first() is None
 
+    def test_create_msgpack(self, published):
+        upload_id = published.upload_id
+
+        result = click.testing.CliRunner().invoke(
+            cli, ['admin', 'uploads', 'msgpacked', upload_id], catch_exceptions=False, obj=utils.POPO())
+
+        assert result.exit_code == 0
+        assert 'writing msgpack archive' in result.stdout
+
     def test_index(self, published):
         upload_id = published.upload_id
         calc = Calc.objects(upload_id=upload_id).first()