Updates based on discussion in PR #226

* Rename Config.METRICS > ALLOWED_METRICS and change from list to set
* Add Metric dataclass and update MetricsClient class to use it
* Update calls in Workflow class to use Metric dataclass
* Update dependencies

Author: ehanson8

dsc/config.py

@@ -7,12 +7,12 @@
 
 METRICS_NAMESPACE = "dso"
 
-METRICS = [
+ALLOWED_METRICS = {
     "item_submitted",  # item submitted to DSS
     "submission_error",  # error during submission to DSS
     "ingested_item",  # item ingested successfully into DSpace
     "ingest_error",  # error during attempted item ingest into DSpace
-]
+}
 
 
 class Config:

dsc/utils/aws/__init__.py

@@ -1,6 +1,6 @@
-from dsc.utils.aws.metrics import MetricsClient
+from dsc.utils.aws.metrics import Metric, MetricsClient
 from dsc.utils.aws.s3 import S3Client
 from dsc.utils.aws.ses import SESClient
 from dsc.utils.aws.sqs import SQSClient
 
-__all__ = ["MetricsClient", "S3Client", "SESClient", "SQSClient"]
+__all__ = ["Metric", "MetricsClient", "S3Client", "SESClient", "SQSClient"]

dsc/utils/aws/metrics.py

@@ -3,10 +3,11 @@
 from __future__ import annotations
 
 import logging
+from dataclasses import dataclass
 
 import boto3
 
-from dsc.config import METRICS, METRICS_NAMESPACE
+from dsc.config import METRICS_NAMESPACE
 
 logger = logging.getLogger(__name__)
 
@@ -43,84 +44,66 @@
 )
 
 
+@dataclass
+class Metric:
+    """A class representing a single metric to be published to CloudWatch."""
+
+    name: str
+    value: int
+    unit: str
+    dimensions: dict[str, str] | None = None
+
+
 class MetricsClient:
     """A simple client to record metrics to AWS CloudWatch."""
 
-    def __init__(self) -> None:
+    def __init__(self, allowed_metrics: set[str] | None = None) -> None:
         """Initialize the MetricsClient."""
-        self.cloudwatch = boto3.client("cloudwatch")
-        self.batch_metrics: list[dict] = []
+        self.namespace = METRICS_NAMESPACE
+        self.allowed_metrics: set[str] | None = allowed_metrics
+        self._cloudwatch = boto3.client("cloudwatch")
+        self.batch_metrics: list[Metric] = []
 
     def publish_single_metric(
         self,
-        metric_name: str,
-        value: int,
-        unit: str,
-        metric_dimensions: dict[str, str] | None = None,
+        metric: Metric,
     ) -> None:
-        """Publish a single metric to CloudWatch.
-
-        Args:
-            metric_name: The name of the metric to publish.
-            value: The value of the metric.
-            unit: The unit of the metric.
-            metric_dimensions: Optional dictionary of dimension names and values.
-
-        Raises:
-            ValueError: If unit is invalid.
-        """
-        metric_data = self._validate_and_build_metric_data(
-            metric_name, value, unit, metric_dimensions
-        )
-        self._push_metric_data([metric_data])
+        """Publish a single metric to CloudWatch."""
+        self._validate_metric(metric)
+        self._push_metric_data([metric])
 
-    def _validate_and_build_metric_data(
+    def _validate_metric(
         self,
-        metric_name: str,
-        value: int,
-        unit: str,
-        metric_dimensions: dict[str, str] | None = None,
-    ) -> dict:
-        """Validate and build a metric data dictionary for CloudWatch.
+        metric: Metric,
+    ) -> bool:
+        """Validate that a metric has required fields and allowed unit.
 
         Args:
-            metric_name: The name of the metric.
-            value: The value of the metric.
-            unit: The unit of the metric.
-            metric_dimensions: Optional dictionary of dimension names and values.
-
-        Returns:
-            A metric data dictionary formatted for CloudWatch.
+            metric: The Metric instance to validate.
         """
-        self._approved_metric(metric_name)
-        self._validate_unit(unit)
-        dimensions = [
-            {"Name": name, "Value": dim_value}
-            for name, dim_value in (
-                metric_dimensions.items() if metric_dimensions else []
+        if not all(hasattr(metric, attr) for attr in ["name", "value", "unit"]):
+            raise ValueError(
+                f"Metric must have 'name', 'value', and 'unit' attributes. Invalid "
+                f"metric: {metric}"
             )
-        ]
-        return {
-            "MetricName": metric_name,
-            "Value": value,
-            "Unit": unit,
-            "Dimensions": dimensions,
-        }
+        self._allowed_metric(metric.name)
+        self._validate_metric_unit(metric.unit)
+        return True
 
-    def _approved_metric(self, metric_name: str) -> bool:
-        """Check if a metric name is in the approved list of metrics for the application.
+    def _allowed_metric(self, metric_name: str) -> bool:
+        """Check if a metric name is in the allowed list of metrics for the application.
 
         Args:
             metric_name: The name of the metric to check.
         """
-        if metric_name not in METRICS:
+        if self.allowed_metrics and metric_name not in self.allowed_metrics:
             raise ValueError(
-                f"Metric name '{metric_name}' is not in the approved list of metrics: "
-                f"{', '.join(METRICS)}"
+                f"Metric name '{metric_name}' is not in the allowed list of metrics: "
+                f"{', '.join(self.allowed_metrics)}"
             )
         return True
 
-    def _validate_unit(self, unit: str) -> None:
+    def _validate_metric_unit(self, unit: str) -> bool:
         """Validate that metric unit is allowed by AWS CloudWatch.
 
         Args:
@@ -133,66 +116,75 @@ def _validate_unit(self, unit: str) -> None:
             raise ValueError(
                 f"Invalid unit '{unit}'. Must be one of: {', '.join(UNIT_VALUES)}"
             )
+        return True
 
-    def _push_metric_data(self, metric_data: list[dict]) -> None:
-        """Push metric data to CloudWatch.
+    def _push_metric_data(self, metrics: list[Metric]) -> None:
+        """Push metrics to CloudWatch.
 
         Args:
-            metric_data: List of metric dictionaries to push.
+            metrics: List of metric instances to push.
         """
+        if not metrics:
+            logger.info("No metrics to publish.")
+            return
+
         try:
-            self.cloudwatch.put_metric_data(
-                Namespace=METRICS_NAMESPACE, MetricData=metric_data
+            metric_data = []
+            for metric in metrics:
+                metric_dict = {
+                    "MetricName": metric.name,
+                    "Value": metric.value,
+                    "Unit": metric.unit,
+                }
+                if metric.dimensions:
+                    metric_dict["Dimensions"] = [
+                        {"Name": key, "Value": value}
+                        for key, value in metric.dimensions.items()
+                    ]
+                metric_data.append(metric_dict)
+
+            self._cloudwatch.put_metric_data(
+                Namespace=self.namespace,
+                MetricData=metric_data,
+            )
+            logger.info(
+                f"Published {len(metrics)} metric(s) to CloudWatch namespace "
+                f"'{self.namespace}'."
             )
-            logger.info(f"Published metric with {metric_data} to CloudWatch.")
         except Exception:
             logger.exception(
-                f"Failed to publish metric with {metric_data} to CloudWatch."
+                f"Failed to publish {len(metrics)} metric(s) to CloudWatch: "
             )
+            raise
 
-    def add_metric_to_batch(
-        self,
-        metric_name: str,
-        value: int,
-        unit: str,
-        metric_dimensions: dict[str, str] | None = None,
-    ) -> None:
-        """Add a metric to the batch for later publishing.
+    def add_metric_to_batch(self, metric: Metric) -> None:
+        """Add a metric to the batch queue.
 
         Args:
-            metric_name: The name of the metric.
-            value: The value of the metric.
-            unit: The unit of the metric.
-            metric_dimensions: Optional dictionary of dimension names and values.
-
-        Raises:
-            ValueError: If unit is invalid.
+            metric: The metric to add to the batch.
         """
-        metric_data = self._validate_and_build_metric_data(
-            metric_name, value, unit, metric_dimensions
-        )
-        self.batch_metrics.append(metric_data)
+        self._validate_metric(metric)
+        self.batch_metrics.append(metric)
 
     def publish_batch_metrics(self, batch_size: int = 20) -> None:
-        """Publish all accumulated batch metrics to CloudWatch.
+        """Publish a batch of metrics to CloudWatch.
 
-        Raises:
-            ValueError: If any metric has an invalid unit or missing required fields.
+        Clears the batch queue after publishing.
+
+        Args:
+            batch_size: Number of metrics to publish in each batch.
         """
         if not self.batch_metrics:
             logger.info("No metrics to publish.")
             return
 
-        # Validate all metrics before publishing
-        for metric in self.batch_metrics:
-            if not all(key in metric for key in ["MetricName", "Value", "Unit"]):
-                raise ValueError(
-                    f"Each metric must contain 'MetricName', 'Value', and 'Unit'. "
-                    f"Invalid metric: {metric}"
-                )
-            self._approved_metric(metric["MetricName"])
-            self._validate_unit(metric["Unit"])
-
-        for x in range(0, len(self.batch_metrics), batch_size):
-            self._push_metric_data(self.batch_metrics[x : x + batch_size])
-        self.batch_metrics.clear()
+        try:
+            # Re-validate all metrics before publishing to CloudWatch
+            for metric in self.batch_metrics:
+                self._validate_metric(metric)
+
+            for x in range(0, len(self.batch_metrics), batch_size):
+                batch = self.batch_metrics[x : x + batch_size]
+                self._push_metric_data(batch)
+        finally:
+            self.batch_metrics.clear()

dsc/workflows/base/workflow.py

@@ -10,7 +10,7 @@
 import jsonschema
 import jsonschema.exceptions
 
-from dsc.config import Config
+from dsc.config import ALLOWED_METRICS, Config
 from dsc.db.models import ItemSubmissionStatus
 from dsc.exceptions import (
     BatchCreationFailedError,
@@ -19,7 +19,7 @@
 )
 from dsc.item_submission import ItemSubmission
 from dsc.reports import CreateReport, FinalizeReport, SubmitReport
-from dsc.utils.aws import MetricsClient, SESClient, SQSClient
+from dsc.utils.aws import Metric, MetricsClient, SESClient, SQSClient
 from dsc.utils.validate.schemas import RESULT_MESSAGE_ATTRIBUTES, RESULT_MESSAGE_BODY
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -130,7 +130,7 @@ def __init__(self, batch_id: str) -> None:
             "skipped": 0,
             "errors": 0,
         }
-        self.metrics_client = MetricsClient()
+        self.metrics_client = MetricsClient(allowed_metrics=ALLOWED_METRICS)
         self.metrics_dimensions = {
             "application": "dsc",
             "workflow_name": self.workflow_name,
@@ -334,10 +334,12 @@ def submit_items(self, collection_handle: str | None = None) -> list:
                 item_submission.submit_attempts += 1
                 item_submission.upsert_db()
                 self.metrics_client.publish_single_metric(
-                    metric_name="item_submitted",
-                    value=1,
-                    unit="Count",
-                    metric_dimensions=self.metrics_dimensions,
+                    Metric(
+                        name="item_submitted",
+                        value=1,
+                        unit="Count",
+                        dimensions=self.metrics_dimensions,
+                    )
                 )
             except NotImplementedError:
                 raise
@@ -349,10 +351,12 @@ def submit_items(self, collection_handle: str | None = None) -> list:
                 item_submission.upsert_db()
 
                 self.metrics_client.publish_single_metric(
-                    metric_name="submission_error",
-                    value=1,
-                    unit="Count",
-                    metric_dimensions=self.metrics_dimensions,
+                    Metric(
+                        name="submission_error",
+                        value=1,
+                        unit="Count",
+                        dimensions=self.metrics_dimensions,
+                    )
                 )
 
         logger.info(
@@ -441,10 +445,12 @@ def finalize_items(self) -> None:
                 logger.debug(f"Record {log_str} was ingested")
 
                 self.metrics_client.publish_single_metric(
-                    metric_name="ingested_item",
-                    value=1,
-                    unit="Count",
-                    metric_dimensions=self.metrics_dimensions,
+                    Metric(
+                        name="ingested_item",
+                        value=1,
+                        unit="Count",
+                        dimensions=self.metrics_dimensions,
+                    )
                 )
             elif result_message.result_type == "error":
                 item_submission.status = ItemSubmissionStatus.INGEST_FAILED
@@ -453,10 +459,12 @@ def finalize_items(self) -> None:
                 logger.debug(f"Record {log_str} failed to ingest")
 
                 self.metrics_client.publish_single_metric(
-                    metric_name="ingest_error",
-                    value=1,
-                    unit="Count",
-                    metric_dimensions=self.metrics_dimensions,
+                    Metric(
+                        name="ingest_error",
+                        value=1,
+                        unit="Count",
+                        dimensions=self.metrics_dimensions,
+                    )
                 )
 
             else: