Skip to content

Fix Neki restore sizing and cluster size #1402

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eileencodes merged 1 commit into from
Sep 10, 2026
Merged

Fix Neki restore sizing and cluster size #1402

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions internal/cmd/backup/restore.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
2 changes: 1 addition & 1 deletion internal/cmd/backup/restore_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
5 changes: 4 additions & 1 deletion internal/cmd/branch/create.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
98 changes: 98 additions & 0 deletions internal/cmd/branch/create_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)

Expand Down
30 changes: 30 additions & 0 deletions internal/cmd/configprofile/config_profile_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion internal/cmd/configprofile/create.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
2 changes: 1 addition & 1 deletion internal/cmd/configprofile/update.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down