Merge pull request #1233 from planetscale/tc-warning-threshold

Support configuring traffic budget `warning_threshold`

Author: joshmgross

internal/cmd/trafficcontrol/budget_create.go

@@ -11,12 +11,13 @@ import (
 
 func BudgetCreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	var flags struct {
-		name        string
-		mode        string
-		capacity    int
-		rate        int
-		burst       int
-		concurrency int
+		name             string
+		mode             string
+		capacity         int
+		rate             int
+		burst            int
+		concurrency      int
+		warningThreshold int
 	}
 
 	cmd := &cobra.Command{
@@ -57,6 +58,9 @@ func BudgetCreateCmd(ch *cmdutil.Helper) *cobra.Command {
 			if cmd.Flags().Changed("concurrency") {
 				req.Concurrency = &flags.concurrency
 			}
+			if cmd.Flags().Changed("warning-threshold") {
+				req.WarningThreshold = &flags.warningThreshold
+			}
 
 			budget, err := client.TrafficBudgets.Create(ctx, req)
 			if err != nil {
@@ -86,6 +90,7 @@ func BudgetCreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.Flags().IntVar(&flags.rate, "rate", 0, "Rate at which capacity refills, as a percentage of server resources (0-100). Unlimited when not set.")
 	cmd.Flags().IntVar(&flags.burst, "burst", 0, "Maximum capacity a single query can consume (0-6000). Unlimited when not set.")
 	cmd.Flags().IntVar(&flags.concurrency, "concurrency", 0, "Percentage of available worker processes (0-100). Unlimited when not set.")
+	cmd.Flags().IntVar(&flags.warningThreshold, "warning-threshold", 0, "Percentage (0-100) of capacity, burst, or concurrency at which to emit warnings for enforced budgets.")
 
 	cmd.MarkFlagRequired("name") // nolint:errcheck
 

internal/cmd/trafficcontrol/budget_create_test.go

@@ -22,18 +22,19 @@ func TestBudgetCreateCmd(t *testing.T) {
 	p := printer.NewPrinter(&format)
 	p.SetResourceOutput(&buf)
 
-	cap, rate, burst, conc := 80, 50, 60, 40
+	cap, rate, burst, conc, warnTh := 80, 50, 60, 40, 30
 
 	created := &ps.TrafficBudget{
-		ID:          budgetID,
-		Name:        "CPU Limiter",
-		Mode:        "enforce",
-		Capacity:    &cap,
-		Rate:        &rate,
-		Burst:       &burst,
-		Concurrency: &conc,
-		CreatedAt:   time.Date(2025, 6, 15, 10, 0, 0, 0, time.UTC),
-		UpdatedAt:   time.Date(2025, 6, 15, 10, 0, 0, 0, time.UTC),
+		ID:               budgetID,
+		Name:             "CPU Limiter",
+		Mode:             "enforce",
+		Capacity:         &cap,
+		Rate:             &rate,
+		Burst:            &burst,
+		Concurrency:      &conc,
+		WarningThreshold: &warnTh,
+		CreatedAt:        time.Date(2025, 6, 15, 10, 0, 0, 0, time.UTC),
+		UpdatedAt:        time.Date(2025, 6, 15, 10, 0, 0, 0, time.UTC),
 	}
 
 	svc := &mock.TrafficBudgetsService{
@@ -47,6 +48,7 @@ func TestBudgetCreateCmd(t *testing.T) {
 			c.Assert(*req.Rate, qt.Equals, 50)
 			c.Assert(*req.Burst, qt.Equals, 60)
 			c.Assert(*req.Concurrency, qt.Equals, 40)
+			c.Assert(*req.WarningThreshold, qt.Equals, 30)
 			return created, nil
 		},
 	}
@@ -67,6 +69,7 @@ func TestBudgetCreateCmd(t *testing.T) {
 		"--rate", "50",
 		"--burst", "60",
 		"--concurrency", "40",
+		"--warning-threshold", "30",
 	})
 	err := cmd.Execute()
 

internal/cmd/trafficcontrol/budget_update.go

@@ -11,12 +11,13 @@ import (
 
 func BudgetUpdateCmd(ch *cmdutil.Helper) *cobra.Command {
 	var flags struct {
-		name        string
-		mode        string
-		capacity    int
-		rate        int
-		burst       int
-		concurrency int
+		name             string
+		mode             string
+		capacity         int
+		rate             int
+		burst            int
+		concurrency      int
+		warningThreshold int
 	}
 
 	cmd := &cobra.Command{
@@ -59,6 +60,9 @@ func BudgetUpdateCmd(ch *cmdutil.Helper) *cobra.Command {
 			if cmd.Flags().Changed("concurrency") {
 				req.Concurrency = &flags.concurrency
 			}
+			if cmd.Flags().Changed("warning-threshold") {
+				req.WarningThreshold = &flags.warningThreshold
+			}
 
 			end := ch.Printer.PrintProgress(fmt.Sprintf("Updating traffic budget %s in %s/%s",
 				printer.BoldBlue(budgetID), printer.BoldBlue(database), printer.BoldBlue(branch)))
@@ -91,6 +95,7 @@ func BudgetUpdateCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.Flags().IntVar(&flags.rate, "rate", 0, "Rate at which capacity refills, as a percentage of server resources (0-100). Unlimited when not set.")
 	cmd.Flags().IntVar(&flags.burst, "burst", 0, "Maximum capacity a single query can consume (0-6000). Unlimited when not set.")
 	cmd.Flags().IntVar(&flags.concurrency, "concurrency", 0, "Percentage of available worker processes (0-100). Unlimited when not set.")
+	cmd.Flags().IntVar(&flags.warningThreshold, "warning-threshold", 0, "Percentage (0-100) of capacity, burst, or concurrency at which to emit warnings for enforced budgets.")
 
 	return cmd
 }

internal/cmd/trafficcontrol/budget_update_test.go

@@ -45,6 +45,7 @@ func TestBudgetUpdateCmd(t *testing.T) {
 			c.Assert(req.Rate, qt.IsNil)
 			c.Assert(req.Burst, qt.IsNil)
 			c.Assert(req.Concurrency, qt.IsNil)
+			c.Assert(req.WarningThreshold, qt.IsNil)
 			c.Assert(req.Rules, qt.IsNil)
 			return updated, nil
 		},

internal/cmd/trafficcontrol/traffic_control.go

@@ -49,15 +49,16 @@ func TrafficCmd(ch *cmdutil.Helper) *cobra.Command {
 }
 
 type TrafficBudget struct {
-	ID          string `header:"id" json:"id"`
-	Name        string `header:"name" json:"name"`
-	Mode        string `header:"mode" json:"mode"`
-	Capacity    string `header:"capacity" json:"capacity"`
-	Rate        string `header:"rate" json:"rate"`
-	Burst       string `header:"burst" json:"burst"`
-	Concurrency string `header:"concurrency" json:"concurrency"`
-	CreatedAt   int64  `header:"created_at,timestamp(ms|utc|human)" json:"created_at"`
-	UpdatedAt   int64  `header:"updated_at,timestamp(ms|utc|human)" json:"updated_at"`
+	ID               string `header:"id" json:"id"`
+	Name             string `header:"name" json:"name"`
+	Mode             string `header:"mode" json:"mode"`
+	Capacity         string `header:"capacity" json:"capacity"`
+	Rate             string `header:"rate" json:"rate"`
+	Burst            string `header:"burst" json:"burst"`
+	Concurrency      string `header:"concurrency" json:"concurrency"`
+	WarningThreshold string `header:"warning_threshold" json:"warning_threshold"`
+	CreatedAt        int64  `header:"created_at,timestamp(ms|utc|human)" json:"created_at"`
+	UpdatedAt        int64  `header:"updated_at,timestamp(ms|utc|human)" json:"updated_at"`
 
 	orig *ps.TrafficBudget
 }
@@ -72,16 +73,17 @@ func (b *TrafficBudget) MarshalCSVValue() any {
 
 func toTrafficBudget(b *ps.TrafficBudget) *TrafficBudget {
 	return &TrafficBudget{
-		ID:          b.ID,
-		Name:        b.Name,
-		Mode:        b.Mode,
-		Capacity:    formatOptionalInt(b.Capacity),
-		Rate:        formatOptionalInt(b.Rate),
-		Burst:       formatOptionalInt(b.Burst),
-		Concurrency: formatOptionalInt(b.Concurrency),
-		CreatedAt:   printer.GetMilliseconds(b.CreatedAt),
-		UpdatedAt:   printer.GetMilliseconds(b.UpdatedAt),
-		orig:        b,
+		ID:               b.ID,
+		Name:             b.Name,
+		Mode:             b.Mode,
+		Capacity:         formatOptionalInt(b.Capacity),
+		Rate:             formatOptionalInt(b.Rate),
+		Burst:            formatOptionalInt(b.Burst),
+		Concurrency:      formatOptionalInt(b.Concurrency),
+		WarningThreshold: formatOptionalInt(b.WarningThreshold),
+		CreatedAt:        printer.GetMilliseconds(b.CreatedAt),
+		UpdatedAt:        printer.GetMilliseconds(b.UpdatedAt),
+		orig:             b,
 	}
 }