wip

Signed-off-by: Chris Martin <[email protected]>
armadaproject · Feb 21, 2024 · a4c854a · a4c854a
1 parent 10fe780
commit a4c854a
Show file tree

Hide file tree

Showing 10 changed files with 230 additions and 901 deletions.
diff --git a/internal/armada/queryapi/query_api.go b/internal/armada/queryapi/query_api.go
@@ -12,15 +12,15 @@ import (
 )
 
 // JobStateMap is a mapping between database state and api states
-var JobStateMap = map[int16]api.JobStatus{
-	lookout.JobLeasedOrdinal:    api.JobStatus_LEASED,
-	lookout.JobQueuedOrdinal:    api.JobStatus_QUEUED,
-	lookout.JobPendingOrdinal:   api.JobStatus_PENDING,
-	lookout.JobRunningOrdinal:   api.JobStatus_RUNNING,
-	lookout.JobSucceededOrdinal: api.JobStatus_SUCCEEDED,
-	lookout.JobFailedOrdinal:    api.JobStatus_FAILED,
-	lookout.JobCancelledOrdinal: api.JobStatus_CANCELLED,
-	lookout.JobPreemptedOrdinal: api.JobStatus_PREEMPTED,
+var JobStateMap = map[int16]api.JobState{
+	lookout.JobLeasedOrdinal:    api.JobState_LEASED,
+	lookout.JobQueuedOrdinal:    api.JobState_QUEUED,
+	lookout.JobPendingOrdinal:   api.JobState_PENDING,
+	lookout.JobRunningOrdinal:   api.JobState_RUNNING,
+	lookout.JobSucceededOrdinal: api.JobState_SUCCEEDED,
+	lookout.JobFailedOrdinal:    api.JobState_FAILED,
+	lookout.JobCancelledOrdinal: api.JobState_CANCELLED,
+	lookout.JobPreemptedOrdinal: api.JobState_PREEMPTED,
 }
 
 type QueryApi struct {
@@ -43,10 +43,15 @@ func (q *QueryApi) GetJobStatus(ctx context.Context, req *api.JobStatusRequest)
 	}
 	apiStatus, ok := JobStateMap[status]
 	if !ok {
-		apiStatus = api.JobStatus_UNKNOWN
+		apiStatus = api.JobState_UNKNOWN
 	}
 	return &api.JobStatusResponse{
-		JobId:     req.JobId,
-		JobStatus: apiStatus,
+		JobId:    req.JobId,
+		JobState: apiStatus,
 	}, nil
 }
+
+func (q *QueryApi) GetJob(ctx context.Context, request *api.JobRequest) (*api.JobResponse, error) {
+	// TODO implement me
+	panic("implement me")
+}
diff --git a/internal/armada/queryapi/query_api_test.go b/internal/armada/queryapi/query_api_test.go
@@ -24,19 +24,19 @@ func TestGetJobStatus(t *testing.T) {
 	}{
 		"leased job": {
 			jobId:            "leasedJob",
-			expectedResponse: &api.JobStatusResponse{JobId: "leasedJob", JobStatus: api.JobStatus_LEASED},
+			expectedResponse: &api.JobStatusResponse{JobId: "leasedJob", JobState: api.JobState_LEASED},
 		},
 		"running job": {
 			jobId:            "runningJob",
-			expectedResponse: &api.JobStatusResponse{JobId: "runningJob", JobStatus: api.JobStatus_RUNNING},
+			expectedResponse: &api.JobStatusResponse{JobId: "runningJob", JobState: api.JobState_RUNNING},
 		},
 		"completed job": {
 			jobId:            "completedJob",
-			expectedResponse: &api.JobStatusResponse{JobId: "completedJob", JobStatus: api.JobStatus_SUCCEEDED},
+			expectedResponse: &api.JobStatusResponse{JobId: "completedJob", JobState: api.JobState_SUCCEEDED},
 		},
 		"missing job": {
 			jobId:            "missingJob",
-			expectedResponse: &api.JobStatusResponse{JobId: "missingJob", JobStatus: api.JobStatus_UNKNOWN},
+			expectedResponse: &api.JobStatusResponse{JobId: "missingJob", JobState: api.JobState_UNKNOWN},
 		},
 	}
 	for name, tc := range tests {

diff --git a/internal/armada/server.go b/internal/armada/server.go
@@ -263,7 +263,7 @@ func Serve(ctx *armadacontext.Context, config *configuration.ArmadaConfig, healt
 			return errors.WithMessage(err, "error creating QueryApi postgres pool")
 		}
 		queryapiServer := queryapi.New(queryDb)
-		api.RegisterQueryApiServer(grpcServer, queryapiServer)
+		api.RegisterJobsServer(grpcServer, queryapiServer)
 	}
 
 	api.RegisterSubmitServer(grpcServer, pulsarSubmitServer)

diff --git a/pkg/api/api.swagger.go b/pkg/api/api.swagger.go
@@ -1296,7 +1296,11 @@ func SwaggerJsonTemplate() string {
 		"        \"RUNNING\",\n" +
 		"        \"SUCCEEDED\",\n" +
 		"        \"FAILED\",\n" +
-		"        \"UNKNOWN\"\n" +
+		"        \"UNKNOWN\",\n" +
+		"        \"SUBMITTED\",\n" +
+		"        \"LEASED\",\n" +
+		"        \"PREEMPTED\",\n" +
+		"        \"CANCELLED\"\n" +
 		"      ]\n" +
 		"    },\n" +
 		"    \"apiJobSubmitRequest\": {\n" +

diff --git a/pkg/api/api.swagger.json b/pkg/api/api.swagger.json
@@ -1285,7 +1285,11 @@
         "RUNNING",
         "SUCCEEDED",
         "FAILED",
-        "UNKNOWN"
+        "UNKNOWN",
+        "SUBMITTED",
+        "LEASED",
+        "PREEMPTED",
+        "CANCELLED"
       ]
     },
     "apiJobSubmitRequest": {

diff --git a/pkg/api/queryapi.proto → pkg/api/job.proto b/pkg/api/queryapi.proto → pkg/api/job.proto
@@ -1,23 +1,19 @@
 syntax = 'proto3';
-
-package queryapi;
+package api;
 option go_package = "github.com/armadaproject/armada/pkg/api";
 
 import "github.com/gogo/protobuf/gogoproto/gogo.proto";
+import "pkg/api/submit.proto";
 option (gogoproto.goproto_stringer_all) = false;
 option (gogoproto.stringer_all) = true;
 
-enum JobStatus {
-  UNKNOWN = 0;
-  SUBMITTED = 1;
-  QUEUED = 2;
-  LEASED = 3;
-  PENDING = 4;
-  RUNNING = 5;
-  SUCCEEDED = 6;
-  FAILED = 7;
-  PREEMPTED = 8;
-  CANCELLED = 9;
+
+message JobRequest {
+  string job_id = 1;
+}
+
+message JobResponse{
+  Job job = 1;
 }
 
 message JobStatusRequest{
@@ -26,9 +22,10 @@ message JobStatusRequest{
 
 message JobStatusResponse{
   string job_id = 1;
-  JobStatus job_status = 2;
+  JobState job_state = 2;
 }
 
-service QueryApi {
+service Jobs {
   rpc GetJobStatus (JobStatusRequest) returns (JobStatusResponse);
+  rpc GetJob (JobRequest) returns (JobResponse);
 }