Cosmos Go SDK Custom Header Injection

custom_header_injection.go

package main

import (
    "context"
    "fmt"
    "net/http"

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

// Context key type (unexported to avoid collisions)
type ctxKey string

// Key used to signal whether the header should be injected.
// NOTE: This is only needed if you want conditional behavior.
const injectHeaderKey ctxKey = "injectHeader"

// Custom policy that conditionally injects a header based on context.
type conditionalHeaderPolicy struct{}

func (p *conditionalHeaderPolicy) Do(req *policy.Request) (*http.Response, error) {
    // Read the flag from context.
    // If you don't need conditional behavior, you can skip all of this
    // and always inject the header.
    val := req.Raw().Context().Value(injectHeaderKey)

    if shouldInject, ok := val.(bool); ok && shouldInject {
        req.Raw().Header.Set("x-ms-some-header", "some-value")
    }

    return req.Next()
}

func main() {
    endpoint := "https://your-account.documents.azure.com:443/"

    cred, err := azidentity.NewDefaultAzureCredential(nil)
    if err != nil {
        panic(err)
    }

    clientOpts := &azcore.ClientOptions{
        PerCallPolicies: []policy.Policy{
            &conditionalHeaderPolicy{},
        },
    }

    // Cosmos DB client using Entra ID
    client, err := azcosmos.NewClient(endpoint, cred, clientOpts)
    if err != nil {
        panic(err)
    }

    db, err := client.NewDatabase("mydb")
    if err != nil {
        panic(err)
    }

    // ------------------------------------------------------------
    // Example 1: Inject the header
    // ------------------------------------------------------------
    // Only needed if you want conditional behavior.
    ctxWithHeader := context.WithValue(context.Background(), injectHeaderKey, true)

    _, err = db.Read(ctxWithHeader, nil)
    if err != nil {
        fmt.Println("Read failed:", err)
    } else {
        fmt.Println("Read succeeded WITH header injection")
    }

    // ------------------------------------------------------------
    // Example 2: Do NOT inject the header
    // ------------------------------------------------------------
    // Again, only needed if you want conditional behavior.
    ctxWithoutHeader := context.WithValue(context.Background(), injectHeaderKey, false)

    _, err = db.Read(ctxWithoutHeader, nil)
    if err != nil {
        fmt.Println("Read failed:", err)
    } else {
        fmt.Println("Read succeeded WITHOUT header injection")
    }
}