feat: Add support for cross-account metrics in CloudWatch Metrics (#2881)

Co-authored-by: Per Kerssens <[email protected]>

Author: 2 people

athena-cloudwatch-metrics/athena-cloudwatch-metrics-package.yaml

@@ -40,6 +40,10 @@ Parameters:
     Description: "(Optional) An IAM policy ARN to use as the PermissionsBoundary for the created Lambda function's execution role"
     Default: ''
     Type: String
+  IncludeLinkedAccounts:
+    Description: "If set to 'true', metrics from linked accounts will be included in the results."
+    Default: 'false'
+    Type: String
 Conditions:
   HasPermissionsBoundary: !Not [ !Equals [ !Ref PermissionsBoundaryARN, "" ] ]
 Resources:
@@ -51,6 +55,7 @@ Resources:
           disable_spill_encryption: !Ref DisableSpillEncryption
           spill_bucket: !Ref SpillBucket
           spill_prefix: !Ref SpillPrefix
+          include_linked_accounts: !Ref IncludeLinkedAccounts
       FunctionName: !Ref AthenaCatalogName
       Handler: "com.amazonaws.athena.connectors.cloudwatch.metrics.MetricsCompositeHandler"
       CodeUri: "./target/athena-cloudwatch-metrics-2022.47.1.jar"

athena-cloudwatch-metrics/athena-cloudwatch-metrics.yaml

@@ -40,6 +40,10 @@ Parameters:
     Description: "(Optional) An IAM policy ARN to use as the PermissionsBoundary for the created Lambda function's execution role"
     Default: ''
     Type: String
+  IncludeLinkedAccounts:
+    Description: "If set to 'true', metrics from linked accounts will be included in the results."
+    Default: 'false'
+    Type: String
 Conditions:
   HasPermissionsBoundary: !Not [ !Equals [ !Ref PermissionsBoundaryARN, "" ] ]
   IsRegionBAH: !Equals [!Ref "AWS::Region", "me-south-1"]
@@ -53,6 +57,7 @@ Resources:
           disable_spill_encryption: !Ref DisableSpillEncryption
           spill_bucket: !Ref SpillBucket
           spill_prefix: !Ref SpillPrefix
+          include_linked_accounts: !Ref IncludeLinkedAccounts
       FunctionName: !Ref AthenaCatalogName
       PackageType: "Image"
       ImageUri: !Sub

athena-cloudwatch-metrics/pom.xml

@@ -62,6 +62,12 @@
             <version>${log4j2Version}</version>
             <scope>runtime</scope>
         </dependency>
+        <dependency>
+            <groupId>org.junit-pioneer</groupId>
+            <artifactId>junit-pioneer</artifactId>
+            <version>2.3.0</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

athena-cloudwatch-metrics/src/main/java/com/amazonaws/athena/connectors/cloudwatch/metrics/MetricUtils.java

@@ -59,8 +59,22 @@ public class MetricUtils
     //this is a format required by Cloudwatch Metrics
     private static final String METRIC_ID = "m1";
 
+    // Environment variable to control cross-account metrics inclusion
+    private static final String INCLUDE_LINKED_ACCOUNTS_ENV_VAR = "include_linked_accounts";
+
     private MetricUtils() {}
 
+    /**
+     * Checks if cross-account metrics should be included based on environment variable
+     *
+     * @return True if cross-account metrics should be included, false otherwise
+     */
+    private static boolean shouldIncludeLinkedAccounts()
+    {
+        String includeLinkedAccounts = System.getenv(INCLUDE_LINKED_ACCOUNTS_ENV_VAR);
+        return Boolean.parseBoolean(includeLinkedAccounts);
+    }
+
     /**
      * Filters metrics who have at least 1 metric dimension that matches DIMENSION_NAME_FIELD and DIMENSION_VALUE_FIELD filters.
      * This is just an optimization and isn't fully correct. We depend on the calling engine to apply full constraints. Also
@@ -104,6 +118,8 @@ protected static void pushDownPredicate(Constraints constraints, ListMetricsRequ
     {
         Map<String, ValueSet> summary = constraints.getSummary();
 
+        listMetricsRequest.includeLinkedAccounts(shouldIncludeLinkedAccounts());
+
         ValueSet namespaceConstraint = summary.get(NAMESPACE_FIELD);
         if (namespaceConstraint != null && namespaceConstraint.isSingleValue()) {
             listMetricsRequest.namespace(namespaceConstraint.getSingleValue().toString());

athena-cloudwatch-metrics/src/test/java/com/amazonaws/athena/connectors/cloudwatch/metrics/MetricUtilsTest.java

@@ -7,9 +7,9 @@
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- * 
+ *
  *      http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -34,9 +34,10 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import com.google.common.collect.ImmutableList;
 import org.apache.arrow.vector.types.Types;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junitpioneer.jupiter.SetEnvironmentVariable;
 import software.amazon.awssdk.services.cloudwatch.model.Dimension;
 import software.amazon.awssdk.services.cloudwatch.model.DimensionFilter;
 import software.amazon.awssdk.services.cloudwatch.model.GetMetricDataRequest;
@@ -60,21 +61,21 @@
 import static com.amazonaws.athena.connectors.cloudwatch.metrics.tables.Table.STATISTIC_FIELD;
 import static com.amazonaws.athena.connectors.cloudwatch.metrics.tables.Table.TIMESTAMP_FIELD;
 import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT;
-import static org.junit.Assert.*;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class MetricUtilsTest
 {
     private FederatedIdentity identity = new FederatedIdentity("arn", "account", Collections.emptyMap(), Collections.emptyList(), Collections.emptyMap());
     private String catalog = "default";
     private BlockAllocator allocator;
 
-    @Before
+    @BeforeEach
     public void setup()
     {
         allocator = new BlockAllocatorImpl();
     }
 
-    @After
+    @AfterEach
     public void tearDown()
     {
         allocator.close();
@@ -137,6 +138,28 @@ public void pushDownPredicate()
         assertEquals(DimensionFilter.builder().name("match4").value("match5").build(), request.dimensions().get(0));
     }
 
+    @Test
+    @SetEnvironmentVariable(key = "include_linked_accounts", value = "true")
+    public void pushDownPredicateWithLinkedAccountsTrue() throws Exception
+    {
+        ListMetricsRequest.Builder requestBuilder = ListMetricsRequest.builder();
+        MetricUtils.pushDownPredicate(new Constraints(new HashMap<>(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap(), null), requestBuilder);
+        ListMetricsRequest request = requestBuilder.build();
+
+        assertTrue(request.includeLinkedAccounts());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "include_linked_accounts", value = "false")
+    public void pushDownPredicateWithLinkedAccountsFalse() throws Exception
+    {
+        ListMetricsRequest.Builder requestBuilder = ListMetricsRequest.builder();
+        MetricUtils.pushDownPredicate(new Constraints(new HashMap<>(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap(), null), requestBuilder);
+        ListMetricsRequest request = requestBuilder.build();
+
+        assertFalse(request.includeLinkedAccounts());
+    }
+
     @Test
     public void makeGetMetricDataRequest()
     {