Database APIs for NAS updated (#394)

* FINAL PUSH * FIX TESTS * new lock * new lock * small fi * DELET SPACE * deleted ununsed function
kubeflow · Feb 21, 2019 · d6a67ea · d6a67ea
1 parent 3bb8b54
commit d6a67ea
Show file tree

Hide file tree

Showing 8 changed files with 154 additions and 107 deletions.
diff --git a/cmd/manager/main.go b/cmd/manager/main.go
@@ -32,15 +32,12 @@ func (s *server) CreateStudy(ctx context.Context, in *api_pb.CreateStudyRequest)
 	if in == nil || in.StudyConfig == nil {
 		return &api_pb.CreateStudyReply{}, errors.New("StudyConfig is missing.")
 	}
+
 	studyID, err := dbIf.CreateStudy(in.StudyConfig)
 	if err != nil {
 		return &api_pb.CreateStudyReply{}, err
 	}
-	s.SaveStudy(ctx, &api_pb.SaveStudyRequest{
-		StudyName:   in.StudyConfig.Name,
-		Owner:       in.StudyConfig.Owner,
-		Description: "StudyID: " + studyID,
-	})
+
 	return &api_pb.CreateStudyReply{StudyId: studyID}, nil
 }
 
@@ -56,7 +53,12 @@ func (s *server) DeleteStudy(ctx context.Context, in *api_pb.DeleteStudyRequest)
 }
 
 func (s *server) GetStudy(ctx context.Context, in *api_pb.GetStudyRequest) (*api_pb.GetStudyReply, error) {
-	sc, err := dbIf.GetStudyConfig(in.StudyId)
+
+	sc, err := dbIf.GetStudy(in.StudyId)
+
+	if err != nil {
+		return &api_pb.GetStudyReply{}, err
+	}
 	return &api_pb.GetStudyReply{StudyConfig: sc}, err
 }
 
@@ -67,7 +69,7 @@ func (s *server) GetStudyList(ctx context.Context, in *api_pb.GetStudyListReques
 	}
 	result := make([]*api_pb.StudyOverview, len(sl))
 	for i, id := range sl {
-		sc, err := dbIf.GetStudyConfig(id)
+		sc, err := dbIf.GetStudy(id)
 		if err != nil {
 			return &api_pb.GetStudyListReply{}, err
 		}
@@ -149,7 +151,7 @@ func (s *server) GetMetrics(ctx context.Context, in *api_pb.GetMetricsRequest) (
 	if in.StudyId == "" {
 		return &api_pb.GetMetricsReply{}, errors.New("StudyId should be set")
 	}
-	sc, err := dbIf.GetStudyConfig(in.StudyId)
+	sc, err := dbIf.GetStudy(in.StudyId)
 	if err != nil {
 		return &api_pb.GetMetricsReply{}, err
 	}
@@ -214,6 +216,7 @@ func (s *server) UpdateWorkerState(ctx context.Context, in *api_pb.UpdateWorkerS
 func (s *server) GetWorkerFullInfo(ctx context.Context, in *api_pb.GetWorkerFullInfoRequest) (*api_pb.GetWorkerFullInfoReply, error) {
 	return dbIf.GetWorkerFullInfo(in.StudyId, in.TrialId, in.WorkerId, in.OnlyLatestLog)
 }
+
 func (s *server) SetSuggestionParameters(ctx context.Context, in *api_pb.SetSuggestionParametersRequest) (*api_pb.SetSuggestionParametersReply, error) {
 	var err error
 	var id string

diff --git a/cmd/manager/main_test.go b/cmd/manager/main_test.go
@@ -60,7 +60,7 @@ func TestGetStudies(t *testing.T) {
 	}
 	mockDB.EXPECT().GetStudyList().Return(sid, nil)
 	for i := range sid {
-		mockDB.EXPECT().GetStudyConfig(sid[i]).Return(sc[i], nil)
+		mockDB.EXPECT().GetStudy(sid[i]).Return(sc[i], nil)
 	}
 
 	req := &api.GetStudyListRequest{}

diff --git a/pkg/db/db_init.go b/pkg/db/db_init.go
@@ -7,6 +7,7 @@ import (
 
 func (d *dbConn) DBInit() {
 	db := d.db
+
 	_, err := db.Exec(`CREATE TABLE IF NOT EXISTS studies
 		(id CHAR(16) PRIMARY KEY,
 		name VARCHAR(255),
@@ -17,7 +18,10 @@ func (d *dbConn) DBInit() {
 		tags TEXT,
 		objective_value_name VARCHAR(255),
 		metrics TEXT,
-		job_id TEXT)`)
+		nasconfig TEXT,
+		job_id TEXT,
+		job_type TEXT)`)
+
 	if err != nil {
 		log.Fatalf("Error creating studies table: %v", err)
 	}
@@ -37,6 +41,7 @@ func (d *dbConn) DBInit() {
 		parameters TEXT,
 		objective_value VARCHAR(255),
 		tags TEXT,
+		time DATETIME(6),
 		FOREIGN KEY(study_id) REFERENCES studies(id) ON DELETE CASCADE)`)
 	if err != nil {
 		log.Fatalf("Error creating trials table: %v", err)
@@ -95,6 +100,7 @@ func (d *dbConn) DBInit() {
 	if err != nil {
 		log.Fatalf("Error creating earlystop_param table: %v", err)
 	}
+
 }
 
 func (d *dbConn) SelectOne() error {