simulate cosmos db error with custom HTTP transport and log it with custom retry policy

simulate_cosmosdb_error.go

package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	azlog "github.com/Azure/azure-sdk-for-go/sdk/azcore/log"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
)

func init() {
	azlog.SetListener(func(cls azlog.Event, msg string) {

		// Log retry-related events
		switch cls {
		case azlog.EventRetryPolicy:
			fmt.Printf("Retry Policy Event: %s\n", msg)

		}
	})

	// Set logging level to include retries
	// azlog.SetEvents(azlog.EventRetryPolicy, azlog.EventResponse, azlog.EventRequest)
	azlog.SetEvents(azlog.EventRetryPolicy)

}

// CustomTransport403Error implements policy.Transporter to simulate 403 errors only for ReadItem operations
type CustomTransport403Error struct{}

func (t *CustomTransport403Error) Do(req *http.Request) (*http.Response, error) {

	// Check if this is a ReadItem operation (typically a GET request with an item id in the path)
	// ReadItem URLs look like: /dbs/{db}/colls/{container}/docs/{id}
	isReadItemOperation := req.Method == "GET" && strings.Contains(req.URL.Path, "/docs/")

	if isReadItemOperation {
		fmt.Printf("CustomTransport403Error: Simulating 403 error for ReadItem operation: %s\n", req.URL.String())

		// Create a simulated 403 response with sub-status 3
		header := make(http.Header)
		header.Set("x-ms-substatus", "3")
		header.Set("x-ms-activity-id", "readitem-test-activity-id")
		header.Set("x-ms-request-id", "readitem-test-request-id")
		header.Set("Content-Type", "application/json")

		response := &http.Response{
			StatusCode: 403,
			Status:     "403 Forbidden",
			Header:     header,
			Body:       io.NopCloser(strings.NewReader(`{"code": "Forbidden", "message": "Simulated 403 error for ReadItem with sub-status 3"}`)),
			Request:    req,
		}

		// Return both the response and the error so the SDK can handle it properly
		responseErr := azruntime.NewResponseError(response)
		return response, responseErr
	}

	// For all other operations (like account properties), use a fake successful response
	fmt.Printf("CustomTransport403Error: Allowing operation: %s %s\n", req.Method, req.URL.String())

	// Create a fake successful response for account properties and other operations
	header := make(http.Header)
	header.Set("Content-Type", "application/json")
	header.Set("x-ms-activity-id", "success-activity-id")
	header.Set("x-ms-request-id", "success-request-id")

	response := &http.Response{
		StatusCode: 200,
		Status:     "200 OK",
		Header:     header,
		Body:       io.NopCloser(strings.NewReader("")),
		Request:    req,
	}

	return response, nil
}

// RetryLoggingPolicy logs error details during retries
type RetryLoggingPolicy struct{}

func (p *RetryLoggingPolicy) Do(req *policy.Request) (*http.Response, error) {
	// fmt.Println("RetryLoggingPolicy: Starting retry with request URL:", req.Raw().URL.String())

	// Call the next policy in the chain
	resp, err := req.Next()

	// If there's an error, log the details
	if err != nil {
		var azErr *azcore.ResponseError
		if errors.As(err, &azErr) {
			subStatus := azErr.RawResponse.Header.Get("x-ms-substatus")
			if subStatus == "" {
				subStatus = "N/A"
			}

			fmt.Printf("RetryLoggingPolicy: ResponseError during retry - Status: %d, SubStatus: %s, URL: %s\n",
				azErr.StatusCode, subStatus, req.Raw().URL.String())
		} else {
			fmt.Printf("RetryLoggingPolicy: Non-ResponseError during retry - %T: %v, URL: %s\n",
				err, err, req.Raw().URL.String())
		}
	} else if resp != nil && resp.StatusCode >= 400 {
		// Log HTTP error responses even if they don't result in Go errors
		subStatus := resp.Header.Get("x-ms-substatus")
		if subStatus == "" {
			subStatus = "N/A"
		}

		fmt.Printf("RetryLoggingPolicy: HTTP error response - Status: %d, SubStatus: %s, URL: %s\n",
			resp.StatusCode, subStatus, req.Raw().URL.String())
	}

	return resp, err
}

func main() {

	opts := &azcosmos.ClientOptions{
		ClientOptions: azcore.ClientOptions{
			PerRetryPolicies: []policy.Policy{
				&RetryLoggingPolicy{}, // This will log error details during retries
			},
			Transport: &CustomTransport403Error{}, // Use the selective transport to simulate 403 errors only for ReadItem
		},
	}

	creds, _ := azidentity.NewDefaultAzureCredential(nil)
	client, err := azcosmos.NewClient("https://i_dont_exist.documents.azure.com:443", creds, opts)
	if err != nil {
		fmt.Printf("NewClient Error occurred: %v\n", err)
		return
	}

	// Test the ReadItem operation
	container, err := client.NewContainer("dummy", "dummy")
	if err != nil {
		fmt.Printf("NewContainer Error occurred: %v\n", err)
		return
	}
	partitionKey := azcosmos.NewPartitionKeyString("testpk")

	_, err = container.ReadItem(context.Background(), partitionKey, "testid", nil)
	handlerError(err)

}

func handlerError(err error) {
	if err != nil {
		fmt.Println("ReadItem Error occurred")

		// Debug: Print the actual error type
		fmt.Printf("Error type: %T\n", err)
		// fmt.Printf("Error value: %v\n", err)

		var azErr *azcore.ResponseError

		if errors.As(err, &azErr) {
			fmt.Println("Successfully unwrapped to azcore.ResponseError using errors.As")
			fmt.Printf("error status code: %d\n", azErr.StatusCode)

			subStatus := azErr.RawResponse.Header.Get("x-ms-substatus")
			if subStatus == "" {
				subStatus = "N/A"
			}
			fmt.Printf("error sub-status code: %s\n", subStatus)
		}
	}
}