Exports can be filtered by prefix (#1237)

CERT-Polska · Aug 28, 2024 · e4a720b · e4a720b
1 parent 7b8f035
commit e4a720b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 5 deletions.
diff --git a/artemis/api.py b/artemis/api.py
@@ -148,7 +148,7 @@ def stop_and_delete_analysis(analysis_id: str) -> Dict[str, bool]:
 
 
 @router.get("/exports", dependencies=[Depends(verify_api_token)])
-def get_exports() -> List[ReportGenerationTaskModel]:
+def get_exports(tag_prefix: Optional[str] = None) -> List[ReportGenerationTaskModel]:
     """List all exports. An export is a request to create human-readable messages that may be sent to scanned entities."""
     return [
         ReportGenerationTaskModel(
@@ -163,7 +163,7 @@ def get_exports() -> List[ReportGenerationTaskModel]:
             error=task.error,
             alerts=task.alerts,
         )
-        for task in db.list_report_generation_tasks()
+        for task in db.list_report_generation_tasks(tag_prefix=tag_prefix)
     ]
 
 

diff --git a/artemis/db.py b/artemis/db.py
@@ -472,9 +472,15 @@ def get_report_generation_task(self, id: int) -> Optional[ReportGenerationTask]:
         with self.session() as session:
             return session.query(ReportGenerationTask).filter(ReportGenerationTask.id == id).first()  # type: ignore
 
-    def list_report_generation_tasks(self) -> List[ReportGenerationTask]:
+    def list_report_generation_tasks(self, tag_prefix: Optional[str] = None) -> List[ReportGenerationTask]:
         with self.session() as session:
-            return list(session.query(ReportGenerationTask).order_by(ReportGenerationTask.created_at.desc()))
+            query = session.query(ReportGenerationTask)
+            if tag_prefix:
+                if "%" in tag_prefix:
+                    raise NotImplementedError()
+
+                query = query.filter(ReportGenerationTask.tag.like(tag_prefix + "%"))
+            return list(query.order_by(ReportGenerationTask.created_at.desc()))
 
     def delete_report_generation_task(self, id: int) -> None:
         with self.session() as session:

diff --git a/test/e2e/test_exporting.py b/test/e2e/test_exporting.py
@@ -148,7 +148,7 @@ def test_exporting_api(self) -> None:
         self.assertEqual(
             requests.post(
                 BACKEND_URL + "api/export",
-                json={"skip_previously_exported": True, "language": "pl_PL"},
+                json={"skip_previously_exported": True, "language": "pl_PL", "tag": "exporting-api"},
                 headers={"X-Api-Token": "api-token"},
             ).json(),
             {"ok": True},
@@ -162,6 +162,27 @@ def test_exporting_api(self) -> None:
 
             time.sleep(1)
 
+        self.assertEqual(len(requests.get(BACKEND_URL + "api/exports", headers={"X-Api-Token": "api-token"}).json()), 1)
+        self.assertEqual(
+            len(requests.get(BACKEND_URL + "api/exports?tag_prefix=", headers={"X-Api-Token": "api-token"}).json()), 1
+        )
+        self.assertEqual(
+            len(
+                requests.get(
+                    BACKEND_URL + "api/exports?tag_prefix=exporting-a", headers={"X-Api-Token": "api-token"}
+                ).json()
+            ),
+            1,
+        )
+        self.assertEqual(
+            len(
+                requests.get(
+                    BACKEND_URL + "api/exports?tag_prefix=exporting-b", headers={"X-Api-Token": "api-token"}
+                ).json()
+            ),
+            0,
+        )
+
         self.assertEqual(
             data[0].keys(),
             {