chore: Implement full wait functionality for long-running operations (#3305)

* wip

* update templates

* state attribute with JSON naming

* improve not found logic

* delete TODO comment

* revert timeout

* typo

* error if no pending states

Author: lantoli

internal/common/autogen/handle_operations.go

@@ -1,14 +1,18 @@
 package autogen
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"time"
 
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/retry"
+	"github.com/mongodb/terraform-provider-mongodbatlas/internal/common/retrystrategy"
 	"github.com/mongodb/terraform-provider-mongodbatlas/internal/common/validate"
 	"github.com/mongodb/terraform-provider-mongodbatlas/internal/config"
 )
@@ -45,27 +49,27 @@ type HandleCreateReq struct {
 
 func HandleCreate(ctx context.Context, req HandleCreateReq) {
 	d := &req.Resp.Diagnostics
-	reqBody, err := Marshal(req.Plan, false)
+	bodyReq, err := Marshal(req.Plan, false)
 	if err != nil {
 		addError(d, opCreate, errBuildingAPIRequest, err)
 		return
 	}
-	respBody, err := callAPIWithBody(ctx, req.Client, req.CallParams, reqBody)
+	bodyResp, err := callAPIWithBody(ctx, req.Client, req.CallParams, bodyReq)
 	if err != nil {
 		addError(d, opCreate, errCallingAPI, err)
 		return
 	}
 
 	// Use the plan as the base model to set the response state
-	if err := Unmarshal(respBody, req.Plan); err != nil {
+	if err := Unmarshal(bodyResp, req.Plan); err != nil {
 		addError(d, opCreate, errUnmarshallingResponse, err)
 		return
 	}
 	if err := ResolveUnknowns(req.Plan); err != nil {
 		addError(d, opCreate, errResolvingResponse, err)
 		return
 	}
-	if err := handleWait(ctx, req.Wait, req.Client, req.Plan); err != nil {
+	if err := handleWaitCreateUpdate(ctx, req.Wait, req.Client, req.Plan); err != nil {
 		addError(d, opCreate, errWaitingForChanges, err)
 		return
 	}
@@ -81,18 +85,18 @@ type HandleReadReq struct {
 
 func HandleRead(ctx context.Context, req HandleReadReq) {
 	d := &req.Resp.Diagnostics
-	respBody, apiResp, err := callAPIWithoutBody(ctx, req.Client, req.CallParams)
+	bodyResp, apiResp, err := callAPIWithoutBody(ctx, req.Client, req.CallParams)
+	if notFound(bodyResp, apiResp) {
+		req.Resp.State.RemoveResource(ctx)
+		return
+	}
 	if err != nil {
-		if validate.StatusNotFound(apiResp) {
-			req.Resp.State.RemoveResource(ctx)
-			return
-		}
 		addError(d, opRead, errCallingAPI, err)
 		return
 	}
 
 	// Use the current state as the base model to set the response state
-	if err := Unmarshal(respBody, req.State); err != nil {
+	if err := Unmarshal(bodyResp, req.State); err != nil {
 		addError(d, opRead, errUnmarshallingResponse, err)
 		return
 	}
@@ -113,27 +117,27 @@ type HandleUpdateReq struct {
 
 func HandleUpdate(ctx context.Context, req HandleUpdateReq) {
 	d := &req.Resp.Diagnostics
-	reqBody, err := Marshal(req.Plan, true)
+	bodyReq, err := Marshal(req.Plan, true)
 	if err != nil {
 		addError(d, opUpdate, errBuildingAPIRequest, err)
 		return
 	}
-	respBody, err := callAPIWithBody(ctx, req.Client, req.CallParams, reqBody)
+	bodyResp, err := callAPIWithBody(ctx, req.Client, req.CallParams, bodyReq)
 	if err != nil {
 		addError(d, opUpdate, errCallingAPI, err)
 		return
 	}
 
 	// Use the plan as the base model to set the response state
-	if err := Unmarshal(respBody, req.Plan); err != nil {
+	if err := Unmarshal(bodyResp, req.Plan); err != nil {
 		addError(d, opUpdate, errUnmarshallingResponse, err)
 		return
 	}
 	if err := ResolveUnknowns(req.Plan); err != nil {
 		addError(d, opUpdate, errResolvingResponse, err)
 		return
 	}
-	if err := handleWait(ctx, req.Wait, req.Client, req.Plan); err != nil {
+	if err := handleWaitCreateUpdate(ctx, req.Wait, req.Client, req.Plan); err != nil {
 		addError(d, opUpdate, errWaitingForChanges, err)
 		return
 	}
@@ -154,70 +158,128 @@ func HandleDelete(ctx context.Context, req HandleDeleteReq) {
 		addError(d, opDelete, errCallingAPI, err)
 		return
 	}
-	// don't consider wait errors as it can happen that the resource is already deleted
-	_ = handleWait(ctx, req.Wait, req.Client, req.State)
+	if err := handleWaitDelete(ctx, req.Wait, req.Client); err != nil {
+		addError(d, opDelete, errWaitingForChanges, err)
+	}
 }
 
-// handleWait waits until a long-running operation is done if needed.
+// handleWaitCreateUpdate waits until a long-running operation is done if needed.
 // It also updates the model with the latest JSON response from the API.
-func handleWait(ctx context.Context, wait *WaitReq, client *config.MongoDBClient, model any) error {
+func handleWaitCreateUpdate(ctx context.Context, wait *WaitReq, client *config.MongoDBClient, model any) error {
 	if wait == nil {
 		return nil
 	}
-	respBody, err := waitForChanges(ctx, wait, client)
-	if err != nil {
+	bodyResp, err := waitForChanges(ctx, wait, client)
+	if err != nil || isEmptyJSON(bodyResp) {
 		return err
 	}
-	if respBody == nil {
+	if err := Unmarshal(bodyResp, model); err != nil {
+		return err
+	}
+	return ResolveUnknowns(model)
+}
+
+// handleWaitDelete waits until a long-running operation to delete a resource if neeed.
+func handleWaitDelete(ctx context.Context, wait *WaitReq, client *config.MongoDBClient) error {
+	if wait == nil {
 		return nil
 	}
-	if err := Unmarshal(respBody, model); err != nil {
+	if _, err := waitForChanges(ctx, wait, client); err != nil {
 		return err
 	}
-	return ResolveUnknowns(model)
+	return nil
 }
 
-// waitForChanges waits until a long-running operation is done.
-// It returns the latest JSON response from the API so it can be used to update the response state.
-// TODO: This is a basic implementation, it will be replaced in CLOUDP-314960.
-func waitForChanges(ctx context.Context, wait *WaitReq, client *config.MongoDBClient) ([]byte, error) {
-	time.Sleep(time.Duration(wait.TimeoutSeconds) * time.Second) // TODO: TimeoutSeconds is temporarily used to allow time to destroy the resource until autogen long-running operations are supported in CLOUDP-314960
-	respBody, _, err := callAPIWithoutBody(ctx, client, wait.CallParams)
-	return respBody, err
+func addError(d *diag.Diagnostics, opName, errSummary string, err error) {
+	d.AddError(fmt.Sprintf("Error %s in %s", errSummary, opName), err.Error())
 }
 
 // callAPIWithBody makes a request to the API with the given request body and returns the response body.
 // It is used for POST, PUT, and PATCH requests where a request body is required.
-func callAPIWithBody(ctx context.Context, client *config.MongoDBClient, callParams *config.APICallParams, reqBody []byte) ([]byte, error) {
-	callParams.Body = reqBody
-	apiResp, err := client.UntypedAPICall(ctx, callParams)
+func callAPIWithBody(ctx context.Context, client *config.MongoDBClient, callParams *config.APICallParams, bodyReq []byte) ([]byte, error) {
+	apiResp, err := client.UntypedAPICall(ctx, callParams, bodyReq)
 	if err != nil {
 		return nil, err
 	}
-	respBody, err := io.ReadAll(apiResp.Body)
+	bodyResp, err := io.ReadAll(apiResp.Body)
 	apiResp.Body.Close()
 	if err != nil {
 		return nil, err
 	}
-	return respBody, nil
+	return bodyResp, nil
 }
 
 // callAPIWithoutBody makes a request to the API without a request body and returns the response body.
 // It is used for GET or DELETE requests where no request body is required.
 func callAPIWithoutBody(ctx context.Context, client *config.MongoDBClient, callParams *config.APICallParams) ([]byte, *http.Response, error) {
-	callParams.Body = nil
-	apiResp, err := client.UntypedAPICall(ctx, callParams)
+	apiResp, err := client.UntypedAPICall(ctx, callParams, nil)
 	if err != nil {
 		return nil, apiResp, err
 	}
-	respBody, err := io.ReadAll(apiResp.Body)
+	bodyResp, err := io.ReadAll(apiResp.Body)
 	apiResp.Body.Close()
 	if err != nil {
 		return nil, apiResp, err
 	}
-	return respBody, apiResp, nil
+	return bodyResp, apiResp, nil
 }
 
-func addError(d *diag.Diagnostics, opName, errSummary string, err error) {
-	d.AddError(fmt.Sprintf("Error %s in %s", errSummary, opName), err.Error())
+// waitForChanges waits until a long-running operation is done.
+// It returns the latest JSON response from the API so it can be used to update the response state.
+func waitForChanges(ctx context.Context, wait *WaitReq, client *config.MongoDBClient) ([]byte, error) {
+	if len(wait.TargetStates) == 0 {
+		return nil, fmt.Errorf("wait must have at least one target state, pending states: %v", wait.PendingStates)
+	}
+	stateConf := retry.StateChangeConf{
+		Target:     wait.TargetStates,
+		Pending:    wait.PendingStates,
+		Timeout:    time.Duration(wait.TimeoutSeconds) * time.Second,
+		MinTimeout: time.Duration(wait.MinTimeoutSeconds) * time.Second,
+		Delay:      time.Duration(wait.DelaySeconds) * time.Second,
+		Refresh:    refreshFunc(ctx, wait, client),
+	}
+	bodyResp, err := stateConf.WaitForStateContext(ctx)
+	if err != nil || bodyResp == nil {
+		return nil, err
+	}
+	return bodyResp.([]byte), err
 }
+
+// refreshFunc retries until a target state or error happens.
+// It uses a special state value of "DELETED" when the API returns 404 or empty object
+func refreshFunc(ctx context.Context, wait *WaitReq, client *config.MongoDBClient) retry.StateRefreshFunc {
+	return func() (result any, state string, err error) {
+		bodyResp, httpResp, err := callAPIWithoutBody(ctx, client, wait.CallParams)
+		if notFound(bodyResp, httpResp) {
+			return emptyJSON, retrystrategy.RetryStrategyDeletedState, nil
+		}
+		if err != nil {
+			return nil, "", err
+		}
+		var objJSON map[string]any
+		if err := json.Unmarshal(bodyResp, &objJSON); err != nil {
+			return nil, "", err
+		}
+		stateValAny, found := objJSON[wait.StateAttribute]
+		if !found {
+			return nil, "", fmt.Errorf("wait state attribute not found: %s", wait.StateAttribute)
+		}
+		stateValStr, ok := stateValAny.(string)
+		if !ok {
+			return nil, "", fmt.Errorf("wait state attribute value is not a string, attribute name: %s, value: %s", wait.StateAttribute, stateValAny)
+		}
+		return bodyResp, stateValStr, nil
+	}
+}
+
+// notFound returns if the resource is not found (API response is 404 or response body is empty JSON).
+// That is because some resources like search_deployment can return an ok status code with empty json when resource doesn't exist.
+func notFound(bodyResp []byte, apiResp *http.Response) bool {
+	return validate.StatusNotFound(apiResp) || isEmptyJSON(bodyResp)
+}
+
+func isEmptyJSON(raw []byte) bool {
+	return len(raw) == 0 || bytes.Equal(raw, emptyJSON)
+}
+
+var emptyJSON = []byte("{}")

internal/config/client.go

@@ -241,10 +241,9 @@ type APICallParams struct {
 	RelativePath  string
 	PathParams    map[string]string
 	Method        string
-	Body          []byte
 }
 
-func (c *MongoDBClient) UntypedAPICall(ctx context.Context, params *APICallParams) (*http.Response, error) {
+func (c *MongoDBClient) UntypedAPICall(ctx context.Context, params *APICallParams, bodyReq []byte) (*http.Response, error) {
 	localBasePath, _ := c.AtlasV2.GetConfig().ServerURLWithContext(ctx, "")
 	localVarPath := localBasePath + params.RelativePath
 
@@ -256,11 +255,11 @@ func (c *MongoDBClient) UntypedAPICall(ctx context.Context, params *APICallParam
 	headerParams["Content-Type"] = params.VersionHeader
 	headerParams["Accept"] = params.VersionHeader
 
-	var reqBody any
-	if params.Body != nil { // if nil slice is sent with application/json content type SDK method returns an error
-		reqBody = params.Body
+	var bodyPost any
+	if bodyReq != nil { // if nil slice is sent with application/json content type SDK method returns an error
+		bodyPost = bodyReq
 	}
-	apiReq, err := c.AtlasV2.PrepareRequest(ctx, localVarPath, params.Method, reqBody, headerParams, nil, nil, nil)
+	apiReq, err := c.AtlasV2.PrepareRequest(ctx, localVarPath, params.Method, bodyPost, headerParams, nil, nil, nil)
 	if err != nil {
 		return nil, err
 	}