Add Update{Study,Trial} (#269)

Only tested with unit tests.

Signed-off-by: IWAMOTO Toshihiro <[email protected]>

pkg/db/interface.go

@@ -46,11 +46,13 @@ type VizierDBInterface interface {
 	GetStudyConfig(string) (*api.StudyConfig, error)
 	GetStudyList() ([]string, error)
 	CreateStudy(*api.StudyConfig) (string, error)
+	UpdateStudy(string, *api.StudyConfig) error
 	DeleteStudy(string) error
 
 	GetTrial(string) (*api.Trial, error)
 	GetTrialList(string) ([]*api.Trial, error)
 	CreateTrial(*api.Trial) error
+	UpdateTrial(*api.Trial) error
 	DeleteTrial(string) error
 
 	GetWorker(string) (*api.Worker, error)
@@ -307,6 +309,30 @@ func (d *dbConn) CreateStudy(in *api.StudyConfig) (string, error) {
 	return studyID, nil
 }
 
+// UpdateStudy updates the corresponding row in the DB.
+// It only updates name, owner, tags and job_id.
+// Other columns are silently ignored.
+func (d *dbConn) UpdateStudy(studyID string, in *api.StudyConfig) error {
+	var err error
+
+	tags := make([]string, len(in.Tags))
+	for i, elem := range in.Tags {
+		tags[i], err = (&jsonpb.Marshaler{}).MarshalToString(elem)
+		if err != nil {
+			log.Printf("Error marshalling %v: %v", elem, err)
+			continue
+		}
+	}
+	_, err = d.db.Exec(`UPDATE studies SET name = ?, owner = ?, tags = ?,
+                job_id = ? WHERE id = ?`,
+		in.Name,
+		in.Owner,
+		strings.Join(tags, ",\n"),
+		in.JobId,
+		studyID)
+	return err
+}
+
 func (d *dbConn) DeleteStudy(id string) error {
 	_, err := d.db.Exec("DELETE FROM studies WHERE id = ?", id)
 	return err
@@ -395,9 +421,7 @@ func (d *dbConn) GetTrialList(id string) ([]*api.Trial, error) {
 	return trials, err
 }
 
-func (d *dbConn) CreateTrial(trial *api.Trial) error {
-	// This function sets trial.id, unlike old dbInsertTrials().
-	// Users should not overwrite trial.id
+func marshalTrial(trial *api.Trial) ([]string, []string, error) {
 	var err, lastErr error
 
 	params := make([]string, len(trial.ParameterSet))
@@ -418,12 +442,20 @@ func (d *dbConn) CreateTrial(trial *api.Trial) error {
 			lastErr = err
 		}
 	}
+	return params, tags, lastErr
+}
+
+// CreateTrial stores into the trials DB table.
+// As a side-effect, it generates and sets trial.TrialId.
+// Users should not overwrite TrialId.
+func (d *dbConn) CreateTrial(trial *api.Trial) error {
+	params, tags, lastErr := marshalTrial(trial)
 
 	var trialID string
 	i := 3
 	for true {
 		trialID = generateRandid()
-		_, err = d.db.Exec("INSERT INTO trials VALUES (?, ?, ?, ?, ?)",
+		_, err := d.db.Exec("INSERT INTO trials VALUES (?, ?, ?, ?, ?)",
 			trialID, trial.StudyId, strings.Join(params, ",\n"),
 			trial.ObjectiveValue, strings.Join(tags, ",\n"))
 		if err == nil {
@@ -440,6 +472,20 @@ func (d *dbConn) CreateTrial(trial *api.Trial) error {
 	return lastErr
 }
 
+// UpdateTrial updates the corresponding row in the DB.
+// It only updates parameters and tags. Other columns are silently ignored.
+func (d *dbConn) UpdateTrial(trial *api.Trial) error {
+	params, tags, lastErr := marshalTrial(trial)
+	_, err := d.db.Exec(`UPDATE trials SET parameters = ?, tags = ?,
+		WHERE id = ?`,
+		strings.Join(params, ",\n"), strings.Join(tags, ",\n"),
+		trial.TrialId)
+	if err != nil {
+		return err
+	}
+	return lastErr
+}
+
 func (d *dbConn) DeleteTrial(id string) error {
 	_, err := d.db.Exec("DELETE FROM trials WHERE id = ?", id)
 	return err

pkg/db/interface_test.go

@@ -135,6 +135,22 @@ func TestGetStudyList(t *testing.T) {
 	}
 }
 
+func TestUpdateStudy(t *testing.T) {
+	studyID := generateRandid()
+	var in api.StudyConfig
+	in.Name = "hoge"
+	in.Owner = "joe"
+	in.JobId = "foobar123"
+
+	mock.ExpectExec(`UPDATE studies SET name = \?, owner = \?, tags = \?,
+                job_id = \? WHERE id = \?`,
+	).WithArgs(in.Name, in.Owner, "", in.JobId, studyID).WillReturnResult(sqlmock.NewResult(1, 1))
+	err := dbInterface.UpdateStudy(studyID, &in)
+	if err != nil {
+		t.Errorf("UpdateStudy error %v", err)
+	}
+}
+
 func TestDeleteStudy(t *testing.T) {
 	studyID := generateRandid()
 	mock.ExpectExec(`DELETE FROM studies WHERE id = \?`).WithArgs(studyID).WillReturnResult(sqlmock.NewResult(1, 1))
@@ -225,6 +241,21 @@ func TestCreateTrial(t *testing.T) {
 	}
 }
 
+func TestUpdateTrial(t *testing.T) {
+	var trial api.Trial
+	trial.TrialId = generateRandid()
+	trial.StudyId = generateRandid()
+	trial.ParameterSet = make([]*api.Parameter, 1)
+	trial.ParameterSet[0] = &api.Parameter{Name: "abc"}
+	mock.ExpectExec(`UPDATE trials SET parameters = \?, tags = \?,
+		WHERE id = \?`,
+	).WithArgs("{\"name\":\"abc\"}", "", trial.TrialId).WillReturnResult(sqlmock.NewResult(1, 1))
+	err := dbInterface.UpdateTrial(&trial)
+	if err != nil {
+		t.Errorf("UpdateTrial error %v", err)
+	}
+}
+
 func TestDeleteTrial(t *testing.T) {
 	id := generateRandid()
 	mock.ExpectExec(`DELETE FROM trials WHERE id = \?`).WithArgs(id).WillReturnResult(sqlmock.NewResult(1, 1))