From 562cb77003d7c651219ca0bd03618c2bf478083b Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Thu, 20 Jun 2024 17:47:39 +0300
Subject: [PATCH] Support getting and clearing task logs using specific metrics

---
 apiserver/apimodels/events.py                 |  3 ++
 apiserver/bll/event/event_bll.py              | 12 +++++++-
 apiserver/schema/services/events.conf         | 21 ++++++++++++++
 apiserver/services/events.py                  |  3 ++
 apiserver/tests/automated/test_task_events.py | 28 +++++++++++++++++++
 5 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/apiserver/apimodels/events.py b/apiserver/apimodels/events.py
index 6ddc29f..052c1bf 100644
--- a/apiserver/apimodels/events.py
+++ b/apiserver/apimodels/events.py
@@ -146,6 +146,7 @@ class LogEventsRequest(TaskEventsRequestBase):
     navigate_earlier: bool = BoolField(default=True)
     from_timestamp: Optional[int] = IntField()
     order: Optional[str] = ActualEnumField(LogOrderEnum)
+    metrics: Sequence[MetricVariants] = ListField(items_types=MetricVariants)
 
 
 class ScalarMetricsIterRawRequest(TaskEventsRequestBase):
@@ -229,3 +230,5 @@ class ClearTaskLogRequest(Base):
     task: str = StringField(required=True)
     threshold_sec = IntField()
     allow_locked = BoolField(default=False)
+    exclude_metrics = ListField(items_types=[str])
+    include_metrics = ListField(items_types=[str])
diff --git a/apiserver/bll/event/event_bll.py b/apiserver/bll/event/event_bll.py
index 0dd1b88..80d8392 100644
--- a/apiserver/bll/event/event_bll.py
+++ b/apiserver/bll/event/event_bll.py
@@ -1227,6 +1227,8 @@ class EventBLL(object):
         task_id: str,
         allow_locked: bool = False,
         threshold_sec: int = None,
+        include_metrics: Sequence[str] = None,
+        exclude_metrics: Sequence[str] = None,
     ):
         self._validate_task_state(
             company_id=company_id, task_id=task_id, allow_locked=allow_locked
@@ -1251,8 +1253,16 @@ class EventBLL(object):
                     }
                 )
                 sort = {"timestamp": {"order": "desc"}}
+
+            if include_metrics:
+                must.append({"terms": {"metric": include_metrics}})
+
+            more_conditions = {}
+            if exclude_metrics:
+                more_conditions = {"must_not": [{"terms": {"metric": exclude_metrics}}]}
+
             es_req = {
-                "query": {"bool": {"must": must}},
+                "query": {"bool": {"must": must, **more_conditions}},
                 **({"sort": sort} if sort else {}),
             }
             es_res = delete_company_events(
diff --git a/apiserver/schema/services/events.conf b/apiserver/schema/services/events.conf
index 01ba7b9..d660d92 100644
--- a/apiserver/schema/services/events.conf
+++ b/apiserver/schema/services/events.conf
@@ -947,6 +947,13 @@ get_task_log {
             }
         }
     }
+    "999.0": ${get_task_log."2.9"} {
+        request.metrics {
+            type: array
+            description: List of metrics and variants
+            items { "$ref": "#/definitions/metric_variants" }
+        }
+    }
 }
 get_task_events {
     "2.1" {
@@ -1705,4 +1712,18 @@ clear_task_log {
             }
         }
     }
+    "999.0": ${clear_task_log."2.19"} {
+        request.properties {
+            include_metrics {
+                type: array
+                description: If passed then only events for these metrics are deleted
+                items: {type: string}
+            }
+            exclude_metrics {
+                type: array
+                description: If passed then events for these metrics are retained
+                items: {type: string}
+            }
+        }
+    }
 }
diff --git a/apiserver/services/events.py b/apiserver/services/events.py
index 0c4f3ec..724a937 100644
--- a/apiserver/services/events.py
+++ b/apiserver/services/events.py
@@ -172,6 +172,7 @@ def get_task_log(call, company_id, request: LogEventsRequest):
         batch_size=request.batch_size,
         navigate_earlier=request.navigate_earlier,
         from_timestamp=request.from_timestamp,
+        metric_variants=_get_metric_variants_from_request(request.metrics),
     )
 
     if request.order and (
@@ -1041,6 +1042,8 @@ def clear_task_log(call: APICall, company_id: str, request: ClearTaskLogRequest)
             task_id=task_id,
             allow_locked=request.allow_locked,
             threshold_sec=request.threshold_sec,
+            exclude_metrics=request.exclude_metrics,
+            include_metrics=request.include_metrics,
         )
     )
 
diff --git a/apiserver/tests/automated/test_task_events.py b/apiserver/tests/automated/test_task_events.py
index d9998ea..0be5d50 100644
--- a/apiserver/tests/automated/test_task_events.py
+++ b/apiserver/tests/automated/test_task_events.py
@@ -346,6 +346,34 @@ class TestTaskEvents(TestService):
         # test order
         self._assert_log_events(task=task, order="asc")
 
+        metric = "metric"
+        variant = "variant"
+        events = [
+            self._create_task_event(
+                "log",
+                task=task,
+                iteration=iter_,
+                timestamp=timestamp + iter_ * 1000,
+                msg=f"This is a log message from test task iter {iter_}",
+                metric=metric,
+                variant=variant,
+            )
+            for iter_ in range(2)
+        ]
+        self.send_batch(events)
+        res = self.api.events.get_task_log(task=task)
+        self.assertEqual(res.total, 12)
+        res = self.api.events.get_task_log(task=task, metrics=[{"metric": metric}])
+        self.assertEqual(res.total, 2)
+
+        # test clear
+        self.api.events.clear_task_log(task=task, exclude_metrics=[metric])
+        res = self.api.events.get_task_log(task=task)
+        self.assertEqual(res.total, 2)
+        self.api.events.clear_task_log(task=task)
+        res = self.api.events.get_task_log(task=task)
+        self.assertEqual(res.total, 0)
+
     def _assert_log_events(
         self,
         task,