cleaned up attribute parsing

Author: daniel-sanche

packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/handlers/gcp_exporter.py

@@ -17,6 +17,7 @@
 from collections import defaultdict
 import logging
 import time
+from typing import Any, Mapping
 
 from google.api.distribution_pb2 import Distribution
 from google.api.metric_pb2 import Metric as GMetric
@@ -42,6 +43,7 @@
     NumberDataPoint,
     PeriodicExportingMetricReader,
 )
+from opentelemetry.util.types import Attributes
 
 from google.cloud.bigtable.data._metrics.handlers.opentelemetry import (
     OpenTelemetryMetricsHandler,
@@ -89,6 +91,29 @@
     for n in INSTRUMENT_NAMES
 ]
 
+# Maps OpenTelemetry data point resource attribute names to Cloud Monitoring MonitoredResource label names
+_RESOURCE_KEY_MAP = {
+    "resource_project": "project_id",
+    "resource_instance": "instance",
+    "resource_cluster": "cluster",
+    "resource_table": "table",
+    "resource_zone": "zone",
+}
+
+
+def _partition_attributes(
+    attributes: Attributes
+) -> tuple[dict[str, str], dict[str, str]]:
+    """Split data point attributes into monitored resource labels and metric labels."""
+    resource_labels = {label_name: "" for label_name in _RESOURCE_KEY_MAP.values()}
+    metric_labels = {}
+    for attr_key, attr_value in attributes.items():
+        if attr_key in _RESOURCE_KEY_MAP:
+            resource_labels[_RESOURCE_KEY_MAP[attr_key]] = str(attr_value)
+        elif not attr_key.startswith("resource_"):
+            metric_labels[attr_key] = str(attr_value)
+    return resource_labels, metric_labels
+
 
 class GoogleCloudMetricsHandler(OpenTelemetryMetricsHandler):
     """
@@ -161,9 +186,10 @@ def export(
                 for metric in scope_metric.metrics:
                     for data_point in metric.data.data_points:
                         if data_point.attributes:
-                            project_id = data_point.attributes.get(
-                                "resource_project", ""
+                            resource_labels, metric_labels = _partition_attributes(
+                                data_point.attributes
                             )
+                            project_id = resource_labels.get("project_id", "")
                             if not project_id:
                                 _LOGGER.warning(
                                     "Missing resource_project attribute for metric %s",
@@ -172,21 +198,7 @@ def export(
                                 continue
                             monitored_resource = MonitoredResource(
                                 type="bigtable_client_raw",
-                                labels={
-                                    "project_id": project_id,
-                                    "instance": data_point.attributes.get(
-                                        "resource_instance", ""
-                                    ),
-                                    "cluster": data_point.attributes.get(
-                                        "resource_cluster", ""
-                                    ),
-                                    "table": data_point.attributes.get(
-                                        "resource_table", ""
-                                    ),
-                                    "zone": data_point.attributes.get(
-                                        "resource_zone", ""
-                                    ),
-                                },
+                                labels=resource_labels,
                             )
                             try:
                                 point = self._to_point(data_point)
@@ -204,11 +216,7 @@ def export(
                                 points=[point],
                                 metric=GMetric(
                                     type=f"{self.prefix}/{metric.name}",
-                                    labels={
-                                        k: str(v)
-                                        for k, v in data_point.attributes.items()
-                                        if not k.startswith("resource_")
-                                    },
+                                    labels=metric_labels,
                                 ),
                                 unit=metric.unit,
                             )

packages/google-cloud-bigtable/tests/unit/data/_metrics/test_gcp_exporter_handler.py

@@ -130,6 +130,34 @@ def _make_one(self, *args, **kwargs):
         with mock.patch("google.auth.default", return_value=(mock.Mock(), "project")):
             return BigtableMetricsExporter(*args, **kwargs)
 
+    def test__partition_attributes(self):
+        from google.cloud.bigtable.data._metrics.handlers.gcp_exporter import (
+            _partition_attributes,
+        )
+
+        attributes = {
+            "resource_project": "proj-1",
+            "resource_instance": "inst-1",
+            "resource_cluster": "cls-1",
+            "resource_table": "tbl-1",
+            "resource_zone": "zone-1",
+            "method": "ReadRows",
+            "status_code": "OK",
+            "resource_custom_unsupported": "should_be_stripped",
+        }
+        res_labels, metric_labels = _partition_attributes(attributes)
+        assert res_labels == {
+            "project_id": "proj-1",
+            "instance": "inst-1",
+            "cluster": "cls-1",
+            "table": "tbl-1",
+            "zone": "zone-1",
+        }
+        assert metric_labels == {
+            "method": "ReadRows",
+            "status_code": "OK",
+        }
+
     def test_ctor_defaults(self):
         from google.cloud.monitoring_v3 import MetricServiceClient