diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 985559f..7ef43f7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,29 +1,28 @@ default_language_version: - python: python3.12 # set for project python version + python: python3.12 +default_stages: + - pre-push repos: - repo: local hooks: - - id: black-apply - name: black-apply - entry: uv run black + - id: ruff-format + name: ruff-format + entry: uv run ruff format --diff language: system pass_filenames: true - types: ["python"] + types: [ "python" ] + - id: mypy name: mypy entry: uv run mypy language: system pass_filenames: true - types: ["python"] - exclude: "tests/" - - id: ruff-apply - name: ruff-apply - entry: uv run ruff check --fix + types: [ "python" ] + exclude: "(tests/|output/|migrations/)" + + - id: ruff-check + name: ruff-check + entry: uv run ruff check language: system pass_filenames: true - types: ["python"] - - id: pip-audit - name: pip-audit - entry: uv run pip-audit - language: system - pass_filenames: false \ No newline at end of file + types: [ "python" ] \ No newline at end of file diff --git a/dsc/config.py b/dsc/config.py index 7cbff6f..77b20a4 100644 --- a/dsc/config.py +++ b/dsc/config.py @@ -5,6 +5,15 @@ import sentry_sdk +METRICS_NAMESPACE = "dso" + +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: REQUIRED_ENV_VARS: Iterable[str] = [ diff --git a/dsc/utils/aws/__init__.py b/dsc/utils/aws/__init__.py index 700b14b..5088c4c 100644 --- a/dsc/utils/aws/__init__.py +++ b/dsc/utils/aws/__init__.py @@ -1,5 +1,6 @@ +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__ = ["S3Client", "SESClient", "SQSClient"] +__all__ = ["Metric", "MetricsClient", "S3Client", "SESClient", "SQSClient"] diff --git a/dsc/utils/aws/metrics.py b/dsc/utils/aws/metrics.py new file mode 100644 index 0000000..613e08b --- /dev/null +++ b/dsc/utils/aws/metrics.py @@ -0,0 +1,236 @@ +"""AWS CloudWatch metrics client for workflow submission tracking.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +import boto3 + +logger = logging.getLogger(__name__) + +CLOUDWATCH_METRICS_LIMIT = 1000 + +UNIT_VALUES = frozenset( + [ + "Bits", + "Bits/Second", + "Bytes", + "Bytes/Second", + "Count", + "Count/Second", + "Gigabits", + "Gigabits/Second", + "Gigabytes", + "Gigabytes/Second", + "Kilobits", + "Kilobits/Second", + "Kilobytes", + "Kilobytes/Second", + "Megabits", + "Megabits/Second", + "Megabytes", + "Megabytes/Second", + "Milliseconds", + "Microseconds", + "None", + "Percent", + "Seconds", + "Terabits", + "Terabits/Second", + "Terabytes", + "Terabytes/Second", + ] +) + + +@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 + namespace: str | None = None + + +class MetricsClient: + """A simple client to record metrics to AWS CloudWatch.""" + + def __init__( + self, namespace: str | None = None, allowed_metrics: set[str] | None = None + ) -> None: + """Initialize the MetricsClient.""" + self.namespace = namespace + self.allowed_metrics: set[str] | None = allowed_metrics + self._cloudwatch = boto3.client("cloudwatch") + self.batch_metrics: list[Metric] = [] + + def publish_metric( + self, + metric: Metric, + ) -> None: + """Publish a single metric to CloudWatch.""" + self._validate_metric(metric) + self._publish_metrics([metric]) + + def _validate_metric( + self, + metric: Metric, + ) -> bool: + """Validate that a metric has required fields and allowed unit. + + Args: + metric: The Metric instance to validate. + """ + 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}" + ) + self._allowed_metric(metric.name) + self._validate_metric_unit(metric.unit) + return True + + 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 self.allowed_metrics and metric_name not in self.allowed_metrics: + raise ValueError( + f"Metric name '{metric_name}' is not in the allowed list of metrics: " + f"{', '.join(self.allowed_metrics)}" + ) + return True + + def _validate_metric_unit(self, unit: str) -> bool: + """Validate that metric unit is allowed by AWS CloudWatch. + + Args: + unit: The unit to validate. + + Raises: + ValueError: If unit is not allowed by AWS CloudWatch. + """ + if unit not in UNIT_VALUES: + allowed_units = ", ".join(sorted(UNIT_VALUES)) + raise ValueError(f"Invalid unit '{unit}'. Must be one of: {allowed_units}") + return True + + def _publish_metrics(self, metrics: list[Metric]) -> None: + """Publish metrics to CloudWatch. + + Automatically chunks metrics if the list exceeds CloudWatch's limit + of 1000 metrics per request. + + Args: + metrics: List of metric instances to publish. + """ + if not metrics: + logger.info("No metrics to publish.") + return + + # Defensively chunk metrics if they exceed CloudWatch's limit + if len(metrics) > CLOUDWATCH_METRICS_LIMIT: + logger.info( + f"Splitting {len(metrics)} metrics into chunks of " + f"{CLOUDWATCH_METRICS_LIMIT} for CloudWatch compliance." + ) + for i in range(0, len(metrics), CLOUDWATCH_METRICS_LIMIT): + chunk = metrics[i : i + CLOUDWATCH_METRICS_LIMIT] + self._publish_metrics(chunk) + return + + try: + # Validate all metrics and ensure consistent namespace + namespaces = set() + metric_data = [] + for metric in metrics: + self._validate_namespace(metric) + selected_namespace = metric.namespace or self.namespace + namespaces.add(selected_namespace) + + 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) + + # Ensure all metrics resolve to the same namespace + if len(namespaces) > 1: + raise ValueError( # noqa: TRY301 + f"Cannot publish metrics with different namespaces in a single " + f"request. Found: {namespaces}" + ) + + selected_namespace = namespaces.pop() + self._cloudwatch.put_metric_data( + Namespace=selected_namespace, + MetricData=metric_data, + ) + logger.info( + f"Published {len(metrics)} metric(s) to CloudWatch namespace " + f"'{selected_namespace}'." + ) + except Exception: + logger.exception( + f"Failed to publish {len(metrics)} metric(s) to CloudWatch: " + ) + raise + + def _validate_namespace(self, metric: Metric) -> bool: + """Validate metric has a namespace or the client has a default namespace.""" + if not metric.namespace and not self.namespace: + raise ValueError( + f"Metric '{metric.name}' must have a namespace if no default " + f"namespace is set for the MetricsClient." + ) + return True + + def add_metrics_to_batch(self, metrics: list[Metric]) -> None: + """Add metrics to the batch queue. + + Args: + metrics: The metrics to add to the batch. + """ + for metric in metrics: + self._validate_metric(metric) + self.batch_metrics.append(metric) + + def publish_metrics_batch(self, batch_size: int = 20) -> None: + """Publish a batch of metrics to CloudWatch. + + Clears the batch queue after successful publishing. + + Args: + batch_size: Number of metrics to publish in each batch. Must be less than + CloudWatch's limit of 1000 metrics per request. + + Raises: + Exception: If publishing fails, metrics remain in the batch queue + for retry or manual handling. + """ + if not self.batch_metrics: + logger.info("No metrics to publish.") + return + + try: + for x in range(0, len(self.batch_metrics), batch_size): + batch = self.batch_metrics[x : x + batch_size] + self._publish_metrics(batch) + except Exception: + # Keep only the unpublished metrics (starting from the failed batch) + self.batch_metrics = self.batch_metrics[x:] + raise + + # Clear only if all batches published successfully + self.batch_metrics.clear() diff --git a/dsc/workflows/base/workflow.py b/dsc/workflows/base/workflow.py index 0844371..9f8bc79 100644 --- a/dsc/workflows/base/workflow.py +++ b/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, METRICS_NAMESPACE, 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 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,6 +130,13 @@ def __init__(self, batch_id: str) -> None: "skipped": 0, "errors": 0, } + self.metrics_client = MetricsClient( + namespace=METRICS_NAMESPACE, allowed_metrics=ALLOWED_METRICS + ) + self.metrics_dimensions = { + "application": "dsc", + "workflow_name": self.workflow_name, + } # cache list of bitstreams self._batch_bitstream_uris: list[str] | None = None @@ -328,6 +335,7 @@ def submit_items(self, collection_handle: str | None = None) -> list: item_submission.status_details = None item_submission.submit_attempts += 1 item_submission.upsert_db() + self._publish_count_metric("item_submitted", f"item {item_identifier}") except NotImplementedError: raise except Exception as exception: # noqa: BLE001 @@ -336,6 +344,9 @@ def submit_items(self, collection_handle: str | None = None) -> list: item_submission.status_details = str(exception) item_submission.submit_attempts += 1 item_submission.upsert_db() + self._publish_count_metric( + "submission_error", f"item {item_submission.item_identifier}" + ) logger.info( f"Submitted messages to the DSS input queue '{CONFIG.sqs_queue_dss_input}' " @@ -421,11 +432,14 @@ def finalize_items(self) -> None: ) sqs_results_summary["ingest_success"] += 1 logger.debug(f"Record {log_str} was ingested") + self._publish_count_metric("ingested_item", f"record {log_str}") elif result_message.result_type == "error": item_submission.status = ItemSubmissionStatus.INGEST_FAILED item_submission.status_details = result_message.error_info sqs_results_summary["ingest_failed"] += 1 logger.debug(f"Record {log_str} failed to ingest") + self._publish_count_metric("ingest_error", f"record {log_str}") + else: item_submission.status = ItemSubmissionStatus.INGEST_UNKNOWN sqs_results_summary["ingest_unknown"] += 1 @@ -451,6 +465,27 @@ def workflow_specific_processing(self) -> None: f"No extra processing for batch based on workflow: '{self.workflow_name}' " ) + def _publish_count_metric(self, metric_name: str, log_data: str) -> None: + """Publish a count metric to CloudWatch. + + Any exceptions are caught and logged. + + Args: + metric_name: The name of the metric to publish. + log_data: Additional data included in the log message. + """ + try: + self.metrics_client.publish_metric( + Metric( + name=metric_name, + value=1, + unit="Count", + dimensions=self.metrics_dimensions, + ) + ) + except Exception: + logger.exception(f"Failed to publish '{metric_name}' metric: {log_data}") + def send_report( self, step: Literal["create", "submit", "finalize"], diff --git a/pyproject.toml b/pyproject.toml index 75e9bb3..0ec2212 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "lxml>=6.0.2", "pandas>=2.3.3", "pynamodb>=6.1.0", - "requests>=2.33.0", + "requests>=2.34.0", "sentry-sdk>=2.50.0", "smart_open", ] diff --git a/uv.lock b/uv.lock index 2d2b41d..0917323 100644 --- a/uv.lock +++ b/uv.lock @@ -588,7 +588,7 @@ requires-dist = [ { name = "lxml", specifier = ">=6.0.2" }, { name = "pandas", specifier = ">=2.3.3" }, { name = "pynamodb", specifier = ">=6.1.0" }, - { name = "requests", specifier = ">=2.33.0" }, + { name = "requests", specifier = ">=2.34.0" }, { name = "sentry-sdk", specifier = ">=2.50.0" }, { name = "smart-open" }, ]