feat: avoiding throttling & improve performance

Author: cmendible

go.mod

@@ -39,7 +39,6 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/postgresql/armpostgresql v1.2.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/postgresql/armpostgresqlflexibleservers v1.1.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/redis/armredis v1.0.0
-	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph v0.9.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.2.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/servicebus/armservicebus v1.2.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/signalr/armsignalr v1.2.0

go.sum

@@ -78,8 +78,6 @@ github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/postgresql/armpostgresqlfl
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/postgresql/armpostgresqlflexibleservers v1.1.0/go.mod h1:nKcJObAisSPDrO9lMuuCBoYY7Ki7ADt8p6XmBhpKNTk=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/redis/armredis v1.0.0 h1:nmpTBgRg1HynngFYICRhceC7s5dmbKN9fJ/XQz/UQ2I=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/redis/armredis v1.0.0/go.mod h1:3yjiOtnkVociBTlF7UZrwAGfJrGaOCsvtVS4HzNajxQ=
-github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph v0.9.0 h1:zLzoX5+W2l95UJoVwiyNS4dX8vHyQ6x2xRLoBBL9wMk=
-github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph v0.9.0/go.mod h1:wVEOJfGTj0oPAUGA1JuRAvz/lxXQsWW16axmHPP47Bk=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.2.0 h1:Dd+RhdJn0OTtVGaeDLZpcumkIVCtA/3/Fo42+eoYvVM=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.2.0/go.mod h1:5kakwfW5CjC9KK+Q4wjXAg+ShuIm2mBMua0ZFj2C8PE=
 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/servicebus/armservicebus v1.2.0 h1:jngSeKBnzC7qIk3rvbWHsLI7eeasEucORHWr2CHX0Yg=

internal/graph/graph.go

@@ -4,71 +4,137 @@
 package graph
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/Azure/azqr/internal/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
-	arg "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/rs/zerolog/log"
 )
 
-type (
-	GraphQuery struct {
-		client *arg.Client
-	}
+// GraphQueryClient provides methods to query Azure Resource Graph using HTTP client.
+type GraphQueryClient struct {
+	httpClient  *http.Client // HTTP client for making requests
+	endpoint    string       // Resource Graph endpoint URL
+	accessToken string       // Bearer token for authentication
+}
+
+// GraphResult holds the data returned from a Resource Graph query.
+type GraphResult struct {
+	Data []interface{} // Query result data
+}
 
-	GraphResult struct {
-		Data  []interface{}
+// QueryRequestOptions represents options for the Resource Graph query.
+type QueryRequestOptions struct {
+	ResultFormat string  `json:"resultFormat,omitempty"` // Format of the result
+	Top          *int32  `json:"top,omitempty"`          // Max number of results
+	SkipToken    *string `json:"skipToken,omitempty"`    // Token for pagination
+}
+
+// QueryRequest represents the payload for a Resource Graph query.
+type QueryRequest struct {
+	Subscriptions []string             `json:"subscriptions"` // List of subscription IDs
+	Query         string               `json:"query"`         // Kusto query string
+	Options       *QueryRequestOptions `json:"options"`       // Query options
+}
+
+// QueryResponse represents the response from the Resource Graph API.
+type QueryResponse struct {
+	Data       []interface{} `json:"data"` // Query result data
+	SkipToken  *string       `json:"skipToken,omitempty"`
+	Quota      int           // Value of x-ms-user-quota-remaining header as int
+	RetryAfter time.Duration // Value of x-ms-user-quota-resets-after header as timespan
+}
+
+// NewGraphQuery creates a new GraphQuery using the provided TokenCredential.
+func NewGraphQuery(cred azcore.TokenCredential) *GraphQueryClient {
+	// Create a new HTTP client with a timeout
+	httpClient := &http.Client{
+		Timeout: 60 * time.Second,
 	}
-)
 
-func NewGraphQuery(cred azcore.TokenCredential) *GraphQuery {
-	client, err := arg.NewClient(cred, nil)
+	// Acquire an access token using the provided credential
+	token, err := cred.GetToken(context.Background(), policy.TokenRequestOptions{
+		Scopes: []string{"https://management.azure.com/.default"},
+	})
 	if err != nil {
-		log.Fatal().Err(err).Msg("Failed to create Resource Graph client")
-		return nil
+		log.Fatal().Err(err).Msg("Failed to acquire Azure access token")
 	}
-	return &GraphQuery{
-		client: client,
+
+	// Set the access token string
+	accessToken := token.Token
+
+	return &GraphQueryClient{
+		httpClient:  httpClient,
+		endpoint:    "https://management.azure.com/providers/Microsoft.ResourceGraph/resources?api-version=2021-03-01",
+		accessToken: accessToken,
 	}
 }
 
-func (q *GraphQuery) Query(ctx context.Context, query string, subscriptions []*string) *GraphResult {
+// Query executes a Resource Graph query for the given subscriptions and query string.
+// It handles batching and pagination.
+func (q *GraphQueryClient) Query(ctx context.Context, query string, subscriptions []*string) *GraphResult {
 	result := GraphResult{
 		Data: make([]interface{}, 0),
 	}
 
+	// Convert []*string to []string for serialization
+	subscriptionIDs := make([]string, len(subscriptions))
+	for i, s := range subscriptions {
+		if s != nil {
+			subscriptionIDs[i] = *s
+		}
+	}
+
 	// Run the query in batches of 300 subscriptions
 	batchSize := 300
-	for i := 0; i < len(subscriptions); i += batchSize {
+	for i := 0; i < len(subscriptionIDs); i += batchSize {
 		j := i + batchSize
-		if j > len(subscriptions) {
-			j = len(subscriptions)
+		if j > len(subscriptionIDs) {
+			j = len(subscriptionIDs)
 		}
 
-		format := arg.ResultFormatObjectArray
-		request := arg.QueryRequest{
-			Subscriptions: subscriptions[i:j],
-			Query:         &query,
-			Options: &arg.QueryRequestOptions{
-				ResultFormat: &format,
-				Top:          to.Ptr(int32(1000)),
-			},
-		}
-
-		if q.client == nil {
-			log.Fatal().Msg("Resource Graph client not initialized")
+		format := "objectArray"
+		options := &QueryRequestOptions{
+			ResultFormat: format,
+			Top:          to.Ptr(int32(1000)),
 		}
 
 		var skipToken *string = nil
 		for ok := true; ok; ok = skipToken != nil {
-			request.Options.SkipToken = skipToken
-			// Run the query and get the results
-			results, err := q.retry(ctx, 3, 10*time.Second, request)
+			options.SkipToken = skipToken
+			request := QueryRequest{
+				Subscriptions: subscriptionIDs[i:j],
+				Query:         query,
+				Options:       options,
+			}
+
+			startTime := time.Now()
+			resp, err := q.retry(ctx, 3, 10*time.Second, request)
+			elapsed := time.Since(startTime).Milliseconds()
 			if err == nil {
-				result.Data = append(result.Data, results.Data.([]interface{})...)
-				skipToken = results.SkipToken
+				result.Data = append(result.Data, resp.Data...)
+				skipToken = resp.SkipToken
+				// If the response contains a skip token, it means there are more results to fetch
+				// so if the request took less than 333ms, we wait to avoid throttling
+				if skipToken != nil && elapsed < 333 {
+					// If the query took less than 333ms, wait to avoid throttling
+					log.Debug().Msgf("Graph query took %d ms, waiting to avoid throttling", elapsed)
+					time.Sleep(time.Duration(400-elapsed) * time.Millisecond)
+				}
+				// Quota limit reached, sleep for the duration specified in the response header
+				if resp.Quota == 0 {
+					duration := resp.RetryAfter
+					log.Debug().Msgf("Graph query quota limit reached. Sleeping for %s", duration)
+					time.Sleep(duration)
+				}
 			} else {
 				log.Fatal().Err(err).Msgf("Failed to run Resource Graph query: %s", query)
 				return nil
@@ -78,18 +144,16 @@ func (q *GraphQuery) Query(ctx context.Context, query string, subscriptions []*s
 	return &result
 }
 
-func (q *GraphQuery) retry(ctx context.Context, attempts int, sleep time.Duration, request arg.QueryRequest) (arg.ClientResourcesResponse, error) {
+// retry executes the Resource Graph query with retries and exponential backoff.
+// Returns the QueryResponse or error.
+func (q *GraphQueryClient) retry(ctx context.Context, attempts int, sleep time.Duration, request QueryRequest) (*QueryResponse, error) {
 	var err error
 	for i := 0; ; i++ {
-		res, err := q.client.Resources(ctx, request, nil)
+		resp, err := q.doRequest(ctx, request)
 		if err == nil {
-			return res, nil
+			return resp, nil
 		}
 
-		// if shouldSkipError(err) {
-		// 	return []azqr.AzureServiceResult{}, nil
-		// }
-
 		errAsString := err.Error()
 
 		if i >= (attempts - 1) {
@@ -102,5 +166,80 @@ func (q *GraphQuery) retry(ctx context.Context, attempts int, sleep time.Duratio
 		time.Sleep(sleep)
 		sleep *= 2
 	}
-	return arg.ClientResourcesResponse{}, err
+	return nil, err
+}
+
+// doRequest sends the HTTP request to the Resource Graph API and returns the response.
+func (q *GraphQueryClient) doRequest(ctx context.Context, request QueryRequest) (*QueryResponse, error) {
+	// Serialize request to JSON
+	body, err := json.Marshal(request)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request: %w", err)
+	}
+
+	// Create HTTP request
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, q.endpoint, bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create HTTP request: %w", err)
+	}
+
+	// Set headers
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+q.accessToken)
+
+	// Send request
+	resp, err := q.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to send HTTP request: %w", err)
+	}
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			log.Fatal().Err(err).Msg("Failed to close response body")
+		}
+	}()
+
+	// Read response body
+	respBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
+	}
+
+	// Check for non-200 status codes
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return nil, fmt.Errorf("received non-2xx status code: %d, body: %s", resp.StatusCode, string(respBody))
+	}
+
+	// Parse response JSON
+	var queryResp QueryResponse
+	if err := json.Unmarshal(respBody, &queryResp); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+
+	// Extract quota headers and set them in the QueryResponse struct
+
+	// Parse x-ms-user-quota-remaining as int
+	quotaStr := resp.Header.Get("x-ms-user-quota-remaining")
+	if quotaStr != "" {
+		quota, err := strconv.Atoi(quotaStr)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse quota header: %w", err)
+		}
+		queryResp.Quota = quota
+	}
+
+	// Parse x-ms-user-quota-resets-after as timespan in "hh:mm:ss" format
+	retryAfterStr := resp.Header.Get("x-ms-user-quota-resets-after")
+	if retryAfterStr != "" {
+		// If time.ParseDuration fails, fallback to manual parsing
+		var h, m, s int
+		_, scanErr := fmt.Sscanf(retryAfterStr, "%d:%d:%d", &h, &m, &s)
+		if scanErr != nil {
+			return nil, fmt.Errorf("failed to parse retry-after header: %w", scanErr)
+		}
+		queryResp.RetryAfter = time.Duration(h)*time.Hour + time.Duration(m)*time.Minute + time.Duration(s)*time.Second
+	}
+
+	log.Debug().Msgf("Graph query quota remaining: %d, Retry after: %s", queryResp.Quota, queryResp.RetryAfter)
+
+	return &queryResp, nil
 }