ARMADA-3283 - Implementing GetActiveQueues on QueryAPI (#393)

* Implementing GetActiveQueuesByPool

* Generating code for proto change

* Changing code to respect proto change

* Removing unused imports

* Making index creation concurrent

* Removing dependency exclusion

Author: mustafai-gr

internal/lookout/schema/migrations/016_add_job_run_pool.sql

@@ -0,0 +1,3 @@
+ALTER TABLE job_run ADD COLUMN pool TEXT;
+
+CREATE INDEX CONCURRENTLY idx_job_run_state_pool ON job_run (job_run_state, pool);

internal/lookoutingester/instructions/instructions.go

@@ -337,6 +337,7 @@ func (c *InstructionConverter) handleJobRunLeased(ts time.Time, event *armadaeve
 		JobId:       event.JobId,
 		Cluster:     util.Truncate(event.ExecutorId, maxClusterLen),
 		Node:        pointer.String(util.Truncate(event.NodeId, maxNodeLen)),
+		Pool:        event.Pool,
 		Leased:      &ts,
 		JobRunState: lookout.JobRunLeasedOrdinal,
 	}

internal/lookoutingester/instructions/instructions_test.go

@@ -65,6 +65,7 @@ var expectedLeasedRun = model.CreateJobRunInstruction{
 	Cluster:     testfixtures.ExecutorId,
 	Leased:      &testfixtures.BaseTime,
 	Node:        pointer.String(testfixtures.NodeName),
+	Pool:        testfixtures.Pool,
 	JobRunState: lookout.JobRunLeasedOrdinal,
 }
 

internal/lookoutingester/lookoutdb/insertion.go

@@ -569,7 +569,8 @@ func (l *LookoutDb) CreateJobRunsBatch(ctx *armadacontext.Context, instructions
 					node          varchar(512),
 					leased        timestamp,
 					pending       timestamp,
-					job_run_state smallint
+					job_run_state smallint,
+					pool 		  text
 				) ON COMMIT DROP;`, tmpTable))
 			if err != nil {
 				l.metrics.RecordDBError(commonmetrics.DBOperationCreateTempTable)
@@ -588,6 +589,7 @@ func (l *LookoutDb) CreateJobRunsBatch(ctx *armadacontext.Context, instructions
 					"leased",
 					"pending",
 					"job_run_state",
+					"pool",
 				},
 				pgx.CopyFromSlice(len(instructions), func(i int) ([]interface{}, error) {
 					return []interface{}{
@@ -598,6 +600,7 @@ func (l *LookoutDb) CreateJobRunsBatch(ctx *armadacontext.Context, instructions
 						instructions[i].Leased,
 						instructions[i].Pending,
 						instructions[i].JobRunState,
+						instructions[i].Pool,
 					}, nil
 				}),
 			)
@@ -615,7 +618,8 @@ func (l *LookoutDb) CreateJobRunsBatch(ctx *armadacontext.Context, instructions
 						node,
 						leased,
 						pending,
-						job_run_state
+						job_run_state,
+						pool
 					) SELECT * from %s
 					ON CONFLICT DO NOTHING`, tmpTable))
 			if err != nil {
@@ -635,8 +639,9 @@ func (l *LookoutDb) CreateJobRunsScalar(ctx *armadacontext.Context, instructions
 			node,
 			leased,
 			pending,
-			job_run_state)
-		VALUES ($1, $2, $3, $4, $5, $6, $7)
+			job_run_state,
+			pool)
+		VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
 		ON CONFLICT DO NOTHING`
 	for _, i := range instructions {
 		err := l.withDatabaseRetryInsert(func() error {
@@ -647,7 +652,8 @@ func (l *LookoutDb) CreateJobRunsScalar(ctx *armadacontext.Context, instructions
 				i.Node,
 				i.Leased,
 				i.Pending,
-				i.JobRunState)
+				i.JobRunState,
+				i.Pool)
 			if err != nil {
 				l.metrics.RecordDBError(commonmetrics.DBOperationInsert)
 			}

internal/lookoutingester/model/model.go

@@ -47,6 +47,7 @@ type CreateJobRunInstruction struct {
 	RunId       string
 	JobId       string
 	Cluster     string
+	Pool        string
 	Node        *string
 	Leased      *time.Time
 	Pending     *time.Time

internal/server/queryapi/database/models.go

@@ -22,3 +22,6 @@ SELECT job_id, state FROM job
 WHERE queue=sqlc.arg(queue)::text
   AND jobset = sqlc.arg(jobset)::text
   AND external_job_uri = sqlc.arg(external_job_uri)::text;
+
+-- name: GetActiveQueuesByPool :many
+SELECT DISTINCT jr.pool, j.queue FROM job j JOIN job_run jr ON j.job_id = jr.job_id WHERE j.state IN (2, 3, 8) AND jr.job_run_state IN (1, 2, 11) ORDER BY jr.pool, j.queue;

internal/server/queryapi/database/query.sql

@@ -242,6 +242,32 @@ func (q *QueryApi) GetJobStatusUsingExternalJobUri(ctx context.Context, req *api
 	}, nil
 }
 
+func (q *QueryApi) GetActiveQueues(ctx context.Context, _ *api.GetActiveQueuesRequest) (*api.GetActiveQueuesResponse, error) {
+	queries := database.New(q.db)
+	activeQueues, err := queries.GetActiveQueuesByPool(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("could not get active queues by pool")
+	}
+	queuesByPool := map[string]*api.ActiveQueues{}
+	for _, result := range activeQueues {
+		if result.Pool == nil {
+			continue
+		}
+
+		pool := *result.Pool
+		if _, ok := queuesByPool[pool]; !ok {
+			queuesByPool[pool] = &api.ActiveQueues{
+				Queues: []string{},
+			}
+		}
+		queuesByPool[pool].Queues = append(queuesByPool[pool].Queues, result.Queue)
+	}
+
+	return &api.GetActiveQueuesResponse{
+		ActiveQueuesByPool: queuesByPool,
+	}, nil
+}
+
 func parseDbJobStateToApi(dbStatus int16) api.JobState {
 	apiStatus, ok := JobStateMap[dbStatus]
 	if !ok {

internal/server/queryapi/database/query.sql.go

@@ -585,6 +585,38 @@ func SwaggerJsonTemplate() string {
 		"        }\n" +
 		"      }\n" +
 		"    },\n" +
+		"    \"/v1/queues/active\": {\n" +
+		"      \"post\": {\n" +
+		"        \"tags\": [\n" +
+		"          \"Jobs\"\n" +
+		"        ],\n" +
+		"        \"operationId\": \"GetActiveQueues\",\n" +
+		"        \"parameters\": [\n" +
+		"          {\n" +
+		"            \"name\": \"body\",\n" +
+		"            \"in\": \"body\",\n" +
+		"            \"required\": true,\n" +
+		"            \"schema\": {\n" +
+		"              \"$ref\": \"#/definitions/apiGetActiveQueuesRequest\"\n" +
+		"            }\n" +
+		"          }\n" +
+		"        ],\n" +
+		"        \"responses\": {\n" +
+		"          \"200\": {\n" +
+		"            \"description\": \"A successful response.\",\n" +
+		"            \"schema\": {\n" +
+		"              \"$ref\": \"#/definitions/apiGetActiveQueuesResponse\"\n" +
+		"            }\n" +
+		"          },\n" +
+		"          \"default\": {\n" +
+		"            \"description\": \"An unexpected error response.\",\n" +
+		"            \"schema\": {\n" +
+		"              \"$ref\": \"#/definitions/runtimeError\"\n" +
+		"            }\n" +
+		"          }\n" +
+		"        }\n" +
+		"      }\n" +
+		"    },\n" +
 		"    \"/v1/run/details\": {\n" +
 		"      \"post\": {\n" +
 		"        \"tags\": [\n" +
@@ -647,6 +679,17 @@ func SwaggerJsonTemplate() string {
 		"        }\n" +
 		"      }\n" +
 		"    },\n" +
+		"    \"apiActiveQueues\": {\n" +
+		"      \"type\": \"object\",\n" +
+		"      \"properties\": {\n" +
+		"        \"queues\": {\n" +
+		"          \"type\": \"array\",\n" +
+		"          \"items\": {\n" +
+		"            \"type\": \"string\"\n" +
+		"          }\n" +
+		"        }\n" +
+		"      }\n" +
+		"    },\n" +
 		"    \"apiBatchQueueCreateResponse\": {\n" +
 		"      \"type\": \"object\",\n" +
 		"      \"properties\": {\n" +
@@ -788,6 +831,20 @@ func SwaggerJsonTemplate() string {
 		"        }\n" +
 		"      }\n" +
 		"    },\n" +
+		"    \"apiGetActiveQueuesRequest\": {\n" +
+		"      \"type\": \"object\"\n" +
+		"    },\n" +
+		"    \"apiGetActiveQueuesResponse\": {\n" +
+		"      \"type\": \"object\",\n" +
+		"      \"properties\": {\n" +
+		"        \"activeQueuesByPool\": {\n" +
+		"          \"type\": \"object\",\n" +
+		"          \"additionalProperties\": {\n" +
+		"            \"$ref\": \"#/definitions/apiActiveQueues\"\n" +
+		"          }\n" +
+		"        }\n" +
+		"      }\n" +
+		"    },\n" +
 		"    \"apiIngressConfig\": {\n" +
 		"      \"type\": \"object\",\n" +
 		"      \"properties\": {\n" +