TRAC-469: add support for wae metrics

Author: rtalvarez

main.go

@@ -24,6 +24,7 @@ var (
 	cfclient  *cf.Client
 	cftimeout time.Duration
 	gql       *GraphQL
+	waeClient *WAEClient
 	log       = logrus.New()
 
 	// kvTrackedNamespaces is the set of KV namespace IDs that get their own
@@ -122,6 +123,7 @@ func fetchMetrics(deniedMetricsSet MetricsSet) {
 		go fetchLoadblancerPoolsHealth(a, &wg)
 		go fetchZeroTrustAnalyticsForAccount(a, &wg)
 		go fetchDNSFirewallAnalytics(a, &wg, deniedMetricsSet)
+		go fetchWorkerWAEAnalytics(a, &wg, deniedMetricsSet)
 	}
 
 	zones := fetchZones(accounts)
@@ -283,6 +285,10 @@ func main() {
 	viper.BindEnv("metrics_denylist")
 	viper.SetDefault("metrics_denylist", "")
 
+	flags.String("wae_dataset_prefix", "bc_worker_metrics_", "Prefix used to discover Cloudflare Workers Analytics Engine datasets via SHOW TABLES. The remainder of each dataset name (with underscores converted to dashes) is emitted as the script_name label.")
+	viper.BindEnv("wae_dataset_prefix")
+	viper.SetDefault("wae_dataset_prefix", "bc_worker_metrics_")
+
 	flags.String("log_level", "info", "log level")
 	viper.BindEnv("log_level")
 	viper.SetDefault("log_level", "info")
@@ -330,6 +336,7 @@ func main() {
 			Transport: middlewares,
 		}
 		gql = NewGraphQLClient(gqlHTTPClient)
+		waeClient = NewWAEClient(gqlHTTPClient)
 	} else if len(viper.GetString("cf_api_email")) > 0 && len(viper.GetString("cf_api_key")) > 0 {
 		cfclient = cf.NewClient(
 			cfoption.WithAPIKey(viper.GetString("cf_api_key")),
@@ -343,6 +350,7 @@ func main() {
 			Transport: middlewares,
 		}
 		gql = NewGraphQLClient(gqlHTTPClient)
+		waeClient = NewWAEClient(gqlHTTPClient)
 	} else {
 		log.Fatal("Please provide CF_API_KEY+CF_API_EMAIL or CF_API_TOKEN")
 	}

prometheus.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"strings"
@@ -75,6 +76,7 @@ const (
 	queueOperationBytesMetricName                MetricName = "cloudflare_queue_operations_bytes"
 	queueOperationLagTimeMetricName              MetricName = "cloudflare_queue_operations_lag_time"
 	queueOperationRetryCountMetricName           MetricName = "cloudflare_queue_operations_retry_count"
+	workerOperationDurationMetricName            MetricName = "cloudflare_worker_operation_duration_seconds"
 )
 
 type MetricsSet map[MetricName]struct{}
@@ -407,6 +409,11 @@ var (
 		Help: "DNS Firewall query count by query type and response code",
 	}, []string{"account_id", "account_name", "dns_firewall_id", "query_type", "response_code"},
 	)
+
+	workerOperationDuration = prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Name: workerOperationDurationMetricName.String(),
+		Help: "Average duration of a worker operation in seconds, sliced by metric_type (worker-defined event type) and label (worker-defined slice)",
+	}, []string{"script_name", "account", "metric_type", "label"})
 )
 
 func buildAllMetricsSet() MetricsSet {
@@ -465,6 +472,7 @@ func buildAllMetricsSet() MetricsSet {
 	allMetricsSet.Add(queueOperationBytesMetricName)
 	allMetricsSet.Add(queueOperationLagTimeMetricName)
 	allMetricsSet.Add(queueOperationRetryCountMetricName)
+	allMetricsSet.Add(workerOperationDurationMetricName)
 	return allMetricsSet
 }
 
@@ -647,6 +655,9 @@ func mustRegisterMetrics(deniedMetrics MetricsSet) {
 	if !deniedMetrics.Has(queueOperationRetryCountMetricName) {
 		prometheus.MustRegister(queueOperationRetryCount)
 	}
+	if !deniedMetrics.Has(workerOperationDurationMetricName) {
+		prometheus.MustRegister(workerOperationDuration)
+	}
 }
 
 func fetchLoadblancerPoolsHealth(account cfaccounts.Account, wg *sync.WaitGroup) {
@@ -1348,6 +1359,89 @@ func addCloudflareTunnelStatus(account cfaccounts.Account) {
 	}
 }
 
+// fetchWorkerWAEAnalytics discovers all WAE datasets in the account that match
+// the configured prefix and queries each one for worker metrics. The dataset
+// name's suffix (with underscores converted to dashes) is emitted as the
+// script_name label so the series are joinable with cloudflare_worker_* metrics.
+func fetchWorkerWAEAnalytics(account cfaccounts.Account, wg *sync.WaitGroup, deniedMetricsSet MetricsSet) {
+	wg.Add(1)
+	defer wg.Done()
+
+	prefix := viper.GetString("wae_dataset_prefix")
+	if prefix == "" {
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), cftimeout)
+	defer cancel()
+
+	listResp, err := waeClient.Query(ctx, account.ID, "SHOW TABLES")
+	if err != nil {
+		log.Error("failed to list WAE datasets for account ", account.ID, ": ", err)
+		return
+	}
+
+	var datasets []string
+	for _, row := range listResp.Data {
+		name, _ := row["dataset"].(string)
+		if strings.HasPrefix(name, prefix) {
+			datasets = append(datasets, name)
+		}
+	}
+	if len(datasets) == 0 {
+		log.Debugf("no WAE datasets matched prefix %q in account %s", prefix, account.ID)
+		return
+	}
+
+	accountName := strings.ToLower(strings.ReplaceAll(account.Name, " ", "-"))
+
+	// Drop this account's previous label combinations so labels that no longer
+	// see traffic (or scripts that no longer exist) disappear instead of going
+	// stale at their last value.
+	workerOperationDuration.DeletePartialMatch(prometheus.Labels{"account": accountName})
+
+	for _, dataset := range datasets {
+		scriptName := strings.ReplaceAll(strings.TrimPrefix(dataset, prefix), "_", "-")
+
+		query := fmt.Sprintf(`
+			SELECT
+				index1 AS metric_type,
+				blob1 AS label,
+				SUM(_sample_interval * double1) / SUM(_sample_interval) AS avg_duration_ms,
+				SUM(_sample_interval) AS total_requests
+			FROM %s
+			WHERE timestamp > NOW() - INTERVAL '5' MINUTE
+			GROUP BY index1, blob1
+			FORMAT JSON`, dataset)
+
+		resp, err := waeClient.Query(ctx, account.ID, query)
+		if err != nil {
+			log.Errorf("failed to fetch WAE analytics from %s (account %s): %v", dataset, account.ID, err)
+			continue
+		}
+
+		if deniedMetricsSet.Has(workerOperationDurationMetricName) {
+			continue
+		}
+
+		for _, row := range resp.Data {
+			metricType, _ := row["metric_type"].(string)
+			label, _ := row["label"].(string)
+			avgDurationMs, _ := row["avg_duration_ms"].(float64)
+			if metricType == "" {
+				continue
+			}
+
+			workerOperationDuration.With(prometheus.Labels{
+				"script_name": scriptName,
+				"account":     accountName,
+				"metric_type": metricType,
+				"label":       label,
+			}).Set(avgDurationMs / 1000.0)
+		}
+	}
+}
+
 // The status of the tunnel.
 // Valid values are:
 //   - inactive (tunnel has never been run)

wae.go

@@ -0,0 +1,71 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+)
+
+const waeAPIBase = "https://api.cloudflare.com/client/v4/accounts"
+
+// WAEClient queries the Workers Analytics Engine SQL API.
+type WAEClient struct {
+	httpClient *http.Client
+}
+
+// WAEResponse is the JSON response from the WAE SQL API.
+// This is NOT the standard Cloudflare v4 API envelope.
+type WAEResponse struct {
+	Meta []WAEColumnMeta          `json:"meta"`
+	Data []map[string]interface{} `json:"data"`
+	Rows int                      `json:"rows"`
+}
+
+type WAEColumnMeta struct {
+	Name string `json:"name"`
+	Type string `json:"type"`
+}
+
+func NewWAEClient(httpClient *http.Client) *WAEClient {
+	if httpClient == nil {
+		httpClient = http.DefaultClient
+	}
+	return &WAEClient{httpClient: httpClient}
+}
+
+// Query executes a SQL query against the WAE API for the given account.
+func (w *WAEClient) Query(ctx context.Context, accountID, query string) (*WAEResponse, error) {
+	url := fmt.Sprintf("%s/%s/analytics_engine/sql", waeAPIBase, accountID)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(query))
+	if err != nil {
+		return nil, fmt.Errorf("creating WAE request: %w", err)
+	}
+
+	log.Debugf("WAE query (account=%s): %s", accountID, query)
+
+	resp, err := w.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("executing WAE query: %w", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("reading WAE response: %w", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("WAE query failed (HTTP %d): %s", resp.StatusCode, body)
+	}
+
+	var result WAEResponse
+	if err := json.Unmarshal(body, &result); err != nil {
+		return nil, fmt.Errorf("parsing WAE response: %w", err)
+	}
+
+	return &result, nil
+}