Skip to content

Commit cad8c27

Browse files
agneumagneum
committed
Merge branch '247-reset-to-specific-snapshot' into 'master'
feat: reset clone state to a specific snapshot (#247) Closes #247 See merge request postgres-ai/database-lab!318
2 parents 4f2c26e + fd1f050 commit cad8c27

File tree

11 files changed

+121
-46
lines changed

11 files changed

+121
-46
lines changed

‎cmd/cli/commands/clone/actions.go

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,11 +158,15 @@ func reset() func(*cli.Context) error {
158158
}
159159

160160
cloneID := cliCtx.Args().First()
161+
resetOptions := types.ResetCloneRequest{
162+
Latest: cliCtx.Bool(cloneResetLatestFlag),
163+
SnapshotID: cliCtx.String(cloneResetSnapshotIDFlag),
164+
}
161165

162166
if cliCtx.Bool("async") {
163-
err = dblabClient.ResetCloneAsync(cliCtx.Context, cloneID)
167+
err = dblabClient.ResetCloneAsync(cliCtx.Context, cloneID, resetOptions)
164168
} else {
165-
err = dblabClient.ResetClone(cliCtx.Context, cloneID)
169+
err = dblabClient.ResetClone(cliCtx.Context, cloneID, resetOptions)
166170
}
167171

168172
if err != nil {

‎cmd/cli/commands/clone/command_list.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,11 @@ import (
1010
"gitlab.com/postgres-ai/database-lab/v2/cmd/cli/commands"
1111
)
1212

13+
const (
14+
cloneResetLatestFlag = "latest"
15+
cloneResetSnapshotIDFlag = "snapshot-id"
16+
)
17+
1318
// CommandList returns available commands for a clones management.
1419
func CommandList() []*cli.Command {
1520
return []*cli.Command{{
@@ -101,6 +106,14 @@ func CommandList() []*cli.Command {
101106
Usage: "run the command asynchronously",
102107
Aliases: []string{"a"},
103108
},
109+
&cli.BoolFlag{
110+
Name: cloneResetLatestFlag,
111+
Usage: "reset clone to the latest available snapshot",
112+
},
113+
&cli.StringFlag{
114+
Name: cloneResetSnapshotIDFlag,
115+
Usage: "snapshot ID used when resetting clone's state",
116+
},
104117
},
105118
},
106119
{

‎pkg/client/dblabapi/clone.go

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -197,10 +197,15 @@ func (c *Client) UpdateClone(ctx context.Context, cloneID string, updateRequest
197197
}
198198

199199
// ResetClone resets a Database Lab clone session.
200-
func (c *Client) ResetClone(ctx context.Context, cloneID string) error {
200+
func (c *Client) ResetClone(ctx context.Context, cloneID string, params types.ResetCloneRequest) error {
201201
u := c.URL(fmt.Sprintf("/clone/%s/reset", cloneID))
202202

203-
request, err := http.NewRequest(http.MethodPost, u.String(), nil)
203+
body := bytes.NewBuffer(nil)
204+
if err := json.NewEncoder(body).Encode(params); err != nil {
205+
return errors.Wrap(err, "failed to encode ResetClone parameters to JSON")
206+
}
207+
208+
request, err := http.NewRequest(http.MethodPost, u.String(), body)
204209
if err != nil {
205210
return errors.Wrap(err, "failed to make a request")
206211
}
@@ -225,10 +230,15 @@ func (c *Client) ResetClone(ctx context.Context, cloneID string) error {
225230
}
226231

227232
// ResetCloneAsync asynchronously resets a Database Lab clone session.
228-
func (c *Client) ResetCloneAsync(ctx context.Context, cloneID string) error {
233+
func (c *Client) ResetCloneAsync(ctx context.Context, cloneID string, params types.ResetCloneRequest) error {
229234
u := c.URL(fmt.Sprintf("/clone/%s/reset", cloneID))
230235

231-
request, err := http.NewRequest(http.MethodPost, u.String(), nil)
236+
body := bytes.NewBuffer(nil)
237+
if err := json.NewEncoder(body).Encode(params); err != nil {
238+
return errors.Wrap(err, "failed to encode ResetClone parameters to JSON")
239+
}
240+
241+
request, err := http.NewRequest(http.MethodPost, u.String(), body)
232242
if err != nil {
233243
return errors.Wrap(err, "failed to make a request")
234244
}

‎pkg/client/dblabapi/clone_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -556,7 +556,7 @@ func TestClientResetClone(t *testing.T) {
556556
c.pollingInterval = time.Millisecond
557557

558558
// Send a request.
559-
err = c.ResetClone(context.Background(), "testCloneID")
559+
err = c.ResetClone(context.Background(), "testCloneID", types.ResetCloneRequest{Latest: true})
560560
require.NoError(t, err)
561561
}
562562

@@ -580,7 +580,7 @@ func TestClientResetCloneAsync(t *testing.T) {
580580
c.client = mockClient
581581

582582
// Send a request.
583-
err = c.ResetCloneAsync(context.Background(), "testCloneID")
583+
err = c.ResetCloneAsync(context.Background(), "testCloneID", types.ResetCloneRequest{Latest: true})
584584
require.NoError(t, err)
585585
}
586586

@@ -611,6 +611,6 @@ func TestClientResetCloneWithFailedRequest(t *testing.T) {
611611
c.client = mockClient
612612

613613
// Send a request.
614-
err = c.ResetClone(context.Background(), "testCloneID")
614+
err = c.ResetClone(context.Background(), "testCloneID", types.ResetCloneRequest{Latest: true, SnapshotID: "test"})
615615
assert.EqualError(t, err, `failed to get response: Check your verification token.`)
616616
}

‎pkg/client/dblabapi/types/clone.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,9 @@ type DatabaseRequest struct {
3131
type SnapshotCloneFieldRequest struct {
3232
ID string `json:"id"`
3333
}
34+
35+
// ResetCloneRequest represents snapshot params of a reset request.
36+
type ResetCloneRequest struct {
37+
SnapshotID string `json:"snapshotID"`
38+
Latest bool `json:"latest"`
39+
}

‎pkg/models/clone.go

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,3 @@ type CloneMetadata struct {
2323
CloningTime float64 `json:"cloningTime"`
2424
MaxIdleMinutes uint `json:"maxIdleMinutes"`
2525
}
26-
27-
// PatchCloneRequest defines a struct for clone updating.
28-
type PatchCloneRequest struct {
29-
Protected bool `json:"protected"`
30-
}

‎pkg/services/cloning/cloning.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,9 @@ type Cloning interface {
3636
CloneConnection(ctx context.Context, cloneID string) (pgxtype.Querier, error)
3737
DestroyClone(string) error
3838
GetClone(string) (*models.Clone, error)
39-
UpdateClone(string, *types.CloneUpdateRequest) (*models.Clone, error)
39+
UpdateClone(string, types.CloneUpdateRequest) (*models.Clone, error)
4040
UpdateCloneStatus(string, models.Status) error
41-
ResetClone(string) error
41+
ResetClone(string, types.ResetCloneRequest) error
4242

4343
GetInstanceState() (*models.InstanceStatus, error)
4444
GetSnapshots() ([]models.Snapshot, error)

‎pkg/services/cloning/mode_base.go

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,7 @@ func (c *baseCloning) GetClone(id string) (*models.Clone, error) {
295295
return w.clone, nil
296296
}
297297

298-
func (c *baseCloning) UpdateClone(id string, patch *types.CloneUpdateRequest) (*models.Clone, error) {
298+
func (c *baseCloning) UpdateClone(id string, patch types.CloneUpdateRequest) (*models.Clone, error) {
299299
w, ok := c.findWrapper(id)
300300
if !ok {
301301
return nil, models.New(models.ErrCodeNotFound, "clone not found")
@@ -328,7 +328,7 @@ func (c *baseCloning) UpdateCloneStatus(cloneID string, status models.Status) er
328328
return nil
329329
}
330330

331-
func (c *baseCloning) ResetClone(cloneID string) error {
331+
func (c *baseCloning) ResetClone(cloneID string, resetOptions types.ResetCloneRequest) error {
332332
w, ok := c.findWrapper(cloneID)
333333
if !ok {
334334
return models.New(models.ErrCodeNotFound, "the clone not found")
@@ -338,6 +338,22 @@ func (c *baseCloning) ResetClone(cloneID string) error {
338338
return models.New(models.ErrCodeNotFound, "clone is not started yet")
339339
}
340340

341+
var snapshotID string
342+
343+
if resetOptions.SnapshotID != "" {
344+
snapshot, err := c.getSnapshotByID(resetOptions.SnapshotID)
345+
if err != nil {
346+
return errors.Wrap(err, "failed to get snapshot ID")
347+
}
348+
349+
snapshotID = snapshot.ID
350+
}
351+
352+
// If the snapshotID variable is empty, the latest snapshot will be chosen.
353+
if snapshotID == "" && !resetOptions.Latest {
354+
snapshotID = w.snapshot.ID
355+
}
356+
341357
if err := c.UpdateCloneStatus(cloneID, models.Status{
342358
Code: models.StatusResetting,
343359
Message: models.CloneMessageResetting,
@@ -346,9 +362,9 @@ func (c *baseCloning) ResetClone(cloneID string) error {
346362
}
347363

348364
go func() {
349-
err := c.provision.ResetSession(w.session, w.snapshot.ID)
365+
snapshot, err := c.provision.ResetSession(w.session, snapshotID)
350366
if err != nil {
351-
log.Errf("Failed to reset a clone: %+v.", err)
367+
log.Errf("Failed to reset clone: %+v.", err)
352368

353369
if updateErr := c.UpdateCloneStatus(cloneID, models.Status{
354370
Code: models.StatusFatal,
@@ -360,6 +376,10 @@ func (c *baseCloning) ResetClone(cloneID string) error {
360376
return
361377
}
362378

379+
c.cloneMutex.Lock()
380+
w.clone.Snapshot = snapshot
381+
c.cloneMutex.Unlock()
382+
363383
if err := c.UpdateCloneStatus(cloneID, models.Status{
364384
Code: models.StatusOK,
365385
Message: models.CloneMessageOK,

‎pkg/services/cloning/mode_mock.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ func (c *mockCloning) GetClone(id string) (*models.Clone, error) {
106106
return clone, nil
107107
}
108108

109-
func (c *mockCloning) UpdateClone(id string, patch *types.CloneUpdateRequest) (*models.Clone, error) {
109+
func (c *mockCloning) UpdateClone(id string, _ types.CloneUpdateRequest) (*models.Clone, error) {
110110
if _, ok := c.clones[id]; !ok {
111111
return nil, errors.New("clone not found")
112112
}
@@ -122,7 +122,7 @@ func (c *mockCloning) UpdateCloneStatus(id string, _ models.Status) error {
122122
return nil
123123
}
124124

125-
func (c *mockCloning) ResetClone(id string) error {
125+
func (c *mockCloning) ResetClone(id string, _ types.ResetCloneRequest) error {
126126
if _, ok := c.clones[id]; !ok {
127127
return errors.New("clone not found")
128128
}

‎pkg/services/provision/mode_local.go

Lines changed: 37 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/pkg/errors"
2222

2323
"gitlab.com/postgres-ai/database-lab/v2/pkg/log"
24+
"gitlab.com/postgres-ai/database-lab/v2/pkg/models"
2425
"gitlab.com/postgres-ai/database-lab/v2/pkg/services/provision/databases/postgres"
2526
"gitlab.com/postgres-ai/database-lab/v2/pkg/services/provision/docker"
2627
"gitlab.com/postgres-ai/database-lab/v2/pkg/services/provision/pool"
@@ -148,7 +149,7 @@ func (p *Provisioner) Reload(cfg Config, dbCfg resources.DB) {
148149
// StartSession starts a new session.
149150
func (p *Provisioner) StartSession(snapshotID string, user resources.EphemeralUser,
150151
extraConfig map[string]string) (*resources.Session, error) {
151-
snapshotID, err := p.getSnapshotID(snapshotID)
152+
snapshot, err := p.getSnapshot(snapshotID)
152153
if err != nil {
153154
return nil, errors.Wrap(err, "failed to get snapshots")
154155
}
@@ -173,8 +174,8 @@ func (p *Provisioner) StartSession(snapshotID string, user resources.EphemeralUs
173174
}
174175
}()
175176

176-
if err := fsm.CreateClone(name, snapshotID); err != nil {
177-
return nil, errors.Wrap(err, "failed to create a clone")
177+
if err := fsm.CreateClone(name, snapshot.ID); err != nil {
178+
return nil, errors.Wrap(err, "failed to create clone")
178179
}
179180

180181
appConfig := p.getAppConfig(fsm.Pool(), name, port)
@@ -228,19 +229,21 @@ func (p *Provisioner) StopSession(session *resources.Session) error {
228229
}
229230

230231
// ResetSession resets an existing session.
231-
func (p *Provisioner) ResetSession(session *resources.Session, snapshotID string) error {
232+
func (p *Provisioner) ResetSession(session *resources.Session, snapshotID string) (*models.Snapshot, error) {
232233
fsm, err := p.pm.GetFSManager(session.Pool)
233234
if err != nil {
234-
return errors.Wrap(err, "failed to find a filesystem manager of this session")
235+
return nil, errors.Wrap(err, "failed to find filesystem manager of this session")
235236
}
236237

237238
name := util.GetCloneName(session.Port)
238239

239-
snapshotID, err = p.getSnapshotID(snapshotID)
240+
snapshot, err := p.getSnapshot(snapshotID)
240241
if err != nil {
241-
return errors.Wrap(err, "failed to get snapshots")
242+
return nil, errors.Wrap(err, "failed to get snapshots")
242243
}
243244

245+
log.Dbg("Snapshot ID to reset session: ", snapshot.ID)
246+
244247
defer func() {
245248
if err != nil {
246249
p.revertSession(name)
@@ -251,26 +254,32 @@ func (p *Provisioner) ResetSession(session *resources.Session, snapshotID string
251254
appConfig.SetExtraConf(session.ExtraConfig)
252255

253256
if err := postgres.Stop(p.runner, fsm.Pool(), name); err != nil {
254-
return errors.Wrap(err, "failed to stop a container")
257+
return nil, errors.Wrap(err, "failed to stop container")
255258
}
256259

257260
if err := fsm.DestroyClone(name); err != nil {
258-
return errors.Wrap(err, "failed to destroy clone")
261+
return nil, errors.Wrap(err, "failed to destroy clone")
259262
}
260263

261-
if err := fsm.CreateClone(name, snapshotID); err != nil {
262-
return errors.Wrap(err, "failed to create a clone")
264+
if err := fsm.CreateClone(name, snapshot.ID); err != nil {
265+
return nil, errors.Wrap(err, "failed to create clone")
263266
}
264267

265268
if err := postgres.Start(p.runner, appConfig); err != nil {
266-
return errors.Wrap(err, "failed to start a container")
269+
return nil, errors.Wrap(err, "failed to start container")
267270
}
268271

269272
if err := p.prepareDB(appConfig, session.EphemeralUser); err != nil {
270-
return errors.Wrap(err, "failed to prepare a database")
273+
return nil, errors.Wrap(err, "failed to prepare database")
271274
}
272275

273-
return nil
276+
snapshotModel := &models.Snapshot{
277+
ID: snapshot.ID,
278+
CreatedAt: util.FormatTime(snapshot.CreatedAt),
279+
DataStateAt: util.FormatTime(snapshot.DataStateAt),
280+
}
281+
282+
return snapshotModel, nil
274283
}
275284

276285
// GetSnapshots provides a snapshot list.
@@ -306,21 +315,27 @@ func (p *Provisioner) revertSession(name string) {
306315
}
307316
}
308317

309-
func (p *Provisioner) getSnapshotID(snapshotID string) (string, error) {
310-
if snapshotID != "" {
311-
return snapshotID, nil
312-
}
313-
318+
func (p *Provisioner) getSnapshot(snapshotID string) (*resources.Snapshot, error) {
314319
snapshots, err := p.GetSnapshots()
315320
if err != nil {
316-
return "", errors.Wrap(err, "failed to get snapshots")
321+
return nil, errors.Wrap(err, "failed to get snapshots")
317322
}
318323

319324
if len(snapshots) == 0 {
320-
return "", errors.New("no snapshots available")
325+
return nil, errors.New("no snapshots available")
326+
}
327+
328+
if snapshotID != "" {
329+
for _, snapshot := range snapshots {
330+
if snapshot.ID == snapshotID {
331+
return &snapshot, nil
332+
}
333+
}
334+
335+
return nil, errors.Errorf("snapshot %q not found", snapshotID)
321336
}
322337

323-
return snapshots[0].ID, nil
338+
return &snapshots[0], nil
324339
}
325340

326341
func (p *Provisioner) initPortPool() error {

0 commit comments

Comments
 (0)