From 0002daef76043af5021f07c2dd3e3a86f620d3d1 Mon Sep 17 00:00:00 2001 From: eileencodes Date: Thu, 10 Sep 2026 11:06:56 -0400 Subject: [PATCH] Fix Neki restore sizing and cluster size Followup to fix bugbot comments on https://github.com/planetscale/cli/pull/1401 Accept Neki restore sizing on --restore-point and normalize cluster-size slugs. --from + --restore-point is a backup restore, so --config-profile/--router were rejected before the backup id was resolved. --cluster-size also skipped ToSizeSKUName, so hyphenated sizes from `pscale size cluster list` disagreed with the profile override path. --- internal/cmd/backup/restore.go | 8 +- internal/cmd/backup/restore_test.go | 2 +- internal/cmd/branch/create.go | 5 +- internal/cmd/branch/create_test.go | 98 +++++++++++++++++++ .../cmd/configprofile/config_profile_test.go | 30 ++++++ internal/cmd/configprofile/create.go | 2 +- internal/cmd/configprofile/update.go | 2 +- 7 files changed, 141 insertions(+), 6 deletions(-) diff --git a/internal/cmd/backup/restore.go b/internal/cmd/backup/restore.go index 08d3a1c2..93baa3e7 100644 --- a/internal/cmd/backup/restore.go +++ b/internal/cmd/backup/restore.go @@ -83,8 +83,12 @@ Preview Neki restore sizes from the source branch with: return ch.Printer.PrintResource(branch.ToDatabaseBranch(newBranch)) } else { clusterName := flags.clusterSize - if db.Kind == planetscale.DatabaseEngineNeki && !cmd.Flags().Changed("cluster-size") { - clusterName = "" + if db.Kind == planetscale.DatabaseEngineNeki { + if !cmd.Flags().Changed("cluster-size") { + clusterName = "" + } else { + clusterName = cmdutil.ToSizeSKUName(flags.clusterSize) + } } createReq := &planetscale.CreatePostgresBranchRequest{ diff --git a/internal/cmd/backup/restore_test.go b/internal/cmd/backup/restore_test.go index 02746f34..ed294303 100644 --- a/internal/cmd/backup/restore_test.go +++ b/internal/cmd/backup/restore_test.go @@ -249,7 +249,7 @@ func TestBackup_RestoreCmd_NekiClusterSize(t *testing.T) { } cmd := RestoreCmd(ch) - cmd.SetArgs([]string{"planetscale", "restore-branch", "mybackup", "--cluster-size", "PS_40"}) + cmd.SetArgs([]string{"planetscale", "restore-branch", "mybackup", "--cluster-size", "PS-40"}) c.Assert(cmd.Execute(), qt.IsNil) c.Assert(svc.CreateFnInvoked, qt.IsTrue) } diff --git a/internal/cmd/branch/create.go b/internal/cmd/branch/create.go index b55fdbe5..cb735a89 100644 --- a/internal/cmd/branch/create.go +++ b/internal/cmd/branch/create.go @@ -122,8 +122,11 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command { clusterSize = "PS_DEV" } } + if db.Kind == ps.DatabaseEngineNeki && clusterSize != "" { + clusterSize = cmdutil.ToSizeSKUName(clusterSize) + } - if err := cmdutil.EnsureNekiRestoreSizing(db.Kind, flags.backupID != "", flags.dataBranching, len(flags.configProfiles) > 0 || len(flags.routers) > 0); err != nil { + if err := cmdutil.EnsureNekiRestoreSizing(db.Kind, flags.backupID != "" || flags.restorePoint != "", flags.dataBranching, len(flags.configProfiles) > 0 || len(flags.routers) > 0); err != nil { return err } diff --git a/internal/cmd/branch/create_test.go b/internal/cmd/branch/create_test.go index a277f090..0595988c 100644 --- a/internal/cmd/branch/create_test.go +++ b/internal/cmd/branch/create_test.go @@ -518,6 +518,35 @@ func TestBranch_CreateCmdNekiRespectsClusterSize(t *testing.T) { c.Assert(buf.String(), qt.JSONEquals, res) } +func TestBranch_CreateCmdNekiNormalizesClusterSize(t *testing.T) { + c := qt.New(t) + + svc := &mock.PostgresBranchesService{ + CreateFn: func(ctx context.Context, req *ps.CreatePostgresBranchRequest) (*ps.PostgresBranch, error) { + c.Assert(req.ClusterName, qt.Equals, "PS_10_AWS_ARM_NEKI") + return &ps.PostgresBranch{Name: "development"}, nil + }, + } + dbSvc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Kind: ps.DatabaseEngineNeki}, nil + }, + } + format := printer.JSON + ch := &cmdutil.Helper{ + Printer: printer.NewPrinter(&format), + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return &ps.Client{PostgresBranches: svc, Databases: dbSvc}, nil + }, + } + + cmd := CreateCmd(ch) + cmd.SetArgs([]string{"planetscale", "development", "--from", "main", "--cluster-size", "PS-10-AWS-ARM-NEKI"}) + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(svc.CreateFnInvoked, qt.IsTrue) +} + func TestBranch_CreateCmdNekiWithWaitPrintsReadyBranch(t *testing.T) { t.Parallel() c := qt.New(t) @@ -1360,6 +1389,75 @@ func TestBranch_CreateCmdNekiRestoreWithSizes(t *testing.T) { c.Assert(buf.String(), qt.JSONEquals, res) } +func TestBranch_CreateCmdNekiRestorePointWithSizes(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" + branch := "restored" + parentBranch := "main" + restorePoint := "2023-01-01T00:00:00Z" + backupID := "backup-id" + replicas := 2 + replicasPerCell := 1 + res := &ps.PostgresBranch{Name: branch} + + backupSvc := &mock.BackupsService{ + ListFn: func(ctx context.Context, req *ps.ListBackupsRequest) ([]*ps.Backup, error) { + c.Assert(req.Organization, qt.Equals, org) + c.Assert(req.Database, qt.Equals, db) + c.Assert(req.Branch, qt.Equals, parentBranch) + return []*ps.Backup{{ + PublicID: backupID, + State: "success", + CompletedAt: time.Date(2022, time.December, 31, 23, 0, 0, 0, time.UTC), + }}, nil + }, + } + svc := &mock.PostgresBranchesService{ + CreateFn: func(ctx context.Context, req *ps.CreatePostgresBranchRequest) (*ps.PostgresBranch, error) { + c.Assert(req.BackupID, qt.Equals, backupID) + c.Assert(req.RestorePoint, qt.Equals, restorePoint) + c.Assert(req.ClusterName, qt.Equals, "") + c.Assert(req.ConfigurationProfileSizes, qt.DeepEquals, []ps.ConfigurationProfileSize{ + {Name: "default", ClusterSize: "PS_40", Replicas: &replicas}, + }) + c.Assert(req.RouterSizes, qt.DeepEquals, []ps.RouterSize{ + {Name: "default", RouterSize: "NKR_20", ReplicasPerCell: &replicasPerCell}, + }) + return res, nil + }, + } + dbSvc := &mock.DatabaseService{ + GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) { + return &ps.Database{Kind: ps.DatabaseEngineNeki}, nil + }, + } + ch := &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: org}, + Client: func() (*ps.Client, error) { + return &ps.Client{PostgresBranches: svc, Databases: dbSvc, Backups: backupSvc}, nil + }, + } + + cmd := CreateCmd(ch) + cmd.SetArgs([]string{ + db, branch, "--from", parentBranch, "--restore-point", restorePoint, + "--config-profile", "name=default,cluster-size=PS-40,replicas=2", + "--router", "name=default,size=NKR-20,replicas-per-cell=1", + }) + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(backupSvc.ListFnInvoked, qt.IsTrue) + c.Assert(svc.CreateFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.JSONEquals, res) +} + func TestBranch_CreateCmdNekiSizingRequiresRestore(t *testing.T) { c := qt.New(t) diff --git a/internal/cmd/configprofile/config_profile_test.go b/internal/cmd/configprofile/config_profile_test.go index f7aa01cc..dd6386ed 100644 --- a/internal/cmd/configprofile/config_profile_test.go +++ b/internal/cmd/configprofile/config_profile_test.go @@ -146,6 +146,21 @@ func TestConfigProfileCreateCmdSendsStorage(t *testing.T) { c.Assert(svc.CreateFnInvoked, qt.IsTrue) } +func TestConfigProfileCreateCmdNormalizesClusterSize(t *testing.T) { + c := qt.New(t) + var out bytes.Buffer + svc := &mock.NekiShardConfigurationProfilesService{CreateFn: func(_ context.Context, req *ps.CreateNekiShardConfigurationProfileRequest) (*ps.NekiShardConfigurationProfile, error) { + c.Assert(req.ClusterSize, qt.IsNotNil) + c.Assert(*req.ClusterSize, qt.Equals, "PS_40") + return testProfile(), nil + }} + + cmd := CreateCmd(configProfileTestHelper(svc, &out)) + cmd.SetArgs([]string{"app", "main", "metal", "--cluster-size", "PS-40"}) + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(svc.CreateFnInvoked, qt.IsTrue) +} + func TestConfigProfileCreateCmdRejectsExtraArguments(t *testing.T) { c := qt.New(t) var out bytes.Buffer @@ -182,6 +197,21 @@ func TestConfigProfileUpdateCmd(t *testing.T) { c.Assert(svc.UpdateFnInvoked, qt.IsTrue) } +func TestConfigProfileUpdateCmdNormalizesClusterSize(t *testing.T) { + c := qt.New(t) + var out bytes.Buffer + svc := &mock.NekiShardConfigurationProfilesService{UpdateFn: func(_ context.Context, req *ps.UpdateNekiShardConfigurationProfileRequest) (*ps.NekiShardConfigurationProfile, error) { + c.Assert(req.ClusterSize, qt.IsNotNil) + c.Assert(*req.ClusterSize, qt.Equals, "PS_40") + return testProfile(), nil + }} + + cmd := UpdateCmd(configProfileTestHelper(svc, &out)) + cmd.SetArgs([]string{"app", "main", "metal", "--cluster-size", "PS-40"}) + c.Assert(cmd.Execute(), qt.IsNil) + c.Assert(svc.UpdateFnInvoked, qt.IsTrue) +} + func TestConfigProfileUpdateCmdSendsStorage(t *testing.T) { c := qt.New(t) var out bytes.Buffer diff --git a/internal/cmd/configprofile/create.go b/internal/cmd/configprofile/create.go index 487d230e..a452ef36 100644 --- a/internal/cmd/configprofile/create.go +++ b/internal/cmd/configprofile/create.go @@ -38,7 +38,7 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command { Database: database, Branch: branch, Name: name, - ClusterSize: stringPointerIfChanged(cmd, "cluster-size", flags.clusterSize), + ClusterSize: stringPointerIfChanged(cmd, "cluster-size", cmdutil.ToSizeSKUName(flags.clusterSize)), Replicas: intPointerIfChanged(cmd, "replicas", flags.replicas), PostgresMajorVersion: stringPointerIfChanged(cmd, "postgres-major-version", flags.major), PostgresMinorVersion: stringPointerIfChanged(cmd, "postgres-minor-version", flags.minor), diff --git a/internal/cmd/configprofile/update.go b/internal/cmd/configprofile/update.go index babacfc6..7a0dede6 100644 --- a/internal/cmd/configprofile/update.go +++ b/internal/cmd/configprofile/update.go @@ -53,7 +53,7 @@ func UpdateCmd(ch *cmdutil.Helper) *cobra.Command { Branch: branch, ConfigurationProfile: profileName, Name: stringPointerIfChanged(cmd, "name", flags.name), - ClusterSize: stringPointerIfChanged(cmd, "cluster-size", flags.clusterSize), + ClusterSize: stringPointerIfChanged(cmd, "cluster-size", cmdutil.ToSizeSKUName(flags.clusterSize)), Replicas: intPointerIfChanged(cmd, "replicas", flags.replicas), Parameters: parameters, PostgresMajorVersion: stringPointerIfChanged(cmd, "postgres-major-version", flags.major),