Add jsonb backend to Lookout (#3408)

internal/common/database/lookout/jobstates.go

@@ -25,7 +25,6 @@ const (
 	JobPreemptedOrdinal = 7
 	JobLeasedOrdinal    = 8
 
-	JobRunLeased           JobRunState = "RUN_LEASED"
 	JobRunPending          JobRunState = "RUN_PENDING"
 	JobRunRunning          JobRunState = "RUN_RUNNING"
 	JobRunSucceeded        JobRunState = "RUN_SUCCEEDED"
@@ -36,6 +35,7 @@ const (
 	JobRunLeaseReturned    JobRunState = "RUN_LEASE_RETURNED"
 	JobRunLeaseExpired     JobRunState = "RUN_LEASE_EXPIRED"
 	JobRunMaxRunsExceeded  JobRunState = "RUN_MAX_RUNS_EXCEEDED"
+	JobRunLeased           JobRunState = "RUN_LEASED"
 
 	JobRunPendingOrdinal          = 1
 	JobRunRunningOrdinal          = 2
@@ -89,4 +89,6 @@ var (
 		JobRunLeaseExpiredOrdinal:     JobRunLeaseExpired,
 		JobRunMaxRunsExceededOrdinal:  JobRunMaxRunsExceeded,
 	}
+
+	JobRunStateOrdinalMap = util.InverseMap(JobRunStateMap)
 )

internal/lookout/ui/src/index.tsx

@@ -35,8 +35,10 @@ import "./index.css"
   const fakeDataEnabled = uiConfig.fakeDataEnabled
 
   const v2TestJobs = fakeDataEnabled ? makeRandomJobs(10000, 42) : []
-  const v2GetJobsService = fakeDataEnabled ? new FakeGetJobsService(v2TestJobs) : new GetJobsService()
-  const v2GroupJobsService = fakeDataEnabled ? new FakeGroupJobsService(v2TestJobs) : new GroupJobsService()
+  const v2GetJobsService = fakeDataEnabled ? new FakeGetJobsService(v2TestJobs) : new GetJobsService(uiConfig.backend)
+  const v2GroupJobsService = fakeDataEnabled
+    ? new FakeGroupJobsService(v2TestJobs)
+    : new GroupJobsService(uiConfig.backend)
   const v2RunErrorService = fakeDataEnabled ? new FakeGetRunErrorService() : new GetRunErrorService()
   const v2LogService = fakeDataEnabled
     ? new FakeLogService()

internal/lookout/ui/src/services/lookoutV2/GetJobsService.ts

@@ -16,6 +16,12 @@ export type GetJobsResponse = {
 }
 
 export class GetJobsService implements IGetJobsService {
+  private backend: string | undefined
+
+  constructor(backend: string | undefined) {
+    this.backend = backend
+  }
+
   async getJobs(
     filters: JobFilter[],
     activeJobSets: boolean,
@@ -24,7 +30,11 @@ export class GetJobsService implements IGetJobsService {
     take: number,
     abortSignal?: AbortSignal,
   ): Promise<GetJobsResponse> {
-    const response = await fetch("/api/v1/jobs", {
+    let path = "/api/v1/jobs"
+    if (this.backend) {
+      path += "?" + new URLSearchParams({ backend: this.backend })
+    }
+    const response = await fetch(path, {
       method: "POST",
       headers: { "Content-Type": "application/json" },
       body: JSON.stringify({
@@ -36,7 +46,6 @@ export class GetJobsService implements IGetJobsService {
       }),
       signal: abortSignal,
     })
-
     const json = await response.json()
     return {
       jobs: json.jobs ?? [],

internal/lookout/ui/src/services/lookoutV2/GroupJobsService.ts

@@ -23,6 +23,12 @@ export type GroupJobsResponse = {
 }
 
 export class GroupJobsService implements IGroupJobsService {
+  private backend: string | undefined
+
+  constructor(backend: string | undefined) {
+    this.backend = backend
+  }
+
   async groupJobs(
     filters: JobFilter[],
     activeJobSets: boolean,
@@ -33,7 +39,11 @@ export class GroupJobsService implements IGroupJobsService {
     take: number,
     abortSignal?: AbortSignal,
   ): Promise<GroupJobsResponse> {
-    const response = await fetch("/api/v1/jobGroups", {
+    let path = "/api/v1/jobGroups"
+    if (this.backend) {
+      path += "?" + new URLSearchParams({ backend: this.backend })
+    }
+    const response = await fetch(path, {
       method: "POST",
       headers: { "Content-Type": "application/json" },
       body: JSON.stringify({

internal/lookout/ui/src/utils.tsx

@@ -24,6 +24,7 @@ interface UIConfig {
   oidcEnabled: boolean
   oidc?: OidcConfig
   commandSpecs: CommandSpec[]
+  backend: string | undefined
 }
 
 export type RequestStatus = "Loading" | "Idle"
@@ -52,6 +53,7 @@ export async function getUIConfig(): Promise<UIConfig> {
     oidcEnabled: false,
     oidc: undefined,
     commandSpecs: [],
+    backend: undefined,
   }
 
   try {
@@ -76,6 +78,7 @@ export async function getUIConfig(): Promise<UIConfig> {
         })
       }
     }
+    if (json.Backend) config.backend = json.Backend
   } catch (e) {
     console.error(e)
   }
@@ -91,6 +94,9 @@ export async function getUIConfig(): Promise<UIConfig> {
 
   if (window.location.pathname === "/oidc") config.oidcEnabled = true
 
+  const backend = searchParams.get("backend")
+  if (backend) config.backend = backend
+
   return config
 }
 

internal/lookoutv2/application.go

@@ -31,8 +31,10 @@ func Serve(configuration configuration.LookoutV2Config) error {
 		return err
 	}
 
-	getJobsRepo := repository.NewSqlGetJobsRepository(db)
-	groupJobsRepo := repository.NewSqlGroupJobsRepository(db)
+	getJobsRepo := repository.NewSqlGetJobsRepository(db, false)
+	getJobsJsonbRepo := repository.NewSqlGetJobsRepository(db, true)
+	groupJobsRepo := repository.NewSqlGroupJobsRepository(db, false)
+	groupJobsJsonbRepo := repository.NewSqlGroupJobsRepository(db, true)
 	decompressor := compress.NewThreadSafeZlibDecompressor()
 	getJobRunErrorRepo := repository.NewSqlGetJobRunErrorRepository(db, decompressor)
 	getJobSpecRepo := repository.NewSqlGetJobSpecRepository(db, decompressor)
@@ -54,7 +56,11 @@ func Serve(configuration configuration.LookoutV2Config) error {
 		func(params operations.GetJobsParams) middleware.Responder {
 			filters := util.Map(params.GetJobsRequest.Filters, conversions.FromSwaggerFilter)
 			order := conversions.FromSwaggerOrder(params.GetJobsRequest.Order)
-			result, err := getJobsRepo.GetJobs(
+			repo := getJobsRepo
+			if backend := params.Backend; backend != nil && *backend == "jsonb" {
+				repo = getJobsJsonbRepo
+			}
+			result, err := repo.GetJobs(
 				armadacontext.New(params.HTTPRequest.Context(), logger),
 				filters,
 				params.GetJobsRequest.ActiveJobSets,
@@ -75,7 +81,11 @@ func Serve(configuration configuration.LookoutV2Config) error {
 		func(params operations.GroupJobsParams) middleware.Responder {
 			filters := util.Map(params.GroupJobsRequest.Filters, conversions.FromSwaggerFilter)
 			order := conversions.FromSwaggerOrder(params.GroupJobsRequest.Order)
-			result, err := groupJobsRepo.GroupBy(
+			repo := groupJobsRepo
+			if backend := params.Backend; backend != nil && *backend == "jsonb" {
+				repo = groupJobsJsonbRepo
+			}
+			result, err := repo.GroupBy(
 				armadacontext.New(params.HTTPRequest.Context(), logger),
 				filters,
 				params.GroupJobsRequest.ActiveJobSets,

internal/lookoutv2/conversions/convert.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/go-openapi/strfmt"
 
+	"github.com/armadaproject/armada/internal/common/database/lookout"
 	"github.com/armadaproject/armada/internal/lookoutv2/gen/models"
 	"github.com/armadaproject/armada/internal/lookoutv2/gen/restapi/operations"
 	"github.com/armadaproject/armada/internal/lookoutv2/model"
@@ -17,7 +18,7 @@ func ToSwaggerJob(job *model.Job) *models.Job {
 	}
 	return &models.Job{
 		Annotations:        job.Annotations,
-		Cancelled:          toSwaggerTimePtr(job.Cancelled),
+		Cancelled:          ToSwaggerTime(job.Cancelled),
 		CPU:                job.Cpu,
 		Duplicate:          job.Duplicate,
 		EphemeralStorage:   job.EphemeralStorage,
@@ -43,13 +44,13 @@ func ToSwaggerRun(run *model.Run) *models.Run {
 	return &models.Run{
 		Cluster:     run.Cluster,
 		ExitCode:    run.ExitCode,
-		Finished:    toSwaggerTimePtr(run.Finished),
-		JobRunState: run.JobRunState,
+		Finished:    PostgreSQLTimeToSwaggerTime(run.Finished),
+		JobRunState: string(lookout.JobRunStateMap[run.JobRunState]),
 		Node:        run.Node,
-		Leased:      toSwaggerTimePtr(run.Leased),
-		Pending:     toSwaggerTimePtr(run.Pending),
+		Leased:      PostgreSQLTimeToSwaggerTime(run.Leased),
+		Pending:     PostgreSQLTimeToSwaggerTime(run.Pending),
 		RunID:       run.RunId,
-		Started:     toSwaggerTimePtr(run.Started),
+		Started:     PostgreSQLTimeToSwaggerTime(run.Started),
 	}
 }
 
@@ -90,10 +91,18 @@ func FromSwaggerGroupedField(groupedField *operations.GroupJobsParamsBodyGrouped
 	}
 }
 
-func toSwaggerTimePtr(ts *time.Time) *strfmt.DateTime {
-	if ts == nil {
+func ToSwaggerTime(t *time.Time) *strfmt.DateTime {
+	if t == nil {
 		return nil
 	}
-	swaggerTs := strfmt.DateTime(*ts)
-	return &swaggerTs
+	s := strfmt.DateTime(*t)
+	return &s
+}
+
+func PostgreSQLTimeToSwaggerTime(t *model.PostgreSQLTime) *strfmt.DateTime {
+	if t == nil {
+		return nil
+	}
+	s := strfmt.DateTime(t.Time)
+	return &s
 }

internal/lookoutv2/conversions/convert_test.go

@@ -77,13 +77,13 @@ var (
 			{
 				Cluster:     "cluster",
 				ExitCode:    pointer.Int32(322),
-				Finished:    &baseTime,
-				JobRunState: string(lookout.JobRunLeaseReturned),
+				Finished:    model.NewPostgreSQLTime(&baseTime),
+				JobRunState: lookout.JobRunStateOrdinalMap[lookout.JobRunLeaseReturned],
 				Node:        pointer.String("node"),
-				Leased:      &baseTime,
-				Pending:     &baseTime,
+				Leased:      model.NewPostgreSQLTime(&baseTime),
+				Pending:     model.NewPostgreSQLTime(&baseTime),
 				RunId:       "run-id",
-				Started:     &baseTime,
+				Started:     model.NewPostgreSQLTime(&baseTime),
 			},
 		},
 		State:     string(lookout.JobFailed),