Fix: save if in DB

yin-zt · Nov 13, 2023 · e0ed60c · e0ed60c
1 parent 1119180
commit e0ed60c
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/datastore/postgres.go b/datastore/postgres.go
@@ -57,6 +57,7 @@ type taskRecord struct {
 	SubJob      []byte         `db:"subjob"`
 	SubJobID    string         `db:"subjob_id"`
 	GPUs        string         `db:"gpus"`
+	IF          string         `db:"if_"`
 }
 
 type jobRecord struct {
@@ -204,6 +205,7 @@ func (r taskRecord) toTask() (*tork.Task, error) {
 		Description: r.Description,
 		SubJob:      subjob,
 		GPUs:        r.GPUs,
+		If:          r.IF,
 	}, nil
 }
 
@@ -412,12 +414,13 @@ func (ds *PostgresDatastore) CreateTask(ctx context.Context, t *tork.Task) error
 			networks, -- $32
 			files_, -- $33
 			registry, -- $34
-			gpus -- $35
+			gpus, -- $35
+			if_ -- $36
 		  ) 
 	      values (
 			$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,
 		    $15,$16,$17,$18,$19,$20,$21,$22,$23,$24,$25,$26,
-			$27,$28,$29,$30,$31,$32,$33,$34,$35)`
+			$27,$28,$29,$30,$31,$32,$33,$34,$35,$36)`
 	_, err = ds.exec(q,
 		t.ID,                         // $1
 		t.JobID,                      // $2
@@ -454,6 +457,7 @@ func (ds *PostgresDatastore) CreateTask(ctx context.Context, t *tork.Task) error
 		files,                        // $33
 		registry,                     // $34
 		t.GPUs,                       // $35
+		t.If,                         // $36
 	)
 	if err != nil {
 		return errors.Wrapf(err, "error inserting task to the db")

diff --git a/datastore/postgres_test.go b/datastore/postgres_test.go
@@ -36,6 +36,7 @@ func TestPostgresCreateAndGetTask(t *testing.T) {
 		Files:       map[string]string{"myfile": "hello world"},
 		Registry:    &tork.Registry{Username: "me", Password: "secret"},
 		GPUs:        "all",
+		If:          "true",
 	}
 	err = ds.CreateTask(ctx, &t1)
 	assert.NoError(t, err)
@@ -48,6 +49,7 @@ func TestPostgresCreateAndGetTask(t *testing.T) {
 	assert.Equal(t, "me", t2.Registry.Username)
 	assert.Equal(t, "secret", t2.Registry.Password)
 	assert.Equal(t, "all", t2.GPUs)
+	assert.Equal(t, "true", t2.If)
 	assert.Nil(t, t2.Parallel)
 }
 

diff --git a/db/postgres/schema.go b/db/postgres/schema.go
@@ -86,7 +86,8 @@ CREATE TABLE tasks (
     description   text,
     subjob        jsonb,
     networks      text[],
-    gpus          text
+    gpus          text,
+    if_           text
 );
 
 CREATE INDEX idx_tasks_state ON tasks (state);