Fixed unstaging with update by query. Added repair uploads as admin utility.

infrastructure/utils.http

@@ -24,3 +24,23 @@ content-type: application/json
 # Delete the calc index
 
 DELETE http://localhost:9200/calcs HTTP/1.1
+
+###
+
+POST http://localhost:9200/calcs/_update_by_query HTTP/1.1
+content-type: application/json
+
+{
+    "script": {
+        "inline": "ctx._source.staging=false",
+        "lang": "painless"
+    },
+    "query": {
+        "match": {
+            "upload_id": "epGwyPvaQ720_COyRvyLEg"
+        }
+    }
+}
+
+###
+POST http://localhost:8000/nomadxt/api/admin/repair_uploads HTTP/1.1

nomad/api.py

@@ -582,6 +582,15 @@ def get_calc(upload_hash, calc_hash):
         abort(500, message='Could not accessing the archive.')
 
 
+@app.route('%s/admin/<string:operation>' % base_path, methods=['POST'])
+def call_admin_operation(operation):
+    if operation == 'repair_uploads':
+        Upload.repair_all()
+    else:
+        abort(400, message='Unknown operation %s' % operation)
+
+    return 'done', 200
+
 api.add_resource(UploadsRes, '%s/uploads' % base_path)
 api.add_resource(UploadRes, '%s/uploads/<string:upload_id>' % base_path)
 api.add_resource(RepoCalcsRes, '%s/repo' % base_path)

nomad/processing/data.py

@@ -40,7 +40,7 @@ import logging
 from nomad import config, files, utils
 from nomad.repo import RepoCalc
 from nomad.user import User, me
-from nomad.processing.base import Proc, process, task, PENDING, SUCCESS, FAILURE
+from nomad.processing.base import Proc, process, task, PENDING, SUCCESS, FAILURE, RUNNING
 from nomad.parsing import LocalBackend, parsers, parser_dict
 from nomad.normalizing import normalizers
 from nomad.utils import get_logger, lnr
@@ -217,8 +217,7 @@ class Upload(Proc):
 
     meta: Any = {
         'indexes': [
-            'upload_hash',
-            'user_id'
+            'upload_hash', 'user_id', 'status'
         ]
     }
 
@@ -410,5 +409,26 @@ class Upload(Proc):
     def failed_calcs(self):
         return Calc.objects(upload_id=self.upload_id, status=FAILURE).count()
 
+    @property
+    def pending_calcs(self):
+        return Calc.objects(upload_id=self.upload_id, status=PENDING).count()
+
     def all_calcs(self, start, end, order_by='mainfile'):
         return Calc.objects(upload_id=self.upload_id)[start:end].order_by(order_by)
+
+    @staticmethod
+    def repair_all():
+        """
+        Utitlity function that will look for suspiciously looking conditions in
+        all uncompleted downloads. It ain't a perfect world.
+        """
+        uploads = Upload.objects(status__in=[PENDING, RUNNING])
+        for upload in uploads:
+            completed = upload.processed_calcs
+            total = upload.total
+            pending = upload.pending_calcs
+
+            if completed + pending == total:
+                time.sleep(2)
+                if pending == upload.pending_calcs:
+                    Calc.objects(upload_id=upload.upload_id, status=PENDING).delete()

nomad/repo.py

@@ -164,12 +164,36 @@ class RepoCalc(ElasticDocument):
         """ Deletes all repo entries of the given upload. """
         RepoCalc.search().query('match', upload_id=upload_id).delete()
 
+    @classmethod
+    def unstage(cls, upload_id, staging=False):
+        """ Update the staging property for all repo entries of the given upload. """
+        cls.update_by_query(upload_id, {
+            'inline': 'ctx._source.staging=%s' % ('true' if staging else 'false'),
+            'lang': 'painless'
+        })
+
     @classmethod
     def update_upload(cls, upload_id, **kwargs):
         """ Update all entries of given upload with keyword args. """
         for calc in RepoCalc.search().query('match', upload_id=upload_id):
             calc.update(**kwargs)
 
+    @classmethod
+    def update_by_query(cls, upload_id, script):
+        """ Update all entries of a given upload via elastic script. """
+        index = cls._default_index()
+        doc_type = cls._doc_type.name
+        conn = cls._get_connection()
+        body = {
+            'script': script,
+            'query': {
+                'match': {
+                    'upload_id': upload_id
+                }
+            }
+        }
+        response = conn.update_by_query(index, doc_type=[doc_type], body=body)
+
     @staticmethod
     def es_search(body):
         """ Perform an elasticsearch and not elasticsearch_dsl search on the Calc index. """

tests/test_repo.py

@@ -117,6 +117,6 @@ def test_staging_elastic_calc(example_elastic_calc: RepoCalc):
 
 
 def test_unstage_elastic_calc(example_elastic_calc: RepoCalc):
-    RepoCalc.update_upload(upload_id='test_upload_id', staging=False)
+    RepoCalc.unstage(upload_id='test_upload_id', staging=False)
 
     assert not RepoCalc.get(id='test_upload_hash/test_calc_hash').staging