diff --git a/internal/cmd/trafficcontrol/budget_create.go b/internal/cmd/trafficcontrol/budget_create.go index c4528dd7..15981022 100644 --- a/internal/cmd/trafficcontrol/budget_create.go +++ b/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 diff --git a/internal/cmd/trafficcontrol/budget_create_test.go b/internal/cmd/trafficcontrol/budget_create_test.go index d57ceba9..03f485d4 100644 --- a/internal/cmd/trafficcontrol/budget_create_test.go +++ b/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() diff --git a/internal/cmd/trafficcontrol/budget_update.go b/internal/cmd/trafficcontrol/budget_update.go index d68a1f0f..1e1622fb 100644 --- a/internal/cmd/trafficcontrol/budget_update.go +++ b/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 } diff --git a/internal/cmd/trafficcontrol/budget_update_test.go b/internal/cmd/trafficcontrol/budget_update_test.go index b3f0ad3b..4e33bab4 100644 --- a/internal/cmd/trafficcontrol/budget_update_test.go +++ b/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 }, diff --git a/internal/cmd/trafficcontrol/traffic_control.go b/internal/cmd/trafficcontrol/traffic_control.go index 90e4d883..f80e4365 100644 --- a/internal/cmd/trafficcontrol/traffic_control.go +++ b/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, } }