Overload void addDimension(DimensionSet dimensionSet) to be able to add multi-dimensional dimensions also for non default dimensions.

Author: phipag

powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsLogger.java

@@ -58,12 +58,22 @@ default void addMetric(String key, double value) {
     }
 
     /**
-     * Add a dimension to the metrics logger
+     * Add a dimension to the metrics logger.
+     * This is equivalent to calling {@code addDimension(DimensionSet.of(key, value))}
      *
      * @param key   the name of the dimension
      * @param value the value of the dimension
      */
-    void addDimension(String key, String value);
+    default void addDimension(String key, String value) {
+        addDimension(DimensionSet.of(key, value));
+    }
+    
+    /**
+     * Add a dimension set to the metrics logger
+     *
+     * @param dimensionSet the dimension set to add
+     */
+    void addDimension(DimensionSet dimensionSet);
 
     /**
      * Add metadata to the metrics logger
@@ -79,7 +89,7 @@ default void addMetric(String key, double value) {
      * @param defaultDimensions map of default dimensions
      */
     void setDefaultDimensions(Map<String, String> defaultDimensions);
-    
+
     /**
      * Get the default dimensions for the metrics logger
      *
@@ -118,7 +128,7 @@ default void addMetric(String key, double value) {
      * @param dimensions custom dimensions for this metric (optional)
      */
     void captureColdStartMetric(Context context, DimensionSet dimensions);
-    
+
     /**
      * Capture cold start metric and flush immediately
      *
@@ -127,19 +137,19 @@ default void addMetric(String key, double value) {
     default void captureColdStartMetric(Context context) {
         captureColdStartMetric(context, null);
     }
-    
+
     /**
      * Capture cold start metric without Lambda context and flush immediately
      *
      * @param dimensions custom dimensions for this metric (optional)
      */
     void captureColdStartMetric(DimensionSet dimensions);
-    
+
     /**
      * Capture cold start metric without Lambda context and flush immediately
      */
     default void captureColdStartMetric() {
-        captureColdStartMetric((DimensionSet)null);
+        captureColdStartMetric((DimensionSet) null);
     }
 
     /**
@@ -152,9 +162,9 @@ default void captureColdStartMetric() {
      * @param namespace  the namespace for the metric
      * @param dimensions custom dimensions for this metric (optional)
      */
-    void pushSingleMetric(String name, double value, MetricUnit unit, String namespace, 
-                          DimensionSet dimensions);
-    
+    void pushSingleMetric(String name, double value, MetricUnit unit, String namespace,
+            DimensionSet dimensions);
+
     /**
      * Push a single metric with custom dimensions. This creates a separate metrics context
      * that doesn't affect the default metrics context.
@@ -167,4 +177,4 @@ void pushSingleMetric(String name, double value, MetricUnit unit, String namespa
     default void pushSingleMetric(String name, double value, MetricUnit unit, String namespace) {
         pushSingleMetric(name, value, unit, namespace, null);
     }
-}
+}

powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/internal/EmfMetricsLogger.java

@@ -64,16 +64,23 @@ public void addMetric(String key, double value, MetricUnit unit, MetricResolutio
     }
 
     @Override
-    public void addDimension(String key, String value) {
-        DimensionSet dimensionSet = new DimensionSet();
-        try {
-            dimensionSet.addDimension(key, value);
-            emfLogger.putDimensions(dimensionSet);
-            // Update our local copy of default dimensions
-            defaultDimensions.put(key, value);
-        } catch (Exception e) {
-            // Ignore dimension errors
+    public void addDimension(software.amazon.lambda.powertools.metrics.model.DimensionSet dimensionSet) {
+        if (dimensionSet == null) {
+            throw new IllegalArgumentException("DimensionSet cannot be null");
         }
+
+        DimensionSet emfDimensionSet = new DimensionSet();
+        dimensionSet.getDimensions().forEach((key, val) -> {
+            try {
+                emfDimensionSet.addDimension(key, val);
+                // Update our local copy of default dimensions
+                defaultDimensions.put(key, val);
+            } catch (Exception e) {
+                // Ignore dimension errors
+            }
+        });
+
+        emfLogger.putDimensions(emfDimensionSet);
     }
 
     @Override

powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/internal/EmfMetricsLoggerTest.java

@@ -182,6 +182,50 @@ void shouldAddDimension() throws Exception {
         }
         assertThat(hasDimension).isTrue();
     }
+    
+    @Test
+    void shouldAddDimensionSet() throws Exception {
+        // Given
+        DimensionSet dimensionSet = DimensionSet.of("Dim1", "Value1", "Dim2", "Value2");
+        
+        // When
+        metricsLogger.clearDefaultDimensions(); // Clear default Service dimension first for easier assertions
+        metricsLogger.addDimension(dimensionSet);
+        metricsLogger.addMetric("test-metric", 100);
+        metricsLogger.flush();
+
+        // Then
+        String emfOutput = outputStreamCaptor.toString().trim();
+        JsonNode rootNode = objectMapper.readTree(emfOutput);
+
+        assertThat(rootNode.has("Dim1")).isTrue();
+        assertThat(rootNode.get("Dim1").asText()).isEqualTo("Value1");
+        assertThat(rootNode.has("Dim2")).isTrue();
+        assertThat(rootNode.get("Dim2").asText()).isEqualTo("Value2");
+
+        // Check that the dimensions are in the CloudWatchMetrics section
+        JsonNode dimensions = rootNode.get("_aws").get("CloudWatchMetrics").get(0).get("Dimensions").get(0);
+        boolean hasDim1 = false;
+        boolean hasDim2 = false;
+        for (JsonNode dimension : dimensions) {
+            String dimName = dimension.asText();
+            if (dimName.equals("Dim1")) {
+                hasDim1 = true;
+            } else if (dimName.equals("Dim2")) {
+                hasDim2 = true;
+            }
+        }
+        assertThat(hasDim1).isTrue();
+        assertThat(hasDim2).isTrue();
+    }
+    
+    @Test
+    void shouldThrowExceptionWhenDimensionSetIsNull() {
+        // When/Then
+        assertThatThrownBy(() -> metricsLogger.addDimension((DimensionSet) null))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessage("DimensionSet cannot be null");
+    }
 
     @Test
     void shouldAddMetadata() throws Exception {

powertools-metrics/src/test/java/software/amazon/lambda/powertools/metrics/testutils/TestMetricsLogger.java

@@ -27,7 +27,7 @@ public void addMetric(String key, double value, MetricUnit unit, MetricResolutio
     }
 
     @Override
-    public void addDimension(String key, String value) {
+    public void addDimension(DimensionSet dimensionSet) {
         // Test placeholder
     }