From c1d531234dacbb52a604f90fa1e9799b1a31e222 Mon Sep 17 00:00:00 2001 From: Elom Gomez Date: Mon, 17 Aug 2026 16:01:36 -0500 Subject: [PATCH] Add pscale database throttler show/update for Vitess Expose database-level migration throttler defaults (distinct from per-DR and vtctld throttlers) and reject non-MySQL databases in the CLI. Co-authored-by: Cursor --- AGENTS.md | 13 ++ internal/cmd/database/database.go | 1 + internal/cmd/database/throttler.go | 239 ++++++++++++++++++++++ internal/cmd/database/throttler_test.go | 252 ++++++++++++++++++++++++ internal/mock/database.go | 16 ++ internal/planetscale/databases.go | 55 ++++++ internal/planetscale/databases_test.go | 61 ++++++ 7 files changed, 637 insertions(+) create mode 100644 internal/cmd/database/throttler.go create mode 100644 internal/cmd/database/throttler_test.go diff --git a/AGENTS.md b/AGENTS.md index c402a9f5..3ae35861 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -247,6 +247,19 @@ Caveats: - `outliers`/`calls` need `pg_stat_statements` on PostgreSQL; if missing, use `pscale insights queries` instead (no extension needed). - Rule of thumb: start with `insights` (traffic-aware, historical), use `inspect` for live state (locks, in-flight queries) and physical layout (sizes, bloat, index usage). +## Vitess database throttler + +Database-level default for future deploy request migrations (not per-DR, not tablet/vtctld): + +```bash +pscale database throttler show --org --format json +pscale database throttler update --org --format json --ratio 25 +pscale database throttler update --org --format json \ + --configuration main=10 --configuration sharded=40 +``` + +`--ratio` is 0–95 (0 disables throttling; 95 is slowest). Use either `--ratio` or `--configuration keyspace=ratio`, not both. Vitess only. + ## Vitess deploy requests (inspect + throttler) Core lifecycle is already covered (`list/create/show/diff/review/deploy/apply/edit/cancel/close/revert/skip-revert`). These inspect commands are read-only: diff --git a/internal/cmd/database/database.go b/internal/cmd/database/database.go index 821a6592..4f53ac12 100644 --- a/internal/cmd/database/database.go +++ b/internal/cmd/database/database.go @@ -30,6 +30,7 @@ func DatabaseCmd(ch *cmdutil.Helper) *cobra.Command { cmd.AddCommand(DeleteCmd(ch)) cmd.AddCommand(ShowCmd(ch)) cmd.AddCommand(UpdateCmd(ch)) + cmd.AddCommand(ThrottlerCmd(ch)) cmd.AddCommand(IPRestrictionCmd(ch)) cmd.AddCommand(DumpCmd(ch)) cmd.AddCommand(RestoreCmd(ch)) diff --git a/internal/cmd/database/throttler.go b/internal/cmd/database/throttler.go new file mode 100644 index 00000000..e5aa1977 --- /dev/null +++ b/internal/cmd/database/throttler.go @@ -0,0 +1,239 @@ +package database + +import ( + "context" + "encoding/json" + "fmt" + "strconv" + "strings" + + "github.com/planetscale/cli/internal/cmdutil" + "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" + + "github.com/spf13/cobra" +) + +// ThrottlerCmd groups database-level Vitess migration throttler commands. +func ThrottlerCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "throttler ", + Short: "Show or update database throttler configuration", + Long: `Show or update database-level Vitess migration throttler configuration. + +This sets the default throttler for future deploy requests on the database. +It is not the per-deploy-request throttler (pscale deploy-request throttler) +and not the tablet/vtctld throttler (pscale branch vtctld throttler).`, + } + + cmd.AddCommand(ThrottlerShowCmd(ch)) + cmd.AddCommand(ThrottlerUpdateCmd(ch)) + return cmd +} + +// ThrottlerShowCmd shows database-level throttler configuration. +func ThrottlerShowCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "show ", + Short: "Show database throttler configuration", + Args: cmdutil.RequiredArgs("database"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database := args[0] + + client, err := ch.Client() + if err != nil { + return err + } + + if err := requireVitessDatabase(ctx, ch, client, database); err != nil { + return err + } + + end := ch.Printer.PrintProgress(fmt.Sprintf("Fetching throttler config for database %s", + printer.BoldBlue(database))) + defer end() + + throttler, err := client.Databases.GetThrottler(ctx, &planetscale.GetDatabaseThrottlerRequest{ + Organization: ch.Config.Organization, + Database: database, + }) + if err != nil { + switch cmdutil.ErrCode(err) { + case planetscale.ErrNotFound: + return fmt.Errorf("database %s does not exist in organization %s", + printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization)) + default: + return cmdutil.HandleError(err) + } + } + end() + + return ch.Printer.PrintResource(toThrottler(throttler)) + }, + } + + return cmd +} + +// ThrottlerUpdateCmd updates database-level throttler configuration. +func ThrottlerUpdateCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + ratio int + configurations []string + } + + cmd := &cobra.Command{ + Use: "update ", + Short: "Update database throttler configuration", + Long: `Update database-level Vitess migration throttler configuration. + +Use --ratio to apply one ratio (0-95) to all eligible keyspaces. +Use --configuration keyspace=ratio (repeatable) for per-keyspace ratios. +Pass exactly one of these modes. 0 effectively disables throttling; 95 slows migrations the most. + +This is the database default for future deploy requests, not a single deploy request.`, + Args: cmdutil.RequiredArgs("database"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database := args[0] + + ratioSet := cmd.Flags().Changed("ratio") + if !ratioSet && len(flags.configurations) == 0 { + return fmt.Errorf("must specify --ratio or --configuration") + } + if ratioSet && len(flags.configurations) > 0 { + return fmt.Errorf("cannot use both --ratio and --configuration; pick one mode") + } + + updateReq := &planetscale.UpdateDatabaseThrottlerRequest{ + Organization: ch.Config.Organization, + Database: database, + } + + if ratioSet { + if flags.ratio < 0 || flags.ratio > 95 { + return fmt.Errorf("--ratio must be between 0 and 95, got %d", flags.ratio) + } + updateReq.Ratio = &flags.ratio + } + + if len(flags.configurations) > 0 { + configs, err := parseThrottlerConfigurations(flags.configurations) + if err != nil { + return err + } + updateReq.Configurations = configs + } + + client, err := ch.Client() + if err != nil { + return err + } + + if err := requireVitessDatabase(ctx, ch, client, database); err != nil { + return err + } + + end := ch.Printer.PrintProgress(fmt.Sprintf("Updating throttler config for database %s", + printer.BoldBlue(database))) + defer end() + + throttler, err := client.Databases.UpdateThrottler(ctx, updateReq) + if err != nil { + switch cmdutil.ErrCode(err) { + case planetscale.ErrNotFound: + return fmt.Errorf("database %s does not exist in organization %s", + printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization)) + default: + return cmdutil.HandleError(err) + } + } + end() + + if ch.Printer.Format() == printer.Human { + ch.Printer.Printf("Updated throttler configuration for database %s.\n", + printer.BoldBlue(database)) + return nil + } + + return ch.Printer.PrintResource(toThrottler(throttler)) + }, + } + + cmd.Flags().IntVar(&flags.ratio, "ratio", 0, "Throttler ratio 0-95 applied to all eligible keyspaces") + cmd.Flags().StringArrayVar(&flags.configurations, "configuration", nil, "Per-keyspace ratio as keyspace=ratio (repeatable)") + + return cmd +} + +func requireVitessDatabase(ctx context.Context, ch *cmdutil.Helper, client *planetscale.Client, database string) error { + db, err := client.Databases.Get(ctx, &planetscale.GetDatabaseRequest{ + Organization: ch.Config.Organization, + Database: database, + }) + if err != nil { + switch cmdutil.ErrCode(err) { + case planetscale.ErrNotFound: + return fmt.Errorf("database %s does not exist in organization %s", + printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization)) + default: + return cmdutil.HandleError(err) + } + } + if db.Kind != planetscale.DatabaseEngineMySQL { + return fmt.Errorf("database throttler is only available for Vitess (MySQL) databases; %s is %s", + printer.BoldBlue(database), printer.BoldBlue(string(db.Kind))) + } + return nil +} + +func parseThrottlerConfigurations(values []string) ([]*planetscale.UpdateThrottlerConfiguration, error) { + configs := make([]*planetscale.UpdateThrottlerConfiguration, 0, len(values)) + for _, value := range values { + keyspace, ratioStr, ok := strings.Cut(value, "=") + if !ok || keyspace == "" || ratioStr == "" { + return nil, fmt.Errorf("invalid --configuration %q: expected keyspace=ratio", value) + } + ratio, err := strconv.Atoi(ratioStr) + if err != nil { + return nil, fmt.Errorf("invalid --configuration %q: ratio must be an integer", value) + } + if ratio < 0 || ratio > 95 { + return nil, fmt.Errorf("invalid --configuration %q: ratio must be between 0 and 95", value) + } + configs = append(configs, &planetscale.UpdateThrottlerConfiguration{ + KeyspaceName: keyspace, + Ratio: ratio, + }) + } + return configs, nil +} + +type ThrottlerRow struct { + Keyspaces string `header:"keyspaces" json:"keyspaces"` + Configurations string `header:"configurations" json:"configurations"` + + orig *planetscale.DatabaseThrottler +} + +func (t *ThrottlerRow) MarshalJSON() ([]byte, error) { + return json.MarshalIndent(t.orig, "", " ") +} + +func (t *ThrottlerRow) MarshalCSVValue() interface{} { + return []*ThrottlerRow{t} +} + +func toThrottler(throttler *planetscale.DatabaseThrottler) *ThrottlerRow { + configs := make([]string, 0, len(throttler.Configurations)) + for _, c := range throttler.Configurations { + configs = append(configs, fmt.Sprintf("%s=%g", c.KeyspaceName, c.Ratio)) + } + + return &ThrottlerRow{ + Keyspaces: strings.Join(throttler.Keyspaces, ", "), + Configurations: strings.Join(configs, ", "), + orig: throttler, + } +} diff --git a/internal/cmd/database/throttler_test.go b/internal/cmd/database/throttler_test.go new file mode 100644 index 00000000..ff0663a9 --- /dev/null +++ b/internal/cmd/database/throttler_test.go @@ -0,0 +1,252 @@ +package database + +import ( + "bytes" + "context" + "testing" + + "github.com/planetscale/cli/internal/cmdutil" + "github.com/planetscale/cli/internal/config" + "github.com/planetscale/cli/internal/mock" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" + + qt "github.com/frankban/quicktest" +) + +func TestDatabase_ThrottlerShowCmd(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + org := "planetscale" + db := "planetscale" + throttler := &ps.DatabaseThrottler{ + Keyspaces: []string{"main"}, + Configurations: []*ps.ThrottlerConfiguration{ + {KeyspaceName: "main", Ratio: 50}, + }, + } + + svc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Name: db, Kind: ps.DatabaseEngineMySQL}, nil + }, + GetThrottlerFn: func(ctx context.Context, req *ps.GetDatabaseThrottlerRequest) (*ps.DatabaseThrottler, error) { + c.Assert(req.Organization, qt.Equals, org) + c.Assert(req.Database, qt.Equals, db) + return throttler, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: org}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: svc}, nil + }, + } + + cmd := ThrottlerShowCmd(ch) + cmd.SetArgs([]string{db}) + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(svc.GetThrottlerFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, throttler) +} + +func TestDatabase_ThrottlerUpdateCmd_Ratio(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + org := "planetscale" + db := "planetscale" + throttler := &ps.DatabaseThrottler{ + Keyspaces: []string{"main"}, + Configurations: []*ps.ThrottlerConfiguration{ + {KeyspaceName: "main", Ratio: 25}, + }, + } + + svc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Name: db, Kind: ps.DatabaseEngineMySQL}, nil + }, + UpdateThrottlerFn: func(ctx context.Context, req *ps.UpdateDatabaseThrottlerRequest) (*ps.DatabaseThrottler, error) { + c.Assert(req.Organization, qt.Equals, org) + c.Assert(req.Database, qt.Equals, db) + c.Assert(req.Ratio, qt.IsNotNil) + c.Assert(*req.Ratio, qt.Equals, 25) + c.Assert(req.Configurations, qt.IsNil) + return throttler, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: org}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: svc}, nil + }, + } + + cmd := ThrottlerUpdateCmd(ch) + cmd.SetArgs([]string{db, "--ratio", "25"}) + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(svc.UpdateThrottlerFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, throttler) +} + +func TestDatabase_ThrottlerUpdateCmd_Configurations(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + format := printer.JSON + p := printer.NewPrinter(&format) + p.SetResourceOutput(&buf) + + org := "planetscale" + db := "planetscale" + throttler := &ps.DatabaseThrottler{ + Keyspaces: []string{"main", "sharded"}, + Configurations: []*ps.ThrottlerConfiguration{ + {KeyspaceName: "main", Ratio: 10}, + {KeyspaceName: "sharded", Ratio: 40}, + }, + } + + svc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Name: db, Kind: ps.DatabaseEngineMySQL}, nil + }, + UpdateThrottlerFn: func(ctx context.Context, req *ps.UpdateDatabaseThrottlerRequest) (*ps.DatabaseThrottler, error) { + c.Assert(req.Ratio, qt.IsNil) + c.Assert(req.Configurations, qt.DeepEquals, []*ps.UpdateThrottlerConfiguration{ + {KeyspaceName: "main", Ratio: 10}, + {KeyspaceName: "sharded", Ratio: 40}, + }) + return throttler, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: org}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: svc}, nil + }, + } + + cmd := ThrottlerUpdateCmd(ch) + cmd.SetArgs([]string{db, "--configuration", "main=10", "--configuration", "sharded=40"}) + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(svc.UpdateThrottlerFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, throttler) +} + +func TestDatabase_ThrottlerUpdateCmd_RequiresFlags(t *testing.T) { + c := qt.New(t) + + format := printer.JSON + p := printer.NewPrinter(&format) + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: &mock.DatabaseService{}}, nil + }, + } + + cmd := ThrottlerUpdateCmd(ch) + cmd.SetArgs([]string{"db"}) + err := cmd.Execute() + c.Assert(err, qt.IsNotNil) + c.Assert(err.Error(), qt.Equals, "must specify --ratio or --configuration") +} + +func TestDatabase_ThrottlerUpdateCmd_MutuallyExclusive(t *testing.T) { + c := qt.New(t) + + format := printer.JSON + p := printer.NewPrinter(&format) + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: &mock.DatabaseService{}}, nil + }, + } + + cmd := ThrottlerUpdateCmd(ch) + cmd.SetArgs([]string{"db", "--ratio", "10", "--configuration", "main=10"}) + err := cmd.Execute() + c.Assert(err, qt.IsNotNil) + c.Assert(err.Error(), qt.Equals, "cannot use both --ratio and --configuration; pick one mode") +} + +func TestDatabase_ThrottlerShowCmd_RejectsPostgres(t *testing.T) { + c := qt.New(t) + + format := printer.JSON + p := printer.NewPrinter(&format) + org := "planetscale" + db := "eagle" + + svc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Name: db, Kind: ps.DatabaseEnginePostgres}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: org}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: svc}, nil + }, + } + + cmd := ThrottlerShowCmd(ch) + cmd.SetArgs([]string{db}) + err := cmd.Execute() + c.Assert(err, qt.IsNotNil) + c.Assert(err.Error(), qt.Contains, "only available for Vitess (MySQL) databases") + c.Assert(err.Error(), qt.Contains, "postgresql") + c.Assert(svc.GetThrottlerFnInvoked, qt.IsFalse) +} + +func TestDatabase_ThrottlerUpdateCmd_RejectsPostgres(t *testing.T) { + c := qt.New(t) + + format := printer.JSON + p := printer.NewPrinter(&format) + org := "planetscale" + db := "eagle" + + svc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Name: db, Kind: ps.DatabaseEnginePostgres}, nil + }, + } + + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: org}, + Client: func() (*ps.Client, error) { + return &ps.Client{Databases: svc}, nil + }, + } + + cmd := ThrottlerUpdateCmd(ch) + cmd.SetArgs([]string{db, "--ratio", "10"}) + err := cmd.Execute() + c.Assert(err, qt.IsNotNil) + c.Assert(err.Error(), qt.Contains, "only available for Vitess (MySQL) databases") + c.Assert(svc.UpdateThrottlerFnInvoked, qt.IsFalse) +} diff --git a/internal/mock/database.go b/internal/mock/database.go index 4a9af780..2a18b13d 100644 --- a/internal/mock/database.go +++ b/internal/mock/database.go @@ -21,6 +21,12 @@ type DatabaseService struct { UpdateSettingsFn func(context.Context, *ps.UpdateDatabaseSettingsRequest) (*ps.Database, error) UpdateSettingsFnInvoked bool + + GetThrottlerFn func(context.Context, *ps.GetDatabaseThrottlerRequest) (*ps.DatabaseThrottler, error) + GetThrottlerFnInvoked bool + + UpdateThrottlerFn func(context.Context, *ps.UpdateDatabaseThrottlerRequest) (*ps.DatabaseThrottler, error) + UpdateThrottlerFnInvoked bool } func (d *DatabaseService) Create(ctx context.Context, req *ps.CreateDatabaseRequest) (*ps.Database, error) { @@ -47,3 +53,13 @@ func (d *DatabaseService) UpdateSettings(ctx context.Context, req *ps.UpdateData d.UpdateSettingsFnInvoked = true return d.UpdateSettingsFn(ctx, req) } + +func (d *DatabaseService) GetThrottler(ctx context.Context, req *ps.GetDatabaseThrottlerRequest) (*ps.DatabaseThrottler, error) { + d.GetThrottlerFnInvoked = true + return d.GetThrottlerFn(ctx, req) +} + +func (d *DatabaseService) UpdateThrottler(ctx context.Context, req *ps.UpdateDatabaseThrottlerRequest) (*ps.DatabaseThrottler, error) { + d.UpdateThrottlerFnInvoked = true + return d.UpdateThrottlerFn(ctx, req) +} diff --git a/internal/planetscale/databases.go b/internal/planetscale/databases.go index 5c36bbee..0eb22f61 100644 --- a/internal/planetscale/databases.go +++ b/internal/planetscale/databases.go @@ -74,6 +74,29 @@ type UpdateDatabaseSettingsRequest struct { DefaultBranch *string `json:"default_branch,omitempty"` } +// GetDatabaseThrottlerRequest gets the database-level Vitess migration throttler. +type GetDatabaseThrottlerRequest struct { + Organization string `json:"-"` + Database string `json:"-"` +} + +// UpdateDatabaseThrottlerRequest updates the database-level Vitess migration throttler. +type UpdateDatabaseThrottlerRequest struct { + Organization string `json:"-"` + Database string `json:"-"` + Ratio *int `json:"ratio,omitempty"` + Configurations []*UpdateThrottlerConfiguration `json:"configurations,omitempty"` +} + +// DatabaseThrottler is the database-level throttler configuration for Vitess +// deploy request migrations. Distinct from per-deploy-request throttler and +// from tablet/vtctld throttler. +type DatabaseThrottler struct { + Keyspaces []string `json:"keyspaces"` + Configurable *ThrottlerConfigurable `json:"configurable"` + Configurations []*ThrottlerConfiguration `json:"configurations"` +} + // DatabaseService is an interface for communicating with the PlanetScale // Databases API endpoint. type DatabasesService interface { @@ -82,6 +105,8 @@ type DatabasesService interface { List(context.Context, *ListDatabasesRequest, ...ListOption) ([]*Database, error) Delete(context.Context, *DeleteDatabaseRequest) (*DatabaseDeletionRequest, error) UpdateSettings(context.Context, *UpdateDatabaseSettingsRequest) (*Database, error) + GetThrottler(context.Context, *GetDatabaseThrottlerRequest) (*DatabaseThrottler, error) + UpdateThrottler(context.Context, *UpdateDatabaseThrottlerRequest) (*DatabaseThrottler, error) } // DatabaseDeletionRequest encapsulates the request for deleting a database from @@ -231,6 +256,36 @@ func (ds *databasesService) UpdateSettings(ctx context.Context, updateReq *Updat return db, nil } +func (ds *databasesService) GetThrottler(ctx context.Context, getReq *GetDatabaseThrottlerRequest) (*DatabaseThrottler, error) { + pathStr := path.Join(databasesAPIPath(getReq.Organization), getReq.Database, "throttler") + req, err := ds.client.newRequest(http.MethodGet, pathStr, nil) + if err != nil { + return nil, fmt.Errorf("error creating request for get database throttler: %w", err) + } + + throttler := &DatabaseThrottler{} + if err := ds.client.do(ctx, req, throttler); err != nil { + return nil, err + } + + return throttler, nil +} + +func (ds *databasesService) UpdateThrottler(ctx context.Context, updateReq *UpdateDatabaseThrottlerRequest) (*DatabaseThrottler, error) { + pathStr := path.Join(databasesAPIPath(updateReq.Organization), updateReq.Database, "throttler") + req, err := ds.client.newRequest(http.MethodPatch, pathStr, updateReq) + if err != nil { + return nil, fmt.Errorf("error creating request for update database throttler: %w", err) + } + + throttler := &DatabaseThrottler{} + if err := ds.client.do(ctx, req, throttler); err != nil { + return nil, err + } + + return throttler, nil +} + func databasesAPIPath(org string) string { return path.Join("v1/organizations", org, "databases") } diff --git a/internal/planetscale/databases_test.go b/internal/planetscale/databases_test.go index 4535b7f7..a150a606 100644 --- a/internal/planetscale/databases_test.go +++ b/internal/planetscale/databases_test.go @@ -580,3 +580,64 @@ func TestDatabases_UpdateSettings(t *testing.T) { c.Assert(db.RequireApprovalForDeploy, qt.IsTrue) c.Assert(db.InsightsRawQueries, qt.IsFalse) } + +func TestDatabases_GetThrottler(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c.Assert(r.Method, qt.Equals, http.MethodGet) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/planetscale-go-test-db/throttler") + w.WriteHeader(200) + _, err := w.Write([]byte(`{"keyspaces":["main"],"configurable":{"id":"db-1","name":"planetscale-go-test-db"},"configurations":[{"keyspace_name":"main","ratio":50}]}`)) + c.Assert(err, qt.IsNil) + })) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + throttler, err := client.Databases.GetThrottler(context.Background(), &GetDatabaseThrottlerRequest{ + Organization: testOrg, + Database: testDatabase, + }) + c.Assert(err, qt.IsNil) + c.Assert(throttler.Keyspaces, qt.DeepEquals, []string{"main"}) + c.Assert(throttler.Configurations, qt.HasLen, 1) + c.Assert(throttler.Configurations[0].Ratio, qt.Equals, float64(50)) +} + +func TestDatabases_UpdateThrottler(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c.Assert(r.Method, qt.Equals, http.MethodPatch) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/planetscale-go-test-db/throttler") + + var body UpdateDatabaseThrottlerRequest + c.Assert(json.NewDecoder(r.Body).Decode(&body), qt.IsNil) + c.Assert(body.Ratio, qt.IsNotNil) + c.Assert(*body.Ratio, qt.Equals, 25) + c.Assert(body.Configurations, qt.DeepEquals, []*UpdateThrottlerConfiguration{ + {KeyspaceName: "main", Ratio: 10}, + }) + + w.WriteHeader(200) + _, err := w.Write([]byte(`{"keyspaces":["main"],"configurations":[{"keyspace_name":"main","ratio":10}]}`)) + c.Assert(err, qt.IsNil) + })) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + ratio := 25 + throttler, err := client.Databases.UpdateThrottler(context.Background(), &UpdateDatabaseThrottlerRequest{ + Organization: testOrg, + Database: testDatabase, + Ratio: &ratio, + Configurations: []*UpdateThrottlerConfiguration{ + {KeyspaceName: "main", Ratio: 10}, + }, + }) + c.Assert(err, qt.IsNil) + c.Assert(throttler.Configurations, qt.HasLen, 1) + c.Assert(throttler.Configurations[0].Ratio, qt.Equals, float64(10)) +}