fix rescan backlog fetching

WIPACrepo · Jan 29, 2025 · 1d6c02e · 1d6c02e
1 parent 8491564
commit 1d6c02e
Showing 2 changed files with 17 additions and 12 deletions.
diff --git a/skydriver/k8s/scan_backlog.py b/skydriver/k8s/scan_backlog.py
@@ -69,7 +69,12 @@ async def get_next(
 
         # grab the scan request object--it has other info
         scan_request_obj = await scan_request_client.find_one(  # type: ignore[attr-defined]
-            {"scan_id": manifest.scan_id}
+            {
+                "$or": [
+                    {"scan_id": manifest.scan_id},
+                    {"rescan_ids": manifest.scan_id},  # one in a list
+                ]
+            }
         )
 
         # grab the k8s
@@ -86,7 +91,7 @@ async def run(
     mongo_client: AsyncIOMotorClient,  # type: ignore[valid-type]
     k8s_batch_api: kubernetes.client.BatchV1Api,
     ewms_rc: RestClient,
-        s3_client: botocore.client.BaseClient,
+    s3_client: botocore.client.BaseClient,
 ) -> None:
     """Error-handling around the scan backlog runner loop."""
     LOGGER.info("Started scan backlog runner.")
@@ -108,7 +113,7 @@ async def _run(
     mongo_client: AsyncIOMotorClient,  # type: ignore[valid-type]
     k8s_batch_api: kubernetes.client.BatchV1Api,
     ewms_rc: RestClient,
-        s3_client: botocore.client.BaseClient,
+    s3_client: botocore.client.BaseClient,
 ) -> None:
     """The (actual) main loop."""
     manifest_client = database.interface.ManifestClient(mongo_client)

diff --git a/skydriver/rest_handlers.py b/skydriver/rest_handlers.py
@@ -614,19 +614,19 @@ async def post(self, scan_id: str) -> None:
         # generate unique scan_id
         new_scan_id = uuid.uuid4().hex
 
-        # grab the original requester's 'scan_request_obj'
+        # grab the 'scan_request_obj'
         scan_request_obj = await self.scan_request_coll.find_one_and_update(
-            {"scan_id": scan_id},
+            {
+                "$or": [
+                    # grab the original requester's 'scan_request_obj'
+                    {"scan_id": scan_id},
+                    # -> backup plan: was this scan_id actually a rescan itself?
+                    {"rescan_ids": scan_id},  # one in a list
+                ]
+            },
             {"$push": {"rescan_ids": new_scan_id}},
             return_document=ReturnDocument.AFTER,
         )
-        # -> backup plan: was this scan_id actually a rescan itself?
-        if not scan_request_obj:
-            scan_request_obj = await self.scan_request_coll.find_one_and_update(
-                {"rescan_ids": scan_id},  # one in a list
-                {"$push": {"rescan_ids": new_scan_id}},
-                return_document=ReturnDocument.AFTER,
-            )
         # -> error: couldn't find it anywhere
         if not scan_request_obj:
             raise web.HTTPError(