Query Aurora instances per cluster (DataDog#31846)

kkpan11 · Dec 6, 2024 · d255f0c · d255f0c
1 parent 858cb1c
commit d255f0c
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 38 deletions.
diff --git a/pkg/databasemonitoring/aws/aurora.go b/pkg/databasemonitoring/aws/aurora.go
@@ -45,45 +45,50 @@ func (c *Client) GetAuroraClusterEndpoints(ctx context.Context, dbClusterIdentif
 	if len(dbClusterIdentifiers) == 0 {
 		return nil, fmt.Errorf("at least one database cluster identifier is required")
 	}
-	clusterInstances, err := c.client.DescribeDBInstances(ctx,
-		&rds.DescribeDBInstancesInput{
-			Filters: []types.Filter{
-				{
-					Name:   aws.String("db-cluster-id"),
-					Values: dbClusterIdentifiers,
-				},
-			},
-		})
-	if err != nil {
-		return nil, fmt.Errorf("error running GetAuroraClusterEndpoints %v", err)
-	}
 	clusters := make(map[string]*AuroraCluster, 0)
-	for _, db := range clusterInstances.DBInstances {
-		if db.Endpoint != nil && db.DBClusterIdentifier != nil {
-			if db.Endpoint.Address == nil || db.DBInstanceStatus == nil || strings.ToLower(*db.DBInstanceStatus) != "available" {
-				continue
-			}
-			// Add to list of instances for the cluster
-			instance := &Instance{
-				Endpoint: *db.Endpoint.Address,
-			}
-			// Set if IAM is configured for the endpoint
-			if db.IAMDatabaseAuthenticationEnabled != nil {
-				instance.IamEnabled = *db.IAMDatabaseAuthenticationEnabled
-			}
-			// Set the port, if it is known
-			if db.Endpoint.Port != nil {
-				instance.Port = *db.Endpoint.Port
-			}
-			if db.Engine != nil {
-				instance.Engine = *db.Engine
-			}
-			if _, ok := clusters[*db.DBClusterIdentifier]; !ok {
-				clusters[*db.DBClusterIdentifier] = &AuroraCluster{
-					Instances: make([]*Instance, 0),
+
+	for _, clusterID := range dbClusterIdentifiers {
+		// TODO: Seth Samuel: This method is not paginated, so if there are more than 100 instances in a cluster, we will only get the first 100
+		// We should add pagination support to this method at some point
+		clusterInstances, err := c.client.DescribeDBInstances(ctx,
+			&rds.DescribeDBInstancesInput{
+				Filters: []types.Filter{
+					{
+						Name:   aws.String("db-cluster-id"),
+						Values: []string{clusterID},
+					},
+				},
+			})
+		if err != nil {
+			return nil, fmt.Errorf("error running GetAuroraClusterEndpoints %v", err)
+		}
+		for _, db := range clusterInstances.DBInstances {
+			if db.Endpoint != nil && db.DBClusterIdentifier != nil {
+				if db.Endpoint.Address == nil || db.DBInstanceStatus == nil || strings.ToLower(*db.DBInstanceStatus) != "available" {
+					continue
+				}
+				// Add to list of instances for the cluster
+				instance := &Instance{
+					Endpoint: *db.Endpoint.Address,
+				}
+				// Set if IAM is configured for the endpoint
+				if db.IAMDatabaseAuthenticationEnabled != nil {
+					instance.IamEnabled = *db.IAMDatabaseAuthenticationEnabled
+				}
+				// Set the port, if it is known
+				if db.Endpoint.Port != nil {
+					instance.Port = *db.Endpoint.Port
+				}
+				if db.Engine != nil {
+					instance.Engine = *db.Engine
+				}
+				if _, ok := clusters[*db.DBClusterIdentifier]; !ok {
+					clusters[*db.DBClusterIdentifier] = &AuroraCluster{
+						Instances: make([]*Instance, 0),
+					}
 				}
+				clusters[*db.DBClusterIdentifier].Instances = append(clusters[*db.DBClusterIdentifier].Instances, instance)
 			}
-			clusters[*db.DBClusterIdentifier].Instances = append(clusters[*db.DBClusterIdentifier].Instances, instance)
 		}
 	}
 	if len(clusters) == 0 {

diff --git a/pkg/databasemonitoring/aws/aurora_test.go b/pkg/databasemonitoring/aws/aurora_test.go
@@ -211,7 +211,7 @@ func TestGetAuroraClusterEndpoints(t *testing.T) {
 		{
 			name: "multiple cluster ids returns single endpoint from API",
 			configureClient: func(k *MockrdsService) {
-				k.EXPECT().DescribeDBInstances(gomock.Any(), createDescribeDBInstancesRequest([]string{"test-cluster", "test-cluster-2"})).Return(&rds.DescribeDBInstancesOutput{
+				k.EXPECT().DescribeDBInstances(gomock.Any(), createDescribeDBInstancesRequest([]string{"test-cluster"})).Return(&rds.DescribeDBInstancesOutput{
 					DBInstances: []types.DBInstance{
 						{
 							Endpoint: &types.Endpoint{
@@ -226,6 +226,9 @@ func TestGetAuroraClusterEndpoints(t *testing.T) {
 						},
 					},
 				}, nil).Times(1)
+				k.EXPECT().DescribeDBInstances(gomock.Any(), createDescribeDBInstancesRequest([]string{"test-cluster-2"})).Return(&rds.DescribeDBInstancesOutput{
+					DBInstances: []types.DBInstance{},
+				}, nil).Times(1)
 			},
 			clusterIDs: []string{"test-cluster", "test-cluster-2"},
 			expectedAuroraClusterEndpoints: map[string]*AuroraCluster{
@@ -244,7 +247,7 @@ func TestGetAuroraClusterEndpoints(t *testing.T) {
 		{
 			name: "multiple cluster ids returns many endpoints from API",
 			configureClient: func(k *MockrdsService) {
-				k.EXPECT().DescribeDBInstances(gomock.Any(), createDescribeDBInstancesRequest([]string{"test-cluster", "test-cluster-2"})).Return(&rds.DescribeDBInstancesOutput{
+				k.EXPECT().DescribeDBInstances(gomock.Any(), createDescribeDBInstancesRequest([]string{"test-cluster"})).Return(&rds.DescribeDBInstancesOutput{
 					DBInstances: []types.DBInstance{
 						{
 							Endpoint: &types.Endpoint{
@@ -268,6 +271,11 @@ func TestGetAuroraClusterEndpoints(t *testing.T) {
 							DBInstanceStatus:                 aws.String("available"),
 							Engine:                           aws.String("aurora-postgresql"),
 						},
+					},
+				}, nil).Times(1)
+
+				k.EXPECT().DescribeDBInstances(gomock.Any(), createDescribeDBInstancesRequest([]string{"test-cluster-2"})).Return(&rds.DescribeDBInstancesOutput{
+					DBInstances: []types.DBInstance{
 						{
 							Endpoint: &types.Endpoint{
 								Address: aws.String("test-endpoint-3"),

diff --git a/releasenotes/notes/query-aurora-instances-per-cluster-8481d23d9e432e1a.yaml b/releasenotes/notes/query-aurora-instances-per-cluster-8481d23d9e432e1a.yaml
@@ -0,0 +1,12 @@
+# Each section from every release note are combined when the
+# CHANGELOG.rst is rendered. So the text needs to be worded so that
+# it does not depend on any information only available in another
+# section. This may mean repeating some details, but each section
+# must be readable independently of the other.
+#
+# Each section note must be formatted as reStructuredText.
+---
+enhancements:
+  - |
+    Query Aurora instances per cluster to allow up to 100 instances per cluster
+    rather than 100 instances total.