[DOP-18570] Fix SparkMetricsListener.onExecutionEnd on Python 3.9 and…

… below
MobileTeleSystems · Aug 28, 2024 · 3c4a1e0 · 3c4a1e0
1 parent c6b09be
commit 3c4a1e0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/onetl/_metrics/listener/execution.py b/onetl/_metrics/listener/execution.py
@@ -40,6 +40,10 @@ class SparkSQLMetricNames(str, Enum):  # noqa: WPS338
     def __str__(self):
         return self.value
 
+    @classmethod
+    def values(cls):
+        return set(cls.__members__.values())
+
 
 @dataclass
 class SparkListenerExecution:

diff --git a/onetl/_metrics/listener/listener.py b/onetl/_metrics/listener/listener.py
@@ -12,6 +12,8 @@
     SparkSQLMetricNames,
 )
 
+KNOWN_METRICS = SparkSQLMetricNames.values()
+
 
 @dataclass
 class SparkMetricsListener(BaseSparkListener):
@@ -81,7 +83,7 @@ def onExecutionEnd(self, event):
             for i in range(metrics.size()):
                 metric = metrics.apply(i)
                 metric_name = metric.name()
-                if metric_name not in SparkSQLMetricNames:
+                if metric_name not in KNOWN_METRICS:
                     continue
                 metric_value = metric_values.get(metric.accumulatorId())
                 if not metric_value.isDefined():