feat(job-run-execution-summary-api): allow unfinished jobs in execution summary (#589)

* add job run summary get state & is finished

* allow unfinished job to appear in job run execution summary

* readd number of upstreams per level

* add truncated & total nodes in response, do not deprecate number of upstream per level

* update makefile

Author: ahmadnaufal

Makefile

@@ -5,7 +5,7 @@ NAME = "github.com/goto/optimus"
 LAST_COMMIT := $(shell git rev-parse --short HEAD)
 LAST_TAG := "$(shell git rev-list --tags --max-count=1)"
 OPMS_VERSION := "$(shell git describe --tags ${LAST_TAG})-next"
-PROTON_COMMIT := "bba736bdc93dabd74e2ce2ec155cee99692c0268"
+PROTON_COMMIT := "ee54e908094ca84e49c9d5f171ce62f3173281ae"
 
 
 .PHONY: build test test-ci generate-proto unit-test-ci integration-test vet coverage clean install lint

config/config_server.go

@@ -199,6 +199,7 @@ type JobExpectatorConfig struct {
 
 type JobExecutionSummaryConfig struct {
 	MaxLineageDepth    int                      `mapstructure:"max_lineage_depth" default:"25"`
+	LineageWindowHours int                      `mapstructure:"lineage_window_hours" default:"24"`
 	HistoricalDuration HistoricalDurationConfig `mapstructure:"historical_duration"`
 }
 

config/loader_test.go

@@ -322,7 +322,8 @@ func (s *ConfigTestSuite) initExpectedServerConfig() {
 		},
 	}
 	s.expectedServerConfig.JobExecutionSummaryConfig = config.JobExecutionSummaryConfig{
-		MaxLineageDepth: 25,
+		MaxLineageDepth:    25,
+		LineageWindowHours: 24,
 		HistoricalDuration: config.HistoricalDurationConfig{
 			LastNRuns:  7,
 			Percentile: 95,

core/scheduler/handler/v1beta1/job_run.go

@@ -68,7 +68,7 @@ type JobRunService interface {
 }
 
 type JobLineageService interface {
-	GetJobExecutionSummary(ctx context.Context, jobSchedules []*scheduler.JobSchedule, numberOfUpstreamPerLevel int) ([]*scheduler.JobRunLineage, error)
+	GetJobExecutionSummary(ctx context.Context, jobSchedules []*scheduler.JobSchedule, opts scheduler.LineageSummaryOptions) ([]*scheduler.JobRunLineage, error)
 }
 
 type ThirdPartySensorService interface {
@@ -590,7 +590,11 @@ func (h JobRunHandler) GetJobRunLineageSummary(ctx context.Context, req *pb.GetJ
 		h.l.Error("error parsing job schedules from request: %s", err)
 		return nil, errors.GRPCErr(err, "unable to parse job schedules from request")
 	}
-	jobRunLineages, err := h.jobLineageService.GetJobExecutionSummary(ctx, targetJobSchedules, int(req.GetNumberOfUpstreamPerLevel()))
+	jobRunLineages, err := h.jobLineageService.GetJobExecutionSummary(ctx, targetJobSchedules, scheduler.LineageSummaryOptions{
+		MaxNodes:           int(req.GetMaxNodes()),
+		TopUpstreamsPerJob: int(req.GetNumberOfUpstreamPerLevel()),
+		WindowHours:        int(req.GetLineageWindowHours()),
+	})
 	if err != nil {
 		h.l.Error("error getting job run lineage summary: %s", err)
 		return nil, errors.GRPCErr(err, "unable to get job run lineage summary")

core/scheduler/handler/v1beta1/job_run_adapter.go

@@ -14,6 +14,14 @@ import (
 	pb "github.com/goto/optimus/protos/gotocompany/optimus/core/v1beta1"
 )
 
+var jobRunStateToProto = map[scheduler.State]pb.JobRunState{
+	scheduler.StateNotScheduled: pb.JobRunState_JOB_RUN_STATE_NOT_SCHEDULED,
+	scheduler.StateWaitUpstream: pb.JobRunState_JOB_RUN_STATE_WAIT_UPSTREAM,
+	scheduler.StateRunning:      pb.JobRunState_JOB_RUN_STATE_RUNNING,
+	scheduler.StateSuccess:      pb.JobRunState_JOB_RUN_STATE_SUCCESS,
+	scheduler.StateFailed:       pb.JobRunState_JOB_RUN_STATE_FAILED,
+}
+
 // buildIdentifySLABreachInputs resolves an IdentifyPotentialSLABreachRequest into
 // the combos (cross product of projects x label groups) to evaluate and the
 // shared predictor config. job_names only applies in legacy single-project mode,
@@ -137,6 +145,30 @@ func fromJobRunLineageSummaryRequest(req *pb.GetJobRunLineageSummaryRequest) ([]
 	return targetJobSchedules, nil
 }
 
+func toJobRunStateProto(state scheduler.State) pb.JobRunState {
+	if pbState, ok := jobRunStateToProto[state]; ok {
+		return pbState
+	}
+
+	return pb.JobRunState_JOB_RUN_STATE_UNSPECIFIED
+}
+
+func toJobRunRefsProto(refs []scheduler.JobRunKey) []*pb.JobRunRef {
+	if len(refs) == 0 {
+		return nil
+	}
+
+	pbRefs := make([]*pb.JobRunRef, 0, len(refs))
+	for _, ref := range refs {
+		pbRefs = append(pbRefs, &pb.JobRunRef{
+			JobName:     ref.JobName.String(),
+			ScheduledAt: timestamppb.New(ref.ScheduledAt),
+		})
+	}
+
+	return pbRefs
+}
+
 func toJobRunLineageSummaryResponse(jobRunLineages []*scheduler.JobRunLineage) *pb.GetJobRunLineageSummaryResponse {
 	var pbJobRunLineages []*pb.JobRunLineageSummary
 	for _, lineage := range jobRunLineages {
@@ -224,6 +256,9 @@ func toJobRunLineageSummaryResponse(jobRunLineages []*scheduler.JobRunLineage) *
 				},
 				Level:              int32(run.Level),
 				DownstreamPathName: run.DownstreamPathName,
+				RunState:           toJobRunStateProto(run.State),
+				IsBlocking:         run.IsBlocking,
+				DownstreamRefs:     toJobRunRefsProto(run.DownstreamRefs),
 			}
 			if run.DelaySummary != nil {
 				pbJobRun.DelaySummary = &pb.JobRunDelaySummary{
@@ -261,6 +296,8 @@ func toJobRunLineageSummaryResponse(jobRunLineages []*scheduler.JobRunLineage) *
 			JobName:     lineage.JobName.String(),
 			ScheduledAt: timestamppb.New(lineage.JobRuns[0].JobRunSummary.ScheduledAt),
 			JobRuns:     pbJobRuns,
+			TotalNodes:  int32(lineage.TotalNodes),
+			Truncated:   lineage.Truncated,
 			ExecutionSummary: &pb.LineageExecutionSummary{
 				TotalScheduledWayTooLateSeconds:     int32(lineage.ExecutionSummary.TotalScheduledWayTooLateSeconds),
 				TotalSystemSchedulingDelaySeconds:   int32(lineage.ExecutionSummary.TotalSystemSchedulingDelaySeconds),

core/scheduler/handler/v1beta1/job_run_test.go

@@ -991,7 +991,7 @@ func TestJobRunHandler(t *testing.T) {
 				NumberOfUpstreamPerLevel: 5,
 			}
 
-			jobLineageService.On("GetJobExecutionSummary", ctx, mock.Anything, 5).
+			jobLineageService.On("GetJobExecutionSummary", ctx, mock.Anything, scheduler.LineageSummaryOptions{TopUpstreamsPerJob: 5}).
 				Return(nil, errors.New("service error"))
 
 			resp, err := handler.GetJobRunLineageSummary(ctx, req)
@@ -1036,14 +1036,83 @@ func TestJobRunHandler(t *testing.T) {
 				},
 			}
 
-			jobLineageService.On("GetJobExecutionSummary", ctx, mock.Anything, 5).
+			jobLineageService.On("GetJobExecutionSummary", ctx, mock.Anything, scheduler.LineageSummaryOptions{TopUpstreamsPerJob: 5}).
 				Return(mockLineages, nil)
 
 			resp, err := handler.GetJobRunLineageSummary(ctx, req)
 			assert.Nil(t, err)
 			assert.NotNil(t, resp)
 			assert.Equal(t, len(mockLineages), len(resp.Jobs))
 		})
+
+		t.Run("should prefer max_nodes over the deprecated per-level field and pass the window", func(t *testing.T) {
+			jobRunService := new(mockJobRunService)
+			jobLineageService := new(mockJobLineageService)
+			defer jobLineageService.AssertExpectations(t)
+
+			handler := v1beta1.NewJobRunHandler(logger, jobRunService, nil, nil, jobLineageService, nil, nil, nil)
+
+			req := &pb.GetJobRunLineageSummaryRequest{
+				TargetJobs:               []*pb.TargetJobRunIdentifier{{JobName: "test-job", ScheduledAt: timestamppb.Now()}},
+				NumberOfUpstreamPerLevel: 5,
+				MaxNodes:                 40,
+				LineageWindowHours:       10,
+			}
+
+			jobLineageService.On("GetJobExecutionSummary", ctx, mock.Anything, scheduler.LineageSummaryOptions{MaxNodes: 40, TopUpstreamsPerJob: 5, WindowHours: 10}).
+				Return([]*scheduler.JobRunLineage{}, nil)
+
+			_, err := handler.GetJobRunLineageSummary(ctx, req)
+			assert.Nil(t, err)
+		})
+
+		t.Run("should map run state, blocking flag and downstream refs onto the response", func(t *testing.T) {
+			jobRunService := new(mockJobRunService)
+			jobLineageService := new(mockJobLineageService)
+			defer jobLineageService.AssertExpectations(t)
+
+			handler := v1beta1.NewJobRunHandler(logger, jobRunService, nil, nil, jobLineageService, nil, nil, nil)
+
+			scheduledAt := timestamppb.Now()
+			upstreamScheduledAt := scheduledAt.AsTime().Add(-1 * time.Hour)
+			req := &pb.GetJobRunLineageSummaryRequest{
+				TargetJobs: []*pb.TargetJobRunIdentifier{{JobName: "test-job", ScheduledAt: scheduledAt}},
+			}
+
+			mockLineages := []*scheduler.JobRunLineage{
+				{
+					JobName:     "test-job",
+					ScheduledAt: scheduledAt.AsTime(),
+					JobRuns: []*scheduler.JobExecutionSummary{
+						{
+							JobName:    "upstream-job",
+							State:      scheduler.StateRunning,
+							IsBlocking: true,
+							DownstreamRefs: []scheduler.JobRunKey{
+								{JobName: "test-job", ScheduledAt: upstreamScheduledAt},
+							},
+							JobRunSummary: &scheduler.JobRunSummary{
+								JobName:     "upstream-job",
+								ScheduledAt: upstreamScheduledAt,
+							},
+						},
+					},
+					ExecutionSummary: &scheduler.LineageExecutionSummary{},
+				},
+			}
+
+			jobLineageService.On("GetJobExecutionSummary", ctx, mock.Anything, scheduler.LineageSummaryOptions{}).Return(mockLineages, nil)
+
+			resp, err := handler.GetJobRunLineageSummary(ctx, req)
+			assert.Nil(t, err)
+
+			pbRun := resp.Jobs[0].JobRuns[0]
+			assert.Equal(t, pb.JobRunState_JOB_RUN_STATE_RUNNING, pbRun.RunState)
+			assert.True(t, pbRun.IsBlocking)
+			assert.Len(t, pbRun.DownstreamRefs, 1)
+			assert.Equal(t, "test-job", pbRun.DownstreamRefs[0].JobName)
+			assert.Equal(t, upstreamScheduledAt.UTC(), pbRun.DownstreamRefs[0].ScheduledAt.AsTime().UTC())
+		})
 	})
 
 	t.Run("GenerateExpectedFinishTime", func(t *testing.T) {
@@ -1135,8 +1204,8 @@ type mockJobLineageService struct {
 	mock.Mock
 }
 
-func (m *mockJobLineageService) GetJobExecutionSummary(ctx context.Context, jobSchedules []*scheduler.JobSchedule, numberOfUpstreamPerLevel int) ([]*scheduler.JobRunLineage, error) {
-	args := m.Called(ctx, jobSchedules, numberOfUpstreamPerLevel)
+func (m *mockJobLineageService) GetJobExecutionSummary(ctx context.Context, jobSchedules []*scheduler.JobSchedule, opts scheduler.LineageSummaryOptions) ([]*scheduler.JobRunLineage, error) {
+	args := m.Called(ctx, jobSchedules, opts)
 	if args.Get(0) == nil {
 		return nil, args.Error(1)
 	}