cleanup between data source and resource

DataDog · Jan 17, 2025 · 0d42495 · 0d42495
1 parent b89cdd4
commit 0d42495
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 43 deletions.
diff --git a/datadog/fwprovider/data_source_datadog_connection.go b/datadog/fwprovider/data_source_datadog_connection.go
@@ -2,7 +2,6 @@ package fwprovider
 
 import (
 	"context"
-	"io"
 
 	"github.com/DataDog/datadog-api-client-go/v2/api/datadogV2"
 	"github.com/hashicorp/terraform-plugin-framework/datasource"
@@ -161,16 +160,9 @@ func (d *connectionDatasource) Read(ctx context.Context, request datasource.Read
 		return
 	}
 
-	conn, httpResponse, err := d.Api.GetActionConnection(d.Auth, state.ID.ValueString())
+	connModel, err := readConnection(d.Api, d.Auth, state.ID.ValueString(), state)
 	if err != nil {
-		body, _ := io.ReadAll(httpResponse.Body)
-		response.Diagnostics.AddError("Could not get connection", string(body))
-		return
-	}
-
-	connModel, err := apiResponseToConnectionModel(conn)
-	if err != nil {
-		response.Diagnostics.AddError(err.Error(), "")
+		response.Diagnostics.AddError("Could not read connection", err.Error())
 		return
 	}
 

diff --git a/datadog/fwprovider/resource_datadog_connection.go b/datadog/fwprovider/resource_datadog_connection.go
@@ -355,39 +355,12 @@ func (r *connectionResource) Read(ctx context.Context, request resource.ReadRequ
 		return
 	}
 
-	conn, httpResponse, err := r.Api.GetActionConnection(r.Auth, state.ID.ValueString())
+	connModel, err := readConnection(r.Api, r.Auth, state.ID.ValueString(), state)
 	if err != nil {
-		if httpResponse != nil {
-			body, err := io.ReadAll(httpResponse.Body)
-			if err != nil {
-				response.Diagnostics.AddError("Could not read API error response", "")
-				return
-			}
-			response.Diagnostics.AddError("Could not get connection", string(body))
-		} else {
-			response.Diagnostics.AddError("Could not get connection", err.Error())
-		}
+		response.Diagnostics.AddError("Could not read connection", err.Error())
 		return
 	}
 
-	connModel, err := apiResponseToConnectionModel(conn)
-	if err != nil {
-		response.Diagnostics.AddError(err.Error(), "")
-		return
-	}
-
-	// The API response does not include the token value, so this code gets it from the state.
-	// This is used to determine whether the token value changed since the last update.
-	if state.HTTP != nil {
-		for _, stateToken := range state.HTTP.TokenAuth.Tokens {
-			for _, responseToken := range connModel.HTTP.TokenAuth.Tokens {
-				if stateToken.Name.Equal(responseToken.Name) {
-					responseToken.Value = stateToken.Value
-				}
-			}
-		}
-	}
-
 	diags = response.State.Set(ctx, connModel)
 	response.Diagnostics.Append(diags...)
 }
@@ -474,10 +447,6 @@ func apiResponseToConnectionModel(connection datadogV2.GetActionConnectionRespon
 
 	if attributes.Integration.AWSIntegration != nil {
 		awsAttr := attributes.Integration.AWSIntegration
-		if awsAttr.Credentials.AWSAssumeRole == nil {
-			err := fmt.Errorf("this provider only supports AWS connections of the assume role type")
-			return nil, err
-		}
 
 		connModel.AWS = &awsConnectionModel{
 			AssumeRole: &awsAssumeRoleConnectionModel{
@@ -753,3 +722,38 @@ func buildHttpDeletions(plan, oldState connectionResourceModel, updateModel *dat
 		updateModel.UrlParameters = append(updateModel.UrlParameters, *paramUpdate)
 	}
 }
+
+// Read logic is shared between data source and resource
+func readConnection(api *datadogV2.ActionConnectionApi, authCtx context.Context, id string, currentState connectionResourceModel) (*connectionResourceModel, error) {
+	conn, httpResponse, err := api.GetActionConnection(authCtx, id)
+	if err != nil {
+		if httpResponse != nil {
+			body, err := io.ReadAll(httpResponse.Body)
+			if err != nil {
+				return nil, fmt.Errorf("could not read error response")
+			}
+			return nil, fmt.Errorf("%s", body)
+		} else {
+			return nil, fmt.Errorf("%s", err.Error())
+		}
+	}
+
+	connModel, err := apiResponseToConnectionModel(conn)
+	if err != nil {
+		return nil, err
+	}
+
+	// The API response does not include the token value, so this code gets it from the state.
+	// This is used to determine whether the token value changed since the last update.
+	if currentState.HTTP != nil {
+		for _, stateToken := range currentState.HTTP.TokenAuth.Tokens {
+			for _, responseToken := range connModel.HTTP.TokenAuth.Tokens {
+				if stateToken.Name.Equal(responseToken.Name) {
+					responseToken.Value = stateToken.Value
+				}
+			}
+		}
+	}
+
+	return connModel, nil
+}