diff --git a/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path.go b/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path.go index 748da047..6ddc1428 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path.go +++ b/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path.go @@ -19,25 +19,19 @@ import ( "strings" "github.com/GoogleCloudPlatform/khi/pkg/common/khictx" - pb "github.com/GoogleCloudPlatform/khi/pkg/generated/khifile/v6" khifilev6 "github.com/GoogleCloudPlatform/khi/pkg/model/khifile/v6" - googlecloudcommon_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudcommon/contract" inspectioncore_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/inspectioncore/contract" ) -// MustComposerEnvironmentTimeline returns the timeline path for a Composer Environment. -func MustComposerEnvironmentTimeline(ctx context.Context, projectID, environmentName string) *khifilev6.TimelinePath { - if projectID == "" { - projectID = "unknown" - } +// MustAirflowTimeline returns the root timeline path for an Airflow environment. +func MustAirflowTimeline(ctx context.Context, environmentName string) *khifilev6.TimelinePath { if environmentName == "" { environmentName = "unknown" } builder := khictx.MustGetValue(ctx, inspectioncore_contract.Builder) - projectPath := googlecloudcommon_contract.MustGCPProjectTimeline(ctx, projectID) - return builder.TimelineAccumulator.GetPath(projectPath, khifilev6.PathSegment{ + return builder.TimelineAccumulator.GetPath(nil, khifilev6.PathSegment{ Name: environmentName, - Type: TimelineTypeComposerEnvironment, + Type: TimelineTypeAirflow, }) } @@ -101,7 +95,7 @@ func MustAirflowComponentsRootTimeline(ctx context.Context, envPath *khifilev6.T } // MustAirflowComponentTimeline returns the timeline path for a specific Airflow component. -func MustAirflowComponentTimeline(ctx context.Context, envPath *khifilev6.TimelinePath, componentType *pb.TimelineType, name string) *khifilev6.TimelinePath { +func MustAirflowComponentTimeline(ctx context.Context, envPath *khifilev6.TimelinePath, name string) *khifilev6.TimelinePath { if name == "" { name = "unknown" } @@ -109,24 +103,16 @@ func MustAirflowComponentTimeline(ctx context.Context, envPath *khifilev6.Timeli builder := khictx.MustGetValue(ctx, inspectioncore_contract.Builder) return builder.TimelineAccumulator.GetPath(compRoot, khifilev6.PathSegment{ Name: name, - Type: componentType, + Type: TimelineTypeAirflowComponent, }) } -// MustAirflowWorkerTimeline returns the timeline path for an Airflow worker. -func MustAirflowWorkerTimeline(ctx context.Context, envPath *khifilev6.TimelinePath, workerHost string) *khifilev6.TimelinePath { - if workerHost == "" { - workerHost = "unknown" - } - return MustAirflowComponentTimeline(ctx, envPath, TimelineTypeAirflowWorker, workerHost) -} - -// MustAirflowDAGProcessorManagerRootTimeline returns the root timeline path for DAG Processor Manager. -func MustAirflowDAGProcessorManagerRootTimeline(ctx context.Context, envPath *khifilev6.TimelinePath) *khifilev6.TimelinePath { +// MustAirflowDAGFilesTimeline returns the root timeline path for DAG files. +func MustAirflowDAGFilesTimeline(ctx context.Context, envPath *khifilev6.TimelinePath) *khifilev6.TimelinePath { builder := khictx.MustGetValue(ctx, inspectioncore_contract.Builder) return builder.TimelineAccumulator.GetPath(envPath, khifilev6.PathSegment{ - Name: "DAG Processor Manager", - Type: TimelineTypeDAGProcessorManager, + Name: "DAG files", + Type: TimelineTypeDAGFiles, }) } @@ -135,7 +121,7 @@ func MustAirflowDAGFileTimeline(ctx context.Context, envPath *khifilev6.Timeline if filePath == "" { filePath = "unknown" } - dpmRoot := MustAirflowDAGProcessorManagerRootTimeline(ctx, envPath) + dpmRoot := MustAirflowDAGFilesTimeline(ctx, envPath) builder := khictx.MustGetValue(ctx, inspectioncore_contract.Builder) trimmedPath := strings.TrimPrefix(filePath, "/home/airflow/gcs/dags/") if trimmedPath == "" { diff --git a/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path_test.go b/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path_test.go index 51021923..c875cafc 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path_test.go +++ b/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_path_test.go @@ -23,61 +23,41 @@ import ( "github.com/google/go-cmp/cmp" ) -func TestMustComposerEnvironmentTimeline(t *testing.T) { +func TestMustAirflowTimeline(t *testing.T) { builder := khifilev6.NewBuilder() ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) testCases := []struct { name string - projectID string environmentName string - wantProject string wantEnv string }{ { name: "valid input", - projectID: "my-project", environmentName: "my-env", - wantProject: "my-project", - wantEnv: "my-env", - }, - { - name: "empty project", - projectID: "", - environmentName: "my-env", - wantProject: "unknown", wantEnv: "my-env", }, { name: "empty env", - projectID: "my-project", - environmentName: "", - wantProject: "my-project", - wantEnv: "unknown", - }, - { - name: "both empty", - projectID: "", environmentName: "", - wantProject: "unknown", wantEnv: "unknown", }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - got := MustComposerEnvironmentTimeline(ctx, tc.projectID, tc.environmentName) + got := MustAirflowTimeline(ctx, tc.environmentName) if got == nil { t.Fatal("expected timeline path to be not nil") } if diff := cmp.Diff(tc.wantEnv, got.Name.Resolve()); diff != "" { - t.Errorf("MustComposerEnvironmentTimeline() environmentName mismatch (-want +got):\n%s", diff) + t.Errorf("MustAirflowTimeline() environmentName mismatch (-want +got):\n%s", diff) } - if got.Parent == nil { - t.Fatal("expected parent timeline path to be not nil") + if got.Parent != nil { + t.Errorf("MustAirflowTimeline() expected root timeline (nil parent), got %v", got.Parent) } - if diff := cmp.Diff(tc.wantProject, got.Parent.Name.Resolve()); diff != "" { - t.Errorf("MustComposerEnvironmentTimeline() projectID mismatch (-want +got):\n%s", diff) + if diff := cmp.Diff(TimelineTypeAirflow.GetId(), got.Type.GetId()); diff != "" { + t.Errorf("MustAirflowTimeline() type ID mismatch (-want +got):\n%s", diff) } }) } @@ -86,7 +66,7 @@ func TestMustComposerEnvironmentTimeline(t *testing.T) { func TestMustAirflowDAGTimeline(t *testing.T) { builder := khifilev6.NewBuilder() ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) - envPath := MustComposerEnvironmentTimeline(ctx, "my-project", "my-env") + envPath := MustAirflowTimeline(ctx, "my-env") testCases := []struct { name string @@ -121,7 +101,7 @@ func TestMustAirflowDAGTimeline(t *testing.T) { func TestMustAirflowDAGRunTimeline(t *testing.T) { builder := khifilev6.NewBuilder() ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) - envPath := MustComposerEnvironmentTimeline(ctx, "my-project", "my-env") + envPath := MustAirflowTimeline(ctx, "my-env") testCases := []struct { name string @@ -182,7 +162,7 @@ func TestMustAirflowDAGRunTimeline(t *testing.T) { func TestMustAirflowTaskInstanceTimeline(t *testing.T) { builder := khifilev6.NewBuilder() ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) - envPath := MustComposerEnvironmentTimeline(ctx, "my-project", "my-env") + envPath := MustAirflowTimeline(ctx, "my-env") runPath := MustAirflowDAGRunTimeline(ctx, envPath, "my-dag", "my-run") testCases := []struct { @@ -218,7 +198,7 @@ func TestMustAirflowTaskInstanceTimeline(t *testing.T) { func TestMustAirflowComponentTimeline(t *testing.T) { builder := khifilev6.NewBuilder() ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) - envPath := MustComposerEnvironmentTimeline(ctx, "my-project", "my-env") + envPath := MustAirflowTimeline(ctx, "my-env") testCases := []struct { name string @@ -239,10 +219,13 @@ func TestMustAirflowComponentTimeline(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - got := MustAirflowComponentTimeline(ctx, envPath, TimelineTypeAirflowScheduler, tc.componentName) + got := MustAirflowComponentTimeline(ctx, envPath, tc.componentName) if got == nil { t.Fatal("expected timeline path to be not nil") } + if diff := cmp.Diff(TimelineTypeAirflowComponent.GetId(), got.Type.GetId()); diff != "" { + t.Errorf("MustAirflowComponentTimeline() timeline type mismatch (-want +got):\n%s", diff) + } if diff := cmp.Diff(tc.wantName, got.Name.Resolve()); diff != "" { t.Errorf("MustAirflowComponentTimeline() componentName mismatch (-want +got):\n%s", diff) } @@ -250,45 +233,10 @@ func TestMustAirflowComponentTimeline(t *testing.T) { } } -func TestMustAirflowWorkerTimeline(t *testing.T) { - builder := khifilev6.NewBuilder() - ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) - envPath := MustComposerEnvironmentTimeline(ctx, "my-project", "my-env") - - testCases := []struct { - name string - workerHost string - wantHost string - }{ - { - name: "valid input", - workerHost: "worker-1", - wantHost: "worker-1", - }, - { - name: "empty worker host", - workerHost: "", - wantHost: "unknown", - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - got := MustAirflowWorkerTimeline(ctx, envPath, tc.workerHost) - if got == nil { - t.Fatal("expected timeline path to be not nil") - } - if diff := cmp.Diff(tc.wantHost, got.Name.Resolve()); diff != "" { - t.Errorf("MustAirflowWorkerTimeline() workerHost mismatch (-want +got):\n%s", diff) - } - }) - } -} - func TestMustAirflowDAGFileTimeline(t *testing.T) { builder := khifilev6.NewBuilder() ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) - envPath := MustComposerEnvironmentTimeline(ctx, "my-project", "my-env") + envPath := MustAirflowTimeline(ctx, "my-env") testCases := []struct { name string @@ -333,7 +281,7 @@ func TestMustAirflowDAGFileTimeline(t *testing.T) { func TestMustAirflowDAGProcessorManagerInstanceTimeline(t *testing.T) { builder := khifilev6.NewBuilder() ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) - envPath := MustComposerEnvironmentTimeline(ctx, "my-project", "my-env") + envPath := MustAirflowTimeline(ctx, "my-env") testCases := []struct { name string diff --git a/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_type.go b/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_type.go index 83ee47ef..c0f94a33 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_type.go +++ b/pkg/task/inspection/googlecloudclustercomposer/contract/timeline_type.go @@ -23,64 +23,60 @@ import ( // when this package is imported. var ( - // TimelineTypeComposerEnvironment is the style for a Cloud Composer environment. - // Background is set to #377e22. - TimelineTypeComposerEnvironment = style.MustRegisterTimelineType( + // TimelineTypeAirflow is the style for an Airflow environment. + TimelineTypeAirflow = style.MustRegisterTimelineType( "Airflow", "Timeline representing a Managed Airflow environment", "settings", 0.6, - style.MustForceConvertSRGBHex("#377e22"), - style.ColorWhite, - style.MustForceConvertSRGBHex("#377e22"), + style.Color{R: 0.82, G: 0.28, B: 0.19, A: 1}, style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, 20, style.AlphabeticalSortPolicy(), ) // TimelineTypeDAGs is the root style for DAGs hierarchy. - // Progressive lighter green background #5cb239. TimelineTypeDAGs = style.MustRegisterTimelineType( - "dags", + "DAGs", "Grouping timeline for Airflow DAGs", "folder", 0.6, - style.MustForceConvertSRGBHex("#5cb239"), - style.ColorWhite, - style.MustForceConvertSRGBHex("#5cb239"), + style.Color{R: 0.93, G: 0.49, B: 0.38, A: 1}, style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, 30, style.AlphabeticalSortPolicy(), ) // TimelineTypeAirflowDAG is the style for a single DAG. - // Progressive lighter green background #89ca6a. TimelineTypeAirflowDAG = style.MustRegisterTimelineType( - "airflow_dag", + "DAG", "Timeline representing an Airflow DAG", "account_tree", 0.6, - style.MustForceConvertSRGBHex("#89ca6a"), - style.ColorWhite, - style.MustForceConvertSRGBHex("#89ca6a"), + style.MustForceConvertSRGBHex("#444444"), style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, 40, - style.AlphabeticalSortPolicy(), + style.ChronologicalSortPolicy(2), ) // TimelineTypeAirflowDAGRun is the style for a DAG run. - // Progressive lighter green background #bce3a5. TimelineTypeAirflowDAGRun = style.MustRegisterTimelineType( - "airflow_dag_run", + "DAG Run", "Timeline representing an Airflow DAG run", "play_circle", 0.6, - style.MustForceConvertSRGBHex("#bce3a5"), + style.MustForceConvertSRGBHex("#CCCCCC"), style.ColorBlack, - style.MustForceConvertSRGBHex("#bce3a5"), + style.MustForceConvertSRGBHex("#F5F5F5"), style.ColorBlack, true, 50, @@ -88,90 +84,92 @@ var ( ) // TimelineTypeAirflowTaskInstance is the style for a TaskInstance. - // As it contains raw log/revisions, its background is set to White. TimelineTypeAirflowTaskInstance = style.MustRegisterTimelineType( - "task", + "Task Instance", "Execution states of the Airflow task instance", "mode_fan", - 0.6, + 0.7, style.ColorWhite, style.ColorBlack, - style.MustForceConvertSRGBHex("#377e22"), - style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, 1501, style.ChronologicalSortPolicy(1), ) // TimelineTypeComponents is the category style for components. - // Progressive lighter green background #5cb239. TimelineTypeComponents = style.MustRegisterTimelineType( - "airflow_components", + "Components", "Grouping timeline for Airflow backend components", "apps", 0.6, - style.MustForceConvertSRGBHex("#5cb239"), - style.ColorWhite, - style.MustForceConvertSRGBHex("#5cb239"), + style.Color{R: 0.93, G: 0.49, B: 0.38, A: 1}, style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, - 60, + 70, style.AlphabeticalSortPolicy(), ) - // TimelineTypeDAGProcessorManager is the category style for DAG Processor Manager stats. - // Progressive lighter green background #5cb239. - TimelineTypeDAGProcessorManager = style.MustRegisterTimelineType( - "dag_files", + // TimelineTypeDAGFiles is the category style for DAG Processor Manager stats. + TimelineTypeDAGFiles = style.MustRegisterTimelineType( + "DAG files", "Grouping timeline for parsed DAG files", "folder", 0.6, - style.MustForceConvertSRGBHex("#5cb239"), - style.ColorWhite, - style.MustForceConvertSRGBHex("#5cb239"), + style.Color{R: 0.93, G: 0.49, B: 0.38, A: 1}, style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, - 70, + 60, style.AlphabeticalSortPolicy(), ) // TimelineTypeDAGFile is the style for a parsed DAG file. - // Progressive lighter green background #89ca6a. TimelineTypeDAGFile = style.MustRegisterTimelineType( - "dag_file", + "DAG File", "Timeline representing an Airflow DAG definition file", "description", 0.6, - style.MustForceConvertSRGBHex("#89ca6a"), - style.ColorWhite, - style.MustForceConvertSRGBHex("#89ca6a"), + style.MustForceConvertSRGBHex("#444444"), style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, 80, style.AlphabeticalSortPolicy(), ) // TimelineTypeDAGProcessorManagerInstance is the style for the manager instance that processed the file. - // As it contains revisions, its background is set to White. TimelineTypeDAGProcessorManagerInstance = style.MustRegisterTimelineType( - "dag_processor_manager_instance", - "Logs of the DAG Processor Manager instance", + "Parser", + "Logs of the DAG Processor Manager instance. Same DAG file can be parsed from multiple DAG Processor Manager instances at the same time thus this is shown as separated timelines.", "terminal", 0.6, style.ColorWhite, style.ColorBlack, - style.MustForceConvertSRGBHex("#A51915"), - style.ColorWhite, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, true, 90, style.AlphabeticalSortPolicy(), ) - // Components specific timelines (Actual log/event containers: White background) - TimelineTypeAirflowScheduler = style.MustRegisterTimelineType("airflow_scheduler", "Logs of the Airflow Scheduler", "schedule", 0.6, style.ColorWhite, style.ColorBlack, style.MustForceConvertSRGBHex("#4285F4"), style.ColorWhite, true, 100, style.AlphabeticalSortPolicy()) - TimelineTypeAirflowWorker = style.MustRegisterTimelineType("airflow_worker", "Logs of the Airflow Worker", "directions_run", 0.6, style.ColorWhite, style.ColorBlack, style.MustForceConvertSRGBHex("#0F9D58"), style.ColorWhite, true, 110, style.AlphabeticalSortPolicy()) - TimelineTypeAirflowDagProcessorManager = style.MustRegisterTimelineType("airflow_dag_processor_manager", "Logs of the Airflow DAG Processor Manager", "summarize", 0.6, style.ColorWhite, style.ColorBlack, style.MustForceConvertSRGBHex("#808080"), style.ColorWhite, true, 115, style.AlphabeticalSortPolicy()) - TimelineTypeAirflowTriggerer = style.MustRegisterTimelineType("airflow_triggerer", "Logs of the Airflow Triggerer", "bolt", 0.6, style.ColorWhite, style.ColorBlack, style.MustForceConvertSRGBHex("#FFBB00"), style.ColorBlack, true, 120, style.AlphabeticalSortPolicy()) - TimelineTypeAirflowWebserver = style.MustRegisterTimelineType("airflow_webserver", "Logs of the Airflow Webserver", "web", 0.6, style.ColorWhite, style.ColorBlack, style.MustForceConvertSRGBHex("#9470DC"), style.ColorWhite, true, 130, style.AlphabeticalSortPolicy()) - TimelineTypeAirflowComponent = style.MustRegisterTimelineType("airflow_component", "Logs of the generic Airflow component", "extension", 0.6, style.ColorWhite, style.ColorBlack, style.MustForceConvertSRGBHex("#808080"), style.ColorWhite, true, 140, style.AlphabeticalSortPolicy()) + // TimelineTypeAirflowComponent is the style for Airflow components. + TimelineTypeAirflowComponent = style.MustRegisterTimelineType( + "Component", + "Logs of the generic Airflow component", + "extension", + 0.6, + style.ColorWhite, + style.ColorBlack, + style.MustForceConvertSRGBHex("#F5F5F5"), + style.ColorBlack, + true, + 100, + style.AlphabeticalSortPolicy(), + ) ) diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper.go b/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper.go index c21dec73..c8742e33 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper.go @@ -166,7 +166,6 @@ func (m *dagProcessorManagerTimelineMapper) LogIngesterTask() taskid.TaskReferen // Dependencies returns additional task dependencies of the mapper. func (m *dagProcessorManagerTimelineMapper) Dependencies() []taskid.UntypedTaskReference { return []taskid.UntypedTaskReference{ - googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref(), googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref(), } } @@ -178,9 +177,8 @@ func (m *dagProcessorManagerTimelineMapper) GroupedLogTask() taskid.TaskReferenc // ProcessLogByGroup is called for each log entry to stage mutations via TimelineChangeSet. func (m *dagProcessorManagerTimelineMapper) ProcessLogByGroup(ctx context.Context, l *log.Log, prevGroupData *DagProcessorState) (*khifilev6.TimelineChangeSet, *DagProcessorState, error) { - clusterIdentity := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref()) environmentName := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref()) - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, clusterIdentity.ProjectID, environmentName) + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, environmentName) commonField, _ := log.GetFieldSet(l, &log.CommonFieldSet{}) mainMessage, err := log.GetFieldSet(l, &googlecloudcommon_contract.GCPMainMessageFieldSet{}) @@ -192,10 +190,10 @@ func (m *dagProcessorManagerTimelineMapper) ProcessLogByGroup(ctx context.Contex parserID := "unknown-parser" if err == nil { if dpmField.SchedulerID != "" { - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowScheduler, dpmField.SchedulerID)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, dpmField.SchedulerID)) parserID = dpmField.SchedulerID } else if dpmField.DagProcessorManagerID != "" { - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowDagProcessorManager, dpmField.DagProcessorManagerID)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, dpmField.DagProcessorManagerID)) parserID = dpmField.DagProcessorManagerID } } diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper_test.go b/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper_test.go index 6decacc4..92f3d9dd 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper_test.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/dag_processor_manager_mapper_test.go @@ -29,7 +29,6 @@ import ( core_contract "github.com/GoogleCloudPlatform/khi/pkg/task/core/contract" googlecloudclustercomposer_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudclustercomposer/contract" googlecloudcommon_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudcommon/contract" - googlecloudk8scommon_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudk8scommon/contract" inspectioncore_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/inspectioncore/contract" "github.com/GoogleCloudPlatform/khi/pkg/testutil/testchangeset" "github.com/google/go-cmp/cmp" @@ -106,7 +105,7 @@ func TestDagProcessorMapperTask_ProcessLogByGroup(t *testing.T) { }, func(t *testing.T, ctx context.Context, cs *khifilev6.TimelineChangeSet) { // Data line - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, "test-project", "test-environment") + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, "test-environment") timelinePath := googlecloudclustercomposer_contract.MustAirflowDAGProcessorManagerInstanceTimeline(ctx, envPath, "/home/airflow/gcs/dags/airflow_monitoring.py", "unknown-parser") testchangeset.AssertTimeline(t, cs). HasRevision(timelinePath, &khifilev6.StagingRevision{ @@ -136,7 +135,6 @@ func TestDagProcessorMapperTask_ProcessLogByGroup(t *testing.T) { ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) taskDependentValues := typedmap.NewTypedMap() - typedmap.Set(taskDependentValues, typedmap.NewTypedKey[googlecloudk8scommon_contract.GoogleCloudClusterIdentity](googlecloudclustercomposer_contract.ClusterIdentityTaskID.ReferenceIDString()), googlecloudk8scommon_contract.GoogleCloudClusterIdentity{ProjectID: "test-project"}) typedmap.Set(taskDependentValues, typedmap.NewTypedKey[string](googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.ReferenceIDString()), "test-environment") ctx = khictx.WithValue(ctx, core_contract.TaskResultMapContextKey, taskDependentValues) diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/other.go b/pkg/task/inspection/googlecloudclustercomposer/impl/other.go index 8a6604a8..c3111a60 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/other.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/other.go @@ -91,7 +91,6 @@ func (m *otherLogToTimelineMapper) LogIngesterTask() taskid.TaskReference[[]*log // Dependencies returns additional task dependencies of the mapper. func (m *otherLogToTimelineMapper) Dependencies() []taskid.UntypedTaskReference { return []taskid.UntypedTaskReference{ - googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref(), googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref(), } } @@ -103,9 +102,8 @@ func (m *otherLogToTimelineMapper) GroupedLogTask() taskid.TaskReference[inspect // ProcessLogByGroup is called for each log entry to stage mutations via TimelineChangeSet. func (m *otherLogToTimelineMapper) ProcessLogByGroup(ctx context.Context, l *log.Log, _ struct{}) (*khifilev6.TimelineChangeSet, struct{}, error) { - clusterIdentity := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref()) environmentName := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref()) - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, clusterIdentity.ProjectID, environmentName) + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, environmentName) composerFieldSet, err := log.GetFieldSet(l, &googlecloudclustercomposer_contract.ComposerFieldSet{}) if err != nil { @@ -120,27 +118,27 @@ func (m *otherLogToTimelineMapper) ProcessLogByGroup(ctx context.Context, l *log mappedToTimeline := false if composerFieldSet.WorkerID != "" { - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowWorker, composerFieldSet.WorkerID)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, composerFieldSet.WorkerID)) mappedToTimeline = true } if composerFieldSet.SchedulerID != "" { - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowScheduler, composerFieldSet.SchedulerID)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, composerFieldSet.SchedulerID)) mappedToTimeline = true } if composerFieldSet.DagProcessorManagerID != "" { - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowDagProcessorManager, composerFieldSet.DagProcessorManagerID)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, composerFieldSet.DagProcessorManagerID)) mappedToTimeline = true } if composerFieldSet.TriggererID != "" { - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowTriggerer, composerFieldSet.TriggererID)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, composerFieldSet.TriggererID)) mappedToTimeline = true } if composerFieldSet.WebserverID != "" { - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowWebserver, composerFieldSet.WebserverID)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, composerFieldSet.WebserverID)) mappedToTimeline = true } @@ -148,7 +146,7 @@ func (m *otherLogToTimelineMapper) ProcessLogByGroup(ctx context.Context, l *log if composerFieldSet.Subservice != "" { componentName = composerFieldSet.Subservice } - cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowComponent, componentName)) + cs.AddEvent(googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, componentName)) } return cs, struct{}{}, nil diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler.go b/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler.go index 029009e0..f533dc07 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler.go @@ -92,7 +92,6 @@ func (m *schedulerLogToTimelineMapper) LogIngesterTask() taskid.TaskReference[[] // Dependencies returns additional task dependencies of the mapper. func (m *schedulerLogToTimelineMapper) Dependencies() []taskid.UntypedTaskReference { return []taskid.UntypedTaskReference{ - googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref(), googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref(), } } @@ -104,16 +103,15 @@ func (m *schedulerLogToTimelineMapper) GroupedLogTask() taskid.TaskReference[ins // ProcessLogByGroup is called for each log entry to stage mutations via TimelineChangeSet. func (m *schedulerLogToTimelineMapper) ProcessLogByGroup(ctx context.Context, l *log.Log, _ struct{}) (*khifilev6.TimelineChangeSet, struct{}, error) { - clusterIdentity := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref()) environmentName := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref()) - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, clusterIdentity.ProjectID, environmentName) + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, environmentName) schedulerField, err := log.GetFieldSet(l, &googlecloudclustercomposer_contract.ComposerFieldSet{}) cs := khifilev6.NewTimelineChangeSet(l) if err == nil { if schedulerField.SchedulerID != "" { - schedulerTimelinePath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowScheduler, schedulerField.SchedulerID) + schedulerTimelinePath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, schedulerField.SchedulerID) cs.AddEvent(schedulerTimelinePath) } } @@ -149,7 +147,7 @@ func (m *schedulerLogToTimelineMapper) ProcessLogByGroup(ctx context.Context, l // If the ti status is zombie, record it on worker if ti.Status() == googlecloudclustercomposer_contract.TASKINSTANCE_ZOMBIE && ti.Host() != "" { - workerTimelinePath := googlecloudclustercomposer_contract.MustAirflowWorkerTimeline(ctx, envPath, ti.Host()) + workerTimelinePath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, ti.Host()) cs.AddEvent(workerTimelinePath) } diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler_test.go b/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler_test.go index 9a7a07a2..04434d42 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler_test.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/scheduler_test.go @@ -28,7 +28,6 @@ import ( core_contract "github.com/GoogleCloudPlatform/khi/pkg/task/core/contract" googlecloudclustercomposer_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudclustercomposer/contract" googlecloudcommon_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudcommon/contract" - googlecloudk8scommon_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudk8scommon/contract" inspectioncore_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/inspectioncore/contract" "github.com/GoogleCloudPlatform/khi/pkg/testutil/testchangeset" "github.com/google/go-cmp/cmp" @@ -57,8 +56,8 @@ func TestAirflowSchedulerMapperTask_ProcessLogByGroup(t *testing.T) { }, ), assert: func(t *testing.T, ctx context.Context, cs *khifilev6.TimelineChangeSet) { - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, "test-project", "test-environment") - schedulerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowScheduler, "airflow-scheduler-7b5f") + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, "test-environment") + schedulerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, "airflow-scheduler-7b5f") ti := googlecloudclustercomposer_contract.NewAirflowTaskInstance("my_dag", "task_id_1", "2023-01-01T00:00:00Z", "1", "worker-1", googlecloudclustercomposer_contract.TASKINSTANCE_SUCCESS) runPath := googlecloudclustercomposer_contract.MustAirflowDAGRunTimeline(ctx, envPath, ti.DagId(), ti.RunId()) tiPath := googlecloudclustercomposer_contract.MustAirflowTaskInstanceTimeline(ctx, runPath, "task_id_1+1") @@ -96,9 +95,9 @@ func TestAirflowSchedulerMapperTask_ProcessLogByGroup(t *testing.T) { }, ), assert: func(t *testing.T, ctx context.Context, cs *khifilev6.TimelineChangeSet) { - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, "test-project", "test-environment") - schedulerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowScheduler, "airflow-scheduler-7b5f") - workerPath := googlecloudclustercomposer_contract.MustAirflowWorkerTimeline(ctx, envPath, "worker-bad") + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, "test-environment") + schedulerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, "airflow-scheduler-7b5f") + workerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, "worker-bad") ti := googlecloudclustercomposer_contract.NewAirflowTaskInstance("my_dag", "task_id_zombie", "2023-01-01T00:00:00Z", "1", "worker-bad", googlecloudclustercomposer_contract.TASKINSTANCE_ZOMBIE) runPath := googlecloudclustercomposer_contract.MustAirflowDAGRunTimeline(ctx, envPath, ti.DagId(), ti.RunId()) tiPath := googlecloudclustercomposer_contract.MustAirflowTaskInstanceTimeline(ctx, runPath, "task_id_zombie+1") @@ -132,8 +131,8 @@ func TestAirflowSchedulerMapperTask_ProcessLogByGroup(t *testing.T) { }, ), assert: func(t *testing.T, ctx context.Context, cs *khifilev6.TimelineChangeSet) { - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, "test-project", "test-environment") - schedulerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, googlecloudclustercomposer_contract.TimelineTypeAirflowScheduler, "airflow-scheduler-7b5f") + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, "test-environment") + schedulerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, "airflow-scheduler-7b5f") testchangeset.AssertTimeline(t, cs). HasEvent(schedulerPath) @@ -148,7 +147,6 @@ func TestAirflowSchedulerMapperTask_ProcessLogByGroup(t *testing.T) { ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) taskDependentValues := typedmap.NewTypedMap() - typedmap.Set(taskDependentValues, typedmap.NewTypedKey[googlecloudk8scommon_contract.GoogleCloudClusterIdentity](googlecloudclustercomposer_contract.ClusterIdentityTaskID.ReferenceIDString()), googlecloudk8scommon_contract.GoogleCloudClusterIdentity{ProjectID: "test-project"}) typedmap.Set(taskDependentValues, typedmap.NewTypedKey[string](googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.ReferenceIDString()), "test-environment") ctx = khictx.WithValue(ctx, core_contract.TaskResultMapContextKey, taskDependentValues) diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/worker.go b/pkg/task/inspection/googlecloudclustercomposer/impl/worker.go index 7b4b2c7e..db99b66a 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/worker.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/worker.go @@ -92,7 +92,6 @@ func (m *workerLogToTimelineMapper) LogIngesterTask() taskid.TaskReference[[]*lo // Dependencies returns additional task dependencies of the mapper. func (m *workerLogToTimelineMapper) Dependencies() []taskid.UntypedTaskReference { return []taskid.UntypedTaskReference{ - googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref(), googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref(), } } @@ -104,16 +103,15 @@ func (m *workerLogToTimelineMapper) GroupedLogTask() taskid.TaskReference[inspec // ProcessLogByGroup is called for each log entry to stage mutations via TimelineChangeSet. func (m *workerLogToTimelineMapper) ProcessLogByGroup(ctx context.Context, l *log.Log, _ struct{}) (*khifilev6.TimelineChangeSet, struct{}, error) { - clusterIdentity := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.ClusterIdentityTaskID.Ref()) environmentName := coretask.GetTaskResult(ctx, googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.Ref()) - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, clusterIdentity.ProjectID, environmentName) + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, environmentName) workerField, err := log.GetFieldSet(l, &googlecloudclustercomposer_contract.ComposerFieldSet{}) cs := khifilev6.NewTimelineChangeSet(l) if err == nil { if workerField.WorkerID != "" { - workerTimelinePath := googlecloudclustercomposer_contract.MustAirflowWorkerTimeline(ctx, envPath, workerField.WorkerID) + workerTimelinePath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, workerField.WorkerID) cs.AddEvent(workerTimelinePath) } } diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go b/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go index 07b87a98..623bc9eb 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go @@ -28,7 +28,6 @@ import ( core_contract "github.com/GoogleCloudPlatform/khi/pkg/task/core/contract" googlecloudclustercomposer_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudclustercomposer/contract" googlecloudcommon_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudcommon/contract" - googlecloudk8scommon_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/googlecloudk8scommon/contract" inspectioncore_contract "github.com/GoogleCloudPlatform/khi/pkg/task/inspection/inspectioncore/contract" "github.com/GoogleCloudPlatform/khi/pkg/testutil/testchangeset" "github.com/google/go-cmp/cmp" @@ -57,8 +56,8 @@ func TestAirflowWorkerMapperTask_ProcessLogByGroup(t *testing.T) { }, ), assert: func(t *testing.T, ctx context.Context, cs *khifilev6.TimelineChangeSet) { - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, "test-project", "test-environment") - workerPath := googlecloudclustercomposer_contract.MustAirflowWorkerTimeline(ctx, envPath, "airflow-worker-abc") + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, "test-environment") + workerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, "airflow-worker-abc") ti := googlecloudclustercomposer_contract.NewAirflowTaskInstance("my_dag", "task_id_1", "2023-01-01T00:00:00Z", "1", "airflow-worker-abc", googlecloudclustercomposer_contract.TASKINSTANCE_RUNNING) runPath := googlecloudclustercomposer_contract.MustAirflowDAGRunTimeline(ctx, envPath, ti.DagId(), ti.RunId()) tiPath := googlecloudclustercomposer_contract.MustAirflowTaskInstanceTimeline(ctx, runPath, "task_id_1+1") @@ -90,8 +89,8 @@ func TestAirflowWorkerMapperTask_ProcessLogByGroup(t *testing.T) { }, ), assert: func(t *testing.T, ctx context.Context, cs *khifilev6.TimelineChangeSet) { - envPath := googlecloudclustercomposer_contract.MustComposerEnvironmentTimeline(ctx, "test-project", "test-environment") - workerPath := googlecloudclustercomposer_contract.MustAirflowWorkerTimeline(ctx, envPath, "airflow-worker-abc") + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, "test-environment") + workerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, "airflow-worker-abc") testchangeset.AssertTimeline(t, cs). HasEvent(workerPath) @@ -106,7 +105,6 @@ func TestAirflowWorkerMapperTask_ProcessLogByGroup(t *testing.T) { ctx := khictx.WithValue(t.Context(), inspectioncore_contract.Builder, builder) taskDependentValues := typedmap.NewTypedMap() - typedmap.Set(taskDependentValues, typedmap.NewTypedKey[googlecloudk8scommon_contract.GoogleCloudClusterIdentity](googlecloudclustercomposer_contract.ClusterIdentityTaskID.ReferenceIDString()), googlecloudk8scommon_contract.GoogleCloudClusterIdentity{ProjectID: "test-project"}) typedmap.Set(taskDependentValues, typedmap.NewTypedKey[string](googlecloudclustercomposer_contract.InputComposerEnvironmentNameTaskID.ReferenceIDString()), "test-environment") ctx = khictx.WithValue(ctx, core_contract.TaskResultMapContextKey, taskDependentValues)