feat(warehouse): BQ test credentials refactoring #2

kiwicom · Feb 18, 2024 · 3366d14 · 3366d14
1 parent 0bf7848
commit 3366d14
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 20 deletions.
diff --git a/client/monte_carlo_client.go b/client/monte_carlo_client.go
@@ -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:"updateCredentials(connectionId: $connectionId, tempCredentialsKey: $tempCredentialsKey)"`
+}
+
 type TagPair struct {
 	Name  string `json:"name"`
 	Value string `json:"value"`

diff --git a/internal/warehouse/bigquery_warehouse.go b/internal/warehouse/bigquery_warehouse.go
@@ -253,27 +253,12 @@ func (r *BigQueryWarehouseResource) Update(ctx context.Context, req resource.Upd
 		}
 	}
 
-	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 := r.updateConnection(ctx, data); 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) {
@@ -389,6 +374,35 @@ func (r *BigQueryWarehouseResource) addConnection(ctx context.Context, data BigQ
 	return &data, diagsResult
 }
 
+func (r *BigQueryWarehouseResource) updateConnection(ctx context.Context, data BigQueryWarehouseResourceModel) (*client.UpdateCredentialsV2, diag.Diagnostics) {
+	var diagsResult diag.Diagnostics
+	testResult, credentialsDiags := r.testCredentials(ctx, data)
+
+	if testResult == nil {
+		diagsResult.Append(credentialsDiags...)
+		return nil, diagsResult
+	}
+
+	updateResult := client.UpdateCredentialsV2{}
+	variables := map[string]interface{}{
+		"connectionId":       client.UUID(data.Credentials.ConnectionUuid.ValueString()),
+		"tempCredentialsKey": testResult.TestBqCredentialsV2.Key,
+	}
+
+	if err := r.client.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
+	}
+}
+
 func bqTestDiagnosticToDiags[T client.BqTestWarnings | client.BqTestErrors](in T) diag.Diagnostics {
 	var diags diag.Diagnostics
 	switch any(in).(type) {