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

Author: agneum

cmd/cli/commands/clone/actions.go

@@ -158,11 +158,15 @@ func reset() func(*cli.Context) error {
 		}
 
 		cloneID := cliCtx.Args().First()
+		resetOptions := types.ResetCloneRequest{
+			Latest:     cliCtx.Bool(cloneResetLatestFlag),
+			SnapshotID: cliCtx.String(cloneResetSnapshotIDFlag),
+		}
 
 		if cliCtx.Bool("async") {
-			err = dblabClient.ResetCloneAsync(cliCtx.Context, cloneID)
+			err = dblabClient.ResetCloneAsync(cliCtx.Context, cloneID, resetOptions)
 		} else {
-			err = dblabClient.ResetClone(cliCtx.Context, cloneID)
+			err = dblabClient.ResetClone(cliCtx.Context, cloneID, resetOptions)
 		}
 
 		if err != nil {

cmd/cli/commands/clone/command_list.go

@@ -10,6 +10,11 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v2/cmd/cli/commands"
 )
 
+const (
+	cloneResetLatestFlag     = "latest"
+	cloneResetSnapshotIDFlag = "snapshot-id"
+)
+
 // CommandList returns available commands for a clones management.
 func CommandList() []*cli.Command {
 	return []*cli.Command{{
@@ -101,6 +106,14 @@ func CommandList() []*cli.Command {
 						Usage:   "run the command asynchronously",
 						Aliases: []string{"a"},
 					},
+					&cli.BoolFlag{
+						Name:  cloneResetLatestFlag,
+						Usage: "reset clone to the latest available snapshot",
+					},
+					&cli.StringFlag{
+						Name:  cloneResetSnapshotIDFlag,
+						Usage: "snapshot ID used when resetting clone's state",
+					},
 				},
 			},
 			{

pkg/client/dblabapi/clone.go

@@ -197,10 +197,15 @@ func (c *Client) UpdateClone(ctx context.Context, cloneID string, updateRequest
 }
 
 // ResetClone resets a Database Lab clone session.
-func (c *Client) ResetClone(ctx context.Context, cloneID string) error {
+func (c *Client) ResetClone(ctx context.Context, cloneID string, params types.ResetCloneRequest) error {
 	u := c.URL(fmt.Sprintf("/clone/%s/reset", cloneID))
 
-	request, err := http.NewRequest(http.MethodPost, u.String(), nil)
+	body := bytes.NewBuffer(nil)
+	if err := json.NewEncoder(body).Encode(params); err != nil {
+		return errors.Wrap(err, "failed to encode ResetClone parameters to JSON")
+	}
+
+	request, err := http.NewRequest(http.MethodPost, u.String(), body)
 	if err != nil {
 		return errors.Wrap(err, "failed to make a request")
 	}
@@ -225,10 +230,15 @@ func (c *Client) ResetClone(ctx context.Context, cloneID string) error {
 }
 
 // ResetCloneAsync asynchronously resets a Database Lab clone session.
-func (c *Client) ResetCloneAsync(ctx context.Context, cloneID string) error {
+func (c *Client) ResetCloneAsync(ctx context.Context, cloneID string, params types.ResetCloneRequest) error {
 	u := c.URL(fmt.Sprintf("/clone/%s/reset", cloneID))
 
-	request, err := http.NewRequest(http.MethodPost, u.String(), nil)
+	body := bytes.NewBuffer(nil)
+	if err := json.NewEncoder(body).Encode(params); err != nil {
+		return errors.Wrap(err, "failed to encode ResetClone parameters to JSON")
+	}
+
+	request, err := http.NewRequest(http.MethodPost, u.String(), body)
 	if err != nil {
 		return errors.Wrap(err, "failed to make a request")
 	}

pkg/client/dblabapi/clone_test.go

@@ -556,7 +556,7 @@ func TestClientResetClone(t *testing.T) {
 	c.pollingInterval = time.Millisecond
 
 	// Send a request.
-	err = c.ResetClone(context.Background(), "testCloneID")
+	err = c.ResetClone(context.Background(), "testCloneID", types.ResetCloneRequest{Latest: true})
 	require.NoError(t, err)
 }
 
@@ -580,7 +580,7 @@ func TestClientResetCloneAsync(t *testing.T) {
 	c.client = mockClient
 
 	// Send a request.
-	err = c.ResetCloneAsync(context.Background(), "testCloneID")
+	err = c.ResetCloneAsync(context.Background(), "testCloneID", types.ResetCloneRequest{Latest: true})
 	require.NoError(t, err)
 }
 
@@ -611,6 +611,6 @@ func TestClientResetCloneWithFailedRequest(t *testing.T) {
 	c.client = mockClient
 
 	// Send a request.
-	err = c.ResetClone(context.Background(), "testCloneID")
+	err = c.ResetClone(context.Background(), "testCloneID", types.ResetCloneRequest{Latest: true, SnapshotID: "test"})
 	assert.EqualError(t, err, `failed to get response: Check your verification token.`)
 }

pkg/client/dblabapi/types/clone.go

@@ -31,3 +31,9 @@ type DatabaseRequest struct {
 type SnapshotCloneFieldRequest struct {
 	ID string `json:"id"`
 }
+
+// ResetCloneRequest represents snapshot params of a reset request.
+type ResetCloneRequest struct {
+	SnapshotID string `json:"snapshotID"`
+	Latest     bool   `json:"latest"`
+}

pkg/models/clone.go

@@ -23,8 +23,3 @@ type CloneMetadata struct {
 	CloningTime     float64 `json:"cloningTime"`
 	MaxIdleMinutes  uint    `json:"maxIdleMinutes"`
 }
-
-// PatchCloneRequest defines a struct for clone updating.
-type PatchCloneRequest struct {
-	Protected bool `json:"protected"`
-}

pkg/services/cloning/cloning.go

@@ -36,9 +36,9 @@ type Cloning interface {
 	CloneConnection(ctx context.Context, cloneID string) (pgxtype.Querier, error)
 	DestroyClone(string) error
 	GetClone(string) (*models.Clone, error)
-	UpdateClone(string, *types.CloneUpdateRequest) (*models.Clone, error)
+	UpdateClone(string, types.CloneUpdateRequest) (*models.Clone, error)
 	UpdateCloneStatus(string, models.Status) error
-	ResetClone(string) error
+	ResetClone(string, types.ResetCloneRequest) error
 
 	GetInstanceState() (*models.InstanceStatus, error)
 	GetSnapshots() ([]models.Snapshot, error)

pkg/services/cloning/mode_base.go

@@ -295,7 +295,7 @@ func (c *baseCloning) GetClone(id string) (*models.Clone, error) {
 	return w.clone, nil
 }
 
-func (c *baseCloning) UpdateClone(id string, patch *types.CloneUpdateRequest) (*models.Clone, error) {
+func (c *baseCloning) UpdateClone(id string, patch types.CloneUpdateRequest) (*models.Clone, error) {
 	w, ok := c.findWrapper(id)
 	if !ok {
 		return nil, models.New(models.ErrCodeNotFound, "clone not found")
@@ -328,7 +328,7 @@ func (c *baseCloning) UpdateCloneStatus(cloneID string, status models.Status) er
 	return nil
 }
 
-func (c *baseCloning) ResetClone(cloneID string) error {
+func (c *baseCloning) ResetClone(cloneID string, resetOptions types.ResetCloneRequest) error {
 	w, ok := c.findWrapper(cloneID)
 	if !ok {
 		return models.New(models.ErrCodeNotFound, "the clone not found")
@@ -338,6 +338,22 @@ func (c *baseCloning) ResetClone(cloneID string) error {
 		return models.New(models.ErrCodeNotFound, "clone is not started yet")
 	}
 
+	var snapshotID string
+
+	if resetOptions.SnapshotID != "" {
+		snapshot, err := c.getSnapshotByID(resetOptions.SnapshotID)
+		if err != nil {
+			return errors.Wrap(err, "failed to get snapshot ID")
+		}
+
+		snapshotID = snapshot.ID
+	}
+
+	// If the snapshotID variable is empty, the latest snapshot will be chosen.
+	if snapshotID == "" && !resetOptions.Latest {
+		snapshotID = w.snapshot.ID
+	}
+
 	if err := c.UpdateCloneStatus(cloneID, models.Status{
 		Code:    models.StatusResetting,
 		Message: models.CloneMessageResetting,
@@ -346,9 +362,9 @@ func (c *baseCloning) ResetClone(cloneID string) error {
 	}
 
 	go func() {
-		err := c.provision.ResetSession(w.session, w.snapshot.ID)
+		snapshot, err := c.provision.ResetSession(w.session, snapshotID)
 		if err != nil {
-			log.Errf("Failed to reset a clone: %+v.", err)
+			log.Errf("Failed to reset clone: %+v.", err)
 
 			if updateErr := c.UpdateCloneStatus(cloneID, models.Status{
 				Code:    models.StatusFatal,
@@ -360,6 +376,10 @@ func (c *baseCloning) ResetClone(cloneID string) error {
 			return
 		}
 
+		c.cloneMutex.Lock()
+		w.clone.Snapshot = snapshot
+		c.cloneMutex.Unlock()
+
 		if err := c.UpdateCloneStatus(cloneID, models.Status{
 			Code:    models.StatusOK,
 			Message: models.CloneMessageOK,

pkg/services/cloning/mode_mock.go

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

pkg/services/provision/mode_local.go

@@ -21,6 +21,7 @@ import (
 	"github.com/pkg/errors"
 
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/log"
+	"gitlab.com/postgres-ai/database-lab/v2/pkg/models"
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/services/provision/databases/postgres"
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/services/provision/docker"
 	"gitlab.com/postgres-ai/database-lab/v2/pkg/services/provision/pool"
@@ -148,7 +149,7 @@ func (p *Provisioner) Reload(cfg Config, dbCfg resources.DB) {
 // StartSession starts a new session.
 func (p *Provisioner) StartSession(snapshotID string, user resources.EphemeralUser,
 	extraConfig map[string]string) (*resources.Session, error) {
-	snapshotID, err := p.getSnapshotID(snapshotID)
+	snapshot, err := p.getSnapshot(snapshotID)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get snapshots")
 	}
@@ -173,8 +174,8 @@ func (p *Provisioner) StartSession(snapshotID string, user resources.EphemeralUs
 		}
 	}()
 
-	if err := fsm.CreateClone(name, snapshotID); err != nil {
-		return nil, errors.Wrap(err, "failed to create a clone")
+	if err := fsm.CreateClone(name, snapshot.ID); err != nil {
+		return nil, errors.Wrap(err, "failed to create clone")
 	}
 
 	appConfig := p.getAppConfig(fsm.Pool(), name, port)
@@ -228,19 +229,21 @@ func (p *Provisioner) StopSession(session *resources.Session) error {
 }
 
 // ResetSession resets an existing session.
-func (p *Provisioner) ResetSession(session *resources.Session, snapshotID string) error {
+func (p *Provisioner) ResetSession(session *resources.Session, snapshotID string) (*models.Snapshot, error) {
 	fsm, err := p.pm.GetFSManager(session.Pool)
 	if err != nil {
-		return errors.Wrap(err, "failed to find a filesystem manager of this session")
+		return nil, errors.Wrap(err, "failed to find filesystem manager of this session")
 	}
 
 	name := util.GetCloneName(session.Port)
 
-	snapshotID, err = p.getSnapshotID(snapshotID)
+	snapshot, err := p.getSnapshot(snapshotID)
 	if err != nil {
-		return errors.Wrap(err, "failed to get snapshots")
+		return nil, errors.Wrap(err, "failed to get snapshots")
 	}
 
+	log.Dbg("Snapshot ID to reset session: ", snapshot.ID)
+
 	defer func() {
 		if err != nil {
 			p.revertSession(name)
@@ -251,26 +254,32 @@ func (p *Provisioner) ResetSession(session *resources.Session, snapshotID string
 	appConfig.SetExtraConf(session.ExtraConfig)
 
 	if err := postgres.Stop(p.runner, fsm.Pool(), name); err != nil {
-		return errors.Wrap(err, "failed to stop a container")
+		return nil, errors.Wrap(err, "failed to stop container")
 	}
 
 	if err := fsm.DestroyClone(name); err != nil {
-		return errors.Wrap(err, "failed to destroy clone")
+		return nil, errors.Wrap(err, "failed to destroy clone")
 	}
 
-	if err := fsm.CreateClone(name, snapshotID); err != nil {
-		return errors.Wrap(err, "failed to create a clone")
+	if err := fsm.CreateClone(name, snapshot.ID); err != nil {
+		return nil, errors.Wrap(err, "failed to create clone")
 	}
 
 	if err := postgres.Start(p.runner, appConfig); err != nil {
-		return errors.Wrap(err, "failed to start a container")
+		return nil, errors.Wrap(err, "failed to start container")
 	}
 
 	if err := p.prepareDB(appConfig, session.EphemeralUser); err != nil {
-		return errors.Wrap(err, "failed to prepare a database")
+		return nil, errors.Wrap(err, "failed to prepare database")
 	}
 
-	return nil
+	snapshotModel := &models.Snapshot{
+		ID:          snapshot.ID,
+		CreatedAt:   util.FormatTime(snapshot.CreatedAt),
+		DataStateAt: util.FormatTime(snapshot.DataStateAt),
+	}
+
+	return snapshotModel, nil
 }
 
 // GetSnapshots provides a snapshot list.
@@ -306,21 +315,27 @@ func (p *Provisioner) revertSession(name string) {
 	}
 }
 
-func (p *Provisioner) getSnapshotID(snapshotID string) (string, error) {
-	if snapshotID != "" {
-		return snapshotID, nil
-	}
-
+func (p *Provisioner) getSnapshot(snapshotID string) (*resources.Snapshot, error) {
 	snapshots, err := p.GetSnapshots()
 	if err != nil {
-		return "", errors.Wrap(err, "failed to get snapshots")
+		return nil, errors.Wrap(err, "failed to get snapshots")
 	}
 
 	if len(snapshots) == 0 {
-		return "", errors.New("no snapshots available")
+		return nil, errors.New("no snapshots available")
+	}
+
+	if snapshotID != "" {
+		for _, snapshot := range snapshots {
+			if snapshot.ID == snapshotID {
+				return &snapshot, nil
+			}
+		}
+
+		return nil, errors.Errorf("snapshot %q not found", snapshotID)
 	}
 
-	return snapshots[0].ID, nil
+	return &snapshots[0], nil
 }
 
 func (p *Provisioner) initPortPool() error {

pkg/srv/routes.go

@@ -101,7 +101,7 @@ func (s *Server) patchClone(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	var patchClone *types.CloneUpdateRequest
+	var patchClone types.CloneUpdateRequest
 	if err := api.ReadJSON(r, &patchClone); err != nil {
 		api.SendBadRequestError(w, r, err.Error())
 
@@ -148,7 +148,19 @@ func (s *Server) resetClone(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	if err := s.Cloning.ResetClone(cloneID); err != nil {
+	var resetOptions types.ResetCloneRequest
+
+	if err := json.NewDecoder(r.Body).Decode(&resetOptions); err != nil {
+		api.SendError(w, r, errors.Wrap(err, "failed to parse request parameters"))
+		return
+	}
+
+	if resetOptions.Latest && resetOptions.SnapshotID != "" {
+		api.SendBadRequestError(w, r, "parameters `latest` and `snapshot ID` must not be specified together")
+		return
+	}
+
+	if err := s.Cloning.ResetClone(cloneID, resetOptions); err != nil {
 		api.SendError(w, r, errors.Wrap(err, "failed to reset clone"))
 		return
 	}