Merge pull request #1222 from planetscale/tc-budget-show

Support showing Traffic Control budgets

Author: joshmgross

go.mod

@@ -25,7 +25,7 @@ require (
 	github.com/mattn/go-shellwords v1.0.12
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
-	github.com/planetscale/planetscale-go v0.155.0
+	github.com/planetscale/planetscale-go v0.156.0
 	github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4
 	github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7
 	github.com/spf13/cobra v1.10.2

go.sum

@@ -184,6 +184,8 @@ github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2v
 github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q=
 github.com/planetscale/planetscale-go v0.155.0 h1:KYFRWFn9d5BeZc++4DF0wS+mlRQ4efrAy+6Zw/1kzXs=
 github.com/planetscale/planetscale-go v0.155.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
+github.com/planetscale/planetscale-go v0.156.0 h1:VThRagwFmthnPZ/A0W8N+bHHQ8+PdqdcLBHunr7p/jY=
+github.com/planetscale/planetscale-go v0.156.0/go.mod h1:paQCI5SgquuoewvMQM7R+r1XJO868bdP6/ihGidYRM0=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4/go.mod h1:M52h5IWxAcbdQ1hSZrLAGQC4ZXslxEsK/Wh9nu3wdWs=
 github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 h1:aRd6vdE1fyuSI4RVj7oCr8lFmgqXvpnPUmN85VbZCp8=

internal/cmd/root.go

@@ -29,6 +29,7 @@ import (
 	"github.com/planetscale/cli/internal/cmd/mcp"
 	"github.com/planetscale/cli/internal/cmd/role"
 	"github.com/planetscale/cli/internal/cmd/size"
+	"github.com/planetscale/cli/internal/cmd/trafficcontrol"
 	"github.com/planetscale/cli/internal/cmd/workflow"
 
 	"github.com/fatih/color"
@@ -320,6 +321,10 @@ func runCmd(ctx context.Context, ver, commit, buildDate string, format *printer.
 	roleCmd.GroupID = "postgres"
 	rootCmd.AddCommand(roleCmd)
 
+	trafficCmd := trafficcontrol.TrafficCmd(ch)
+	trafficCmd.GroupID = "postgres"
+	rootCmd.AddCommand(trafficCmd)
+
 	annotateRequiredFlags(rootCmd)
 
 	return rootCmd.ExecuteContext(ctx)

internal/cmd/trafficcontrol/budget_show.go

@@ -0,0 +1,53 @@
+package trafficcontrol
+
+import (
+	"fmt"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/printer"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+	"github.com/spf13/cobra"
+)
+
+func BudgetShowCmd(ch *cmdutil.Helper) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "show <database> <branch> <budget-id>",
+		Short: "Show a traffic budget",
+		Args:  cmdutil.RequiredArgs("database", "branch", "budget-id"),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			database := args[0]
+			branch := args[1]
+			budgetID := args[2]
+
+			client, err := ch.Client()
+			if err != nil {
+				return err
+			}
+
+			end := ch.Printer.PrintProgress(fmt.Sprintf("Fetching traffic budget %s", printer.BoldBlue(budgetID)))
+			defer end()
+
+			budget, err := client.TrafficBudgets.Get(ctx, &ps.GetTrafficBudgetRequest{
+				Organization: ch.Config.Organization,
+				Database:     database,
+				Branch:       branch,
+				BudgetID:     budgetID,
+			})
+			if err != nil {
+				switch cmdutil.ErrCode(err) {
+				case ps.ErrNotFound:
+					return fmt.Errorf("traffic budget %s does not exist in %s/%s (organization: %s)",
+						printer.BoldBlue(budgetID), printer.BoldBlue(database), printer.BoldBlue(branch), printer.BoldBlue(ch.Config.Organization))
+				default:
+					return cmdutil.HandleError(err)
+				}
+			}
+
+			end()
+			return ch.Printer.PrintResource(toTrafficBudget(budget))
+		},
+	}
+
+	return cmd
+}

internal/cmd/trafficcontrol/budget_show_test.go

@@ -0,0 +1,104 @@
+package trafficcontrol
+
+import (
+	"bytes"
+	"context"
+	"testing"
+	"time"
+
+	qt "github.com/frankban/quicktest"
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/config"
+	"github.com/planetscale/cli/internal/mock"
+	"github.com/planetscale/cli/internal/printer"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+)
+
+func TestBudgetShowCmd(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "planetscale"
+	db := "database"
+	branch := "main"
+	budgetID := "qok87ki4xlau"
+	cap := 80
+
+	budget := &ps.TrafficBudget{
+		ID:        budgetID,
+		Name:      "CPU Limiter",
+		Mode:      "enforce",
+		Capacity:  &cap,
+		Rules:     []ps.TrafficRule{{ID: "rule-1", Kind: "fingerprint"}},
+		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{
+		GetFn: func(ctx context.Context, req *ps.GetTrafficBudgetRequest) (*ps.TrafficBudget, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Branch, qt.Equals, branch)
+			c.Assert(req.BudgetID, qt.Equals, budgetID)
+			return budget, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{TrafficBudgets: svc}, nil
+		},
+	}
+
+	cmd := BudgetShowCmd(ch)
+	cmd.SetArgs([]string{db, branch, budgetID})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.GetFnInvoked, qt.IsTrue)
+
+	res := &TrafficBudget{orig: budget}
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}
+
+func TestBudgetShowCmd_NotFound(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "planetscale"
+	db := "database"
+	branch := "main"
+	budgetID := "qok87ki4xlau"
+
+	svc := &mock.TrafficBudgetsService{
+		GetFn: func(ctx context.Context, req *ps.GetTrafficBudgetRequest) (*ps.TrafficBudget, error) {
+			return nil, &ps.Error{Code: ps.ErrNotFound}
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{TrafficBudgets: svc}, nil
+		},
+	}
+
+	cmd := BudgetShowCmd(ch)
+	cmd.SetArgs([]string{db, branch, budgetID})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err.Error(), qt.Contains, "does not exist")
+	c.Assert(svc.GetFnInvoked, qt.IsTrue)
+}

internal/cmd/trafficcontrol/traffic_control.go

@@ -0,0 +1,79 @@
+package trafficcontrol
+
+import (
+	"encoding/json"
+	"strconv"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/printer"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+	"github.com/spf13/cobra"
+)
+
+func TrafficCmd(ch *cmdutil.Helper) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "traffic-control <command>",
+		Short: "Manage Database Traffic Control™ for a Postgres database branch",
+		Long: "Manage Database Traffic Control™ budgets and rules for a Postgres database branch.\n\n" +
+			"This command is only supported for Postgres databases.",
+		PersistentPreRunE: cmdutil.CheckAuthentication(ch.Config),
+	}
+
+	cmd.PersistentFlags().StringVar(&ch.Config.Organization, "org", ch.Config.Organization, "The organization for the current user")
+	cmd.MarkPersistentFlagRequired("org") // nolint:errcheck
+
+	budgetCmd := &cobra.Command{
+		Use:   "budget <command>",
+		Short: "Manage traffic budgets",
+	}
+	budgetCmd.AddCommand(
+		BudgetShowCmd(ch),
+	)
+
+	cmd.AddCommand(budgetCmd)
+	return cmd
+}
+
+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"`
+
+	orig *ps.TrafficBudget
+}
+
+func (b *TrafficBudget) MarshalJSON() ([]byte, error) {
+	return json.MarshalIndent(b.orig, "", "  ")
+}
+
+func (b *TrafficBudget) MarshalCSVValue() any {
+	return []*TrafficBudget{b}
+}
+
+func formatOptionalInt(v *int) string {
+	if v == nil {
+		return "-"
+	}
+	return strconv.Itoa(*v)
+}
+
+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,
+	}
+}

internal/mock/traffic_budget.go

@@ -0,0 +1,45 @@
+package mock
+
+import (
+	"context"
+
+	ps "github.com/planetscale/planetscale-go/planetscale"
+)
+
+type TrafficBudgetsService struct {
+	ListFn          func(context.Context, *ps.ListTrafficBudgetsRequest) ([]*ps.TrafficBudget, error)
+	ListFnInvoked   bool
+	GetFn           func(context.Context, *ps.GetTrafficBudgetRequest) (*ps.TrafficBudget, error)
+	GetFnInvoked    bool
+	CreateFn        func(context.Context, *ps.CreateTrafficBudgetRequest) (*ps.TrafficBudget, error)
+	CreateFnInvoked bool
+	UpdateFn        func(context.Context, *ps.UpdateTrafficBudgetRequest) (*ps.TrafficBudget, error)
+	UpdateFnInvoked bool
+	DeleteFn        func(context.Context, *ps.DeleteTrafficBudgetRequest) error
+	DeleteFnInvoked bool
+}
+
+func (s *TrafficBudgetsService) List(ctx context.Context, req *ps.ListTrafficBudgetsRequest) ([]*ps.TrafficBudget, error) {
+	s.ListFnInvoked = true
+	return s.ListFn(ctx, req)
+}
+
+func (s *TrafficBudgetsService) Get(ctx context.Context, req *ps.GetTrafficBudgetRequest) (*ps.TrafficBudget, error) {
+	s.GetFnInvoked = true
+	return s.GetFn(ctx, req)
+}
+
+func (s *TrafficBudgetsService) Create(ctx context.Context, req *ps.CreateTrafficBudgetRequest) (*ps.TrafficBudget, error) {
+	s.CreateFnInvoked = true
+	return s.CreateFn(ctx, req)
+}
+
+func (s *TrafficBudgetsService) Update(ctx context.Context, req *ps.UpdateTrafficBudgetRequest) (*ps.TrafficBudget, error) {
+	s.UpdateFnInvoked = true
+	return s.UpdateFn(ctx, req)
+}
+
+func (s *TrafficBudgetsService) Delete(ctx context.Context, req *ps.DeleteTrafficBudgetRequest) error {
+	s.DeleteFnInvoked = true
+	return s.DeleteFn(ctx, req)
+}