From 1e0d691452f14aaa949ccdd17bfb1547fbe40b49 Mon Sep 17 00:00:00 2001
From: James Rasell <jrasell@users.noreply.github.com>
Date: Fri, 30 Jun 2023 08:37:22 +0200
Subject: [PATCH] job: ensure node pool is canonicalized for state restores.
 (#17765)

---
 nomad/state/state_store_restore.go      | 11 +++++++
 nomad/state/state_store_restore_test.go | 40 +++++++++++++++----------
 2 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/nomad/state/state_store_restore.go b/nomad/state/state_store_restore.go
index 254b767c7..166694f6d 100644
--- a/nomad/state/state_store_restore.go
+++ b/nomad/state/state_store_restore.go
@@ -44,6 +44,17 @@ func (r *StateRestore) NodePoolRestore(pool *structs.NodePool) error {
 
 // JobRestore is used to restore a job
 func (r *StateRestore) JobRestore(job *structs.Job) error {
+
+	// When upgrading a cluster pre to post 1.6, the existing jobs will not
+	// have a node pool set. Inserting this into the table will fail, as this
+	// is indexed and cannot be empty.
+	//
+	// This cannot happen within the job canonicalize function, as it would
+	// break the node pools and governance feature.
+	if job.NodePool == "" {
+		job.NodePool = structs.NodePoolDefault
+	}
+
 	if err := r.txn.Insert("jobs", job); err != nil {
 		return fmt.Errorf("job insert failed: %v", err)
 	}
diff --git a/nomad/state/state_store_restore_test.go b/nomad/state/state_store_restore_test.go
index d087f17db..01320b165 100644
--- a/nomad/state/state_store_restore_test.go
+++ b/nomad/state/state_store_restore_test.go
@@ -50,28 +50,36 @@ func TestStateStore_RestoreJob(t *testing.T) {
 	ci.Parallel(t)
 
 	state := testStateStore(t)
-	job := mock.Job()
+	mockJob1 := mock.Job()
 
 	restore, err := state.Restore()
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
+	must.NoError(t, err)
 
-	err = restore.JobRestore(job)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
-	require.NoError(t, restore.Commit())
+	err = restore.JobRestore(mockJob1)
+	must.NoError(t, err)
+	must.NoError(t, restore.Commit())
 
 	ws := memdb.NewWatchSet()
-	out, err := state.JobByID(ws, job.Namespace, job.ID)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
+	out, err := state.JobByID(ws, mockJob1.Namespace, mockJob1.ID)
+	must.NoError(t, err)
+	must.Eq(t, mockJob1, out)
 
-	if !reflect.DeepEqual(out, job) {
-		t.Fatalf("Bad: %#v %#v", out, job)
-	}
+	// Test upgrade to 1.6 or greater to simulate restoring a job which does
+	// not have a node pool set.
+	mockJob2 := mock.Job()
+	mockJob2.NodePool = ""
+
+	restore, err = state.Restore()
+	must.NoError(t, err)
+
+	err = restore.JobRestore(mockJob2)
+	must.NoError(t, err)
+	must.NoError(t, restore.Commit())
+
+	ws = memdb.NewWatchSet()
+	out, err = state.JobByID(ws, mockJob2.Namespace, mockJob2.ID)
+	must.NoError(t, err)
+	must.Eq(t, structs.NodePoolDefault, out.NodePool)
 }
 
 func TestStateStore_RestorePeriodicLaunch(t *testing.T) {