feat(warehouse): BQ test credentials refactoring

client/monte_carlo_client.go

@@ -108,10 +108,10 @@ type GetWarehouse struct {
 	} `json:"getWarehouse"`
 }
 
-const BigQueryConnectionType = "bigquery"
-const BigQueryConnectionTypeResponse = "BIGQUERY"
-const TransactionalConnectionType = "transactional-db"
-const TransactionalConnectionTypeResponse = "TRANSACTIONAL_DB"
+const BqConnectionType = "bigquery"
+const BqConnectionTypeResponse = "BIGQUERY"
+const TrxConnectionType = "transactional-db"
+const TrxConnectionTypeResponse = "TRANSACTIONAL_DB"
 const GetWarehouseQuery string = "query getWarehouse($uuid: UUID) { getWarehouse(uuid: $uuid) { name,connections{uuid,type,createdOn,updatedOn},dataCollector{uuid} } }"
 
 type RemoveConnection struct {
@@ -136,6 +136,13 @@ type UpdateCredentials struct {
 	} `graphql:"updateCredentials(changes: $changes, connectionId: $connectionId, shouldReplace: $shouldReplace, shouldValidate: $shouldValidate)"`
 }
 
+type UpdateCredentialsV2 struct {
+	UpdateCredentialsV2 struct {
+		Success   bool
+		UpdatedAt string
+	} `graphql:"updateCredentialsV2(connectionId: $connectionId, tempCredentialsKey: $tempCredentialsKey)"`
+}
+
 type TagPair struct {
 	Name  string `json:"name"`
 	Value string `json:"value"`

internal/warehouse/bigquery_warehouse.go

@@ -51,6 +51,13 @@ type BqCredentials struct {
 	UpdatedAt         types.String `tfsdk:"updated_at"`
 }
 
+func (m BigQueryWarehouseResourceModel) GetUuid() types.String          { return m.Uuid }
+func (m BigQueryWarehouseResourceModel) GetCollectorUuid() types.String { return m.CollectorUuid }
+func (m BigQueryWarehouseResourceModel) GetName() types.String          { return m.Name }
+func (m BigQueryWarehouseResourceModel) GetConnectionUuid() types.String {
+	return m.Credentials.ConnectionUuid
+}
+
 type BigQueryWarehouseResourceModelV1 struct {
 	Uuid               types.String `tfsdk:"uuid"`
 	ConnectionUuid     types.String `tfsdk:"connection_uuid"`
@@ -136,15 +143,15 @@ func (r *BigQueryWarehouseResource) Create(ctx context.Context, req resource.Cre
 		return
 	}
 
-	result, diags := r.addConnection(ctx, data)
+	result, diags := addConnection(ctx, r.client, r, data, client.BqConnectionType, BqKeyExtractor)
 	resp.Diagnostics.Append(diags...)
 	if result == nil {
 		return
 	}
 
-	data.Uuid = result.Uuid
-	data.Credentials.UpdatedAt = result.Credentials.UpdatedAt
-	data.Credentials.ConnectionUuid = result.Credentials.ConnectionUuid
+	data.Uuid = types.StringValue(result.AddConnection.Connection.Warehouse.Uuid)
+	data.Credentials.UpdatedAt = types.StringValue(result.AddConnection.Connection.CreatedOn)
+	data.Credentials.ConnectionUuid = types.StringValue(result.AddConnection.Connection.Uuid)
 	resp.Diagnostics.Append(resp.State.Set(ctx, &data)...)
 }
 
@@ -195,7 +202,7 @@ func (r *BigQueryWarehouseResource) Read(ctx context.Context, req resource.ReadR
 
 	for _, connection := range getResult.GetWarehouse.Connections {
 		if connection.Uuid == data.Credentials.ConnectionUuid.ValueString() {
-			if connection.Type != client.BigQueryConnectionTypeResponse {
+			if connection.Type != client.BqConnectionTypeResponse {
 				resp.Diagnostics.AddError(
 					fmt.Sprintf("Obtained Warehouse [uuid: %s, connection_uuid: %s] but got unexpected connection "+
 						"type '%s'.", data.Uuid.ValueString(), connection.Uuid, connection.Type),
@@ -243,37 +250,22 @@ func (r *BigQueryWarehouseResource) Update(ctx context.Context, req resource.Upd
 	}
 
 	if data.Credentials.ConnectionUuid.IsUnknown() || data.Credentials.ConnectionUuid.IsNull() {
-		if result, diags := r.addConnection(ctx, data); result != nil {
+		if result, diags := addConnection(ctx, r.client, r, data, client.BqConnectionType, BqKeyExtractor); result != nil {
 			resp.Diagnostics.Append(diags...)
-			data.Credentials.UpdatedAt = result.Credentials.UpdatedAt
-			data.Credentials.ConnectionUuid = result.Credentials.ConnectionUuid
+			data.Credentials.UpdatedAt = types.StringValue(result.AddConnection.Connection.CreatedOn)
+			data.Credentials.ConnectionUuid = types.StringValue(result.AddConnection.Connection.Uuid)
 		} else {
 			resp.Diagnostics.Append(diags...)
 			return
 		}
 	}
 
-	updateResult := client.UpdateCredentials{}
-	variables = map[string]interface{}{
-		"changes":        client.JSONString(data.Credentials.ServiceAccountKey.ValueString()),
-		"connectionId":   client.UUID(data.Credentials.ConnectionUuid.ValueString()),
-		"shouldReplace":  true,
-		"shouldValidate": true,
-	}
-
-	if err := r.client.Mutate(ctx, &updateResult, variables); err != nil {
-		toPrint := fmt.Sprintf("MC client 'UpdateCredentials' mutation result - %s", err.Error())
-		resp.Diagnostics.AddError(toPrint, "")
-		return
-	} else if !updateResult.UpdateCredentials.Success {
-		toPrint := "MC client 'UpdateCredentials' mutation - success = false, " +
-			"connection probably doesnt exists. Rerunning terraform operation usually helps."
-		resp.Diagnostics.AddError(toPrint, "")
-		return
+	if updateResult, diags := updateConnection(ctx, r.client, r, data, BqKeyExtractor); updateResult == nil {
+		resp.Diagnostics.Append(diags...)
+	} else {
+		data.Credentials.UpdatedAt = types.StringValue(updateResult.UpdateCredentialsV2.UpdatedAt)
+		resp.Diagnostics.Append(resp.State.Set(ctx, &data)...)
 	}
-
-	data.Credentials.UpdatedAt = types.StringValue(updateResult.UpdateCredentials.UpdatedAt)
-	resp.Diagnostics.Append(resp.State.Set(ctx, &data)...)
 }
 
 func (r *BigQueryWarehouseResource) Delete(ctx context.Context, req resource.DeleteRequest, resp *resource.DeleteResponse) {
@@ -319,7 +311,7 @@ func (r *BigQueryWarehouseResource) ImportState(ctx context.Context, req resourc
 	}
 }
 
-func (r *BigQueryWarehouseResource) addConnection(ctx context.Context, data BigQueryWarehouseResourceModel) (*BigQueryWarehouseResourceModel, diag.Diagnostics) {
+func (r *BigQueryWarehouseResource) testCredentials(ctx context.Context, data BigQueryWarehouseResourceModel) (*client.TestBqCredentialsV2, diag.Diagnostics) {
 	var diagsResult diag.Diagnostics
 	type BqConnectionDetails map[string]interface{}
 	testResult := client.TestBqCredentialsV2{}
@@ -341,40 +333,9 @@ func (r *BigQueryWarehouseResource) addConnection(ctx context.Context, data BigQ
 		diags = append(diags, bqTestDiagnosticToDiags(testResult.TestBqCredentialsV2.ValidationResult.Errors)...)
 		diagsResult.Append(diags...)
 		return nil, diagsResult
+	} else {
+		return &testResult, diagsResult
 	}
-
-	addResult := client.AddConnection{}
-	var name, createWarehouseType *string = nil, nil
-	warehouseUuid := data.Uuid.ValueStringPointer()
-	collectorUuid := data.CollectorUuid.ValueStringPointer()
-
-	if warehouseUuid == nil || *warehouseUuid == "" {
-		warehouseUuid = nil
-		name = data.Name.ValueStringPointer()
-		temp := client.BigQueryConnectionType
-		createWarehouseType = &temp
-	}
-
-	variables = map[string]interface{}{
-		"dcId":                (*client.UUID)(collectorUuid),
-		"dwId":                (*client.UUID)(warehouseUuid),
-		"key":                 testResult.TestBqCredentialsV2.Key,
-		"jobTypes":            []string{"metadata", "query_logs", "sql_query", "json_schema"},
-		"name":                name,
-		"connectionType":      client.BigQueryConnectionType,
-		"createWarehouseType": createWarehouseType,
-	}
-
-	if err := r.client.Mutate(ctx, &addResult, variables); err != nil {
-		toPrint := fmt.Sprintf("MC client 'AddConnection' mutation result - %s", err.Error())
-		diagsResult.AddError(toPrint, "")
-		return nil, diagsResult
-	}
-
-	data.Uuid = types.StringValue(addResult.AddConnection.Connection.Warehouse.Uuid)
-	data.Credentials.UpdatedAt = types.StringValue(addResult.AddConnection.Connection.CreatedOn)
-	data.Credentials.ConnectionUuid = types.StringValue(addResult.AddConnection.Connection.Uuid)
-	return &data, diagsResult
 }
 
 func bqTestDiagnosticToDiags[T client.BqTestWarnings | client.BqTestErrors](in T) diag.Diagnostics {

internal/warehouse/common.go

@@ -0,0 +1,108 @@
+package warehouse
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/kiwicom/terraform-provider-montecarlo/client"
+
+	"github.com/hashicorp/terraform-plugin-framework/diag"
+	"github.com/hashicorp/terraform-plugin-framework/types"
+)
+
+type WarehouseResource[T WarehouseResourceModel, K TestCredentials] interface {
+	*BigQueryWarehouseResource | *TransactionalWarehouseResource
+	testCredentials(ctx context.Context, data T) (K, diag.Diagnostics)
+}
+
+type WarehouseResourceModel interface {
+	BigQueryWarehouseResourceModel | TransactionalWarehouseResourceModel
+	GetUuid() types.String
+	GetCollectorUuid() types.String
+	GetName() types.String
+	GetConnectionUuid() types.String
+}
+
+type TestCredentials interface {
+	*client.TestBqCredentialsV2 | *client.TestDatabaseCredentials
+}
+
+func BqKeyExtractor(k *client.TestBqCredentialsV2) string {
+	return k.TestBqCredentialsV2.Key
+}
+func TrxKeyExtractor(k *client.TestDatabaseCredentials) string {
+	return k.TestDatabaseCredentials.Key
+}
+
+func addConnection[T WarehouseResource[J, K], J WarehouseResourceModel, K TestCredentials](
+	ctx context.Context, mcClient client.MonteCarloClient, warehouse T, data J, connectionType string, keyExtractor func(K) string,
+) (*client.AddConnection, diag.Diagnostics) {
+	var diagsResult diag.Diagnostics
+	testResult, credentialsDiags := warehouse.testCredentials(ctx, data)
+
+	if testResult == nil {
+		diagsResult.Append(credentialsDiags...)
+		return nil, diagsResult
+	}
+
+	addResult := client.AddConnection{}
+	var name, createWarehouseType *string = nil, nil
+	warehouseUuid := data.GetUuid().ValueStringPointer()
+	collectorUuid := data.GetCollectorUuid().ValueStringPointer()
+
+	if warehouseUuid == nil || *warehouseUuid == "" {
+		warehouseUuid = nil
+		name = data.GetName().ValueStringPointer()
+		temp := connectionType
+		createWarehouseType = &temp
+	}
+
+	variables := map[string]interface{}{
+		"dcId":                (*client.UUID)(collectorUuid),
+		"dwId":                (*client.UUID)(warehouseUuid),
+		"key":                 keyExtractor(testResult),
+		"jobTypes":            []string{"metadata", "query_logs", "sql_query", "json_schema"},
+		"name":                name,
+		"connectionType":      connectionType,
+		"createWarehouseType": createWarehouseType,
+	}
+
+	if err := mcClient.Mutate(ctx, &addResult, variables); err != nil {
+		toPrint := fmt.Sprintf("MC client 'AddConnection' mutation result - %s", err.Error())
+		diagsResult.AddError(toPrint, "")
+		return nil, diagsResult
+	} else {
+		return &addResult, diagsResult
+	}
+}
+
+func updateConnection[T WarehouseResource[J, K], J WarehouseResourceModel, K TestCredentials](
+	ctx context.Context, mcClient client.MonteCarloClient, warehouse T, data J, keyExtractor func(K) string,
+) (*client.UpdateCredentialsV2, diag.Diagnostics) {
+	var diagsResult diag.Diagnostics
+	testResult, credentialsDiags := warehouse.testCredentials(ctx, data)
+
+	if testResult == nil {
+		diagsResult.Append(credentialsDiags...)
+		return nil, diagsResult
+	}
+
+	updateResult := client.UpdateCredentialsV2{}
+	variables := map[string]interface{}{
+		"connectionId":       client.UUID(data.GetConnectionUuid().ValueString()),
+		"tempCredentialsKey": keyExtractor(testResult),
+	}
+
+	if err := mcClient.Mutate(ctx, &updateResult, variables); err != nil {
+		toPrint := fmt.Sprintf("MC client 'UpdateCredentials' mutation result - %s", err.Error())
+		diagsResult.AddError(toPrint, "")
+		return nil, diagsResult
+	} else if !updateResult.UpdateCredentialsV2.Success {
+		toPrint := "MC client 'UpdateCredentials' mutation - success = false, " +
+			"connection probably doesnt exists. Rerunning terraform operation usually helps."
+		diagsResult.AddError(toPrint, "")
+		return nil, diagsResult
+	} else {
+		return &updateResult, diagsResult
+	}
+}

internal/warehouse/testdata/TestAccTransactionalWarehouseResource/create.tf

@@ -34,7 +34,7 @@ variable "pg_password" {
 }
 
 resource "montecarlo_transactional_warehouse" "test" {
-  name                = "name1"
+  name                = "test-warehouse"
   collector_uuid      = "a08d23fc-00a0-4c36-b568-82e9d0e67ad8"
   db_type             = "POSTGRES" # POSTGRES | MYSQL | SQL-SERVER
   deletion_protection = false

internal/warehouse/testdata/TestAccTransactionalWarehouseResource/update.tf

@@ -34,7 +34,7 @@ variable "pg_password" {
 }
 
 resource "montecarlo_transactional_warehouse" "test" {
-  name                = "name1"
+  name                = "test-warehouse-updated"
   collector_uuid      = "a08d23fc-00a0-4c36-b568-82e9d0e67ad8"
   db_type             = "POSTGRES" # POSTGRES | MYSQL | SQL-SERVER
   deletion_protection = false