diff --git a/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset.go b/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset.go index bb248378..f0ceb812 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset.go +++ b/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset.go @@ -16,6 +16,7 @@ package googlecloudclustercomposer_contract import ( "fmt" + "log/slog" "regexp" "strings" @@ -70,7 +71,7 @@ var ( ) func stringToTiState(stateStr string) (Tistate, error) { - switch stateStr { + switch strings.ToLower(stateStr) { case "scheduled": return TASKINSTANCE_SCHEDULED, nil case "queued": @@ -261,108 +262,134 @@ var ( airflowWorkerFinalStateExtractTemplate = regexp.MustCompile(`\[final_state=(?P[a-z_]*)\]`) ) -func (c *ComposerWorkerTaskInstanceFieldSetReader) Read(reader *structured.NodeReader) (log.FieldSet, error) { - if ti, err := c.tryReadFromLabels(reader); err == nil { - return &ComposerWorkerTaskInstanceFieldSet{TaskInstance: ti}, nil +type workerTaskInstanceInfo struct { + dagId string + taskId string + runId string + mapIndex string + workerId string + state Tistate +} + +// merge returns a new workerTaskInstanceInfo where empty fields in i are filled from other, +// and state is updated if other.state is specified. +func (i workerTaskInstanceInfo) merge(other workerTaskInstanceInfo) workerTaskInstanceInfo { + res := i + if res.dagId == "" { + res.dagId = other.dagId } - textPayload, err := reader.ReadString("textPayload") - if err != nil { - return nil, fmt.Errorf("textPayload not found") + if res.taskId == "" { + res.taskId = other.taskId + } + if res.runId == "" { + res.runId = other.runId } + if (res.mapIndex == "" || res.mapIndex == "-1") && other.mapIndex != "" { + res.mapIndex = other.mapIndex + } + if other.workerId != "" { + res.workerId = other.workerId + } + if other.state != TASKINSTANCE_NONE && other.state != "" { + res.state = other.state + } + return res +} +func (c *ComposerWorkerTaskInstanceFieldSetReader) readLabels(reader *structured.NodeReader) workerTaskInstanceInfo { workerId, _ := reader.ReadString("labels.worker_id") + runId, _ := reader.ReadString("labels.run-id") + workflow, _ := reader.ReadString("labels.workflow") + taskId, _ := reader.ReadString("labels.task-id") + mapIndex, _ := reader.ReadString("labels.map-index") + + return workerTaskInstanceInfo{ + dagId: workflow, + taskId: taskId, + runId: runId, + mapIndex: mapIndex, + workerId: workerId, + state: TASKINSTANCE_NONE, + } +} +// parsePayload attempts to parse Airflow task instance details from the log text payload. +func (c *ComposerWorkerTaskInstanceFieldSetReader) parsePayload(textPayload string) workerTaskInstanceInfo { if strings.HasPrefix(textPayload, "Running ") { - matches := airflowWorkerRunningHostTemplate.FindStringSubmatch(textPayload) - if matches != nil { - dagid := matches[airflowWorkerRunningHostTemplate.SubexpIndex("dagid")] - taskid := matches[airflowWorkerRunningHostTemplate.SubexpIndex("taskid")] - runid := matches[airflowWorkerRunningHostTemplate.SubexpIndex("runid")] - host := matches[airflowWorkerRunningHostTemplate.SubexpIndex("host")] + if matches := airflowWorkerRunningHostTemplate.FindStringSubmatch(textPayload); matches != nil { + mapIndex := "" + if i := airflowWorkerRunningHostTemplate.SubexpIndex("mapIndex"); i >= 0 && matches[i] != "" { + mapIndex = matches[i] + } stateStr := matches[airflowWorkerRunningHostTemplate.SubexpIndex("state")] state, err := stringToTiState(stateStr) if err != nil { - return nil, err + slog.Warn(fmt.Sprintf("failed to parse task instance state %q: %v", stateStr, err)) } - mapIndex := "-1" - if i := airflowWorkerRunningHostTemplate.SubexpIndex("mapIndex"); i >= 0 && matches[i] != "" { - mapIndex = matches[i] + return workerTaskInstanceInfo{ + dagId: matches[airflowWorkerRunningHostTemplate.SubexpIndex("dagid")], + taskId: matches[airflowWorkerRunningHostTemplate.SubexpIndex("taskid")], + runId: matches[airflowWorkerRunningHostTemplate.SubexpIndex("runid")], + workerId: matches[airflowWorkerRunningHostTemplate.SubexpIndex("host")], + mapIndex: mapIndex, + state: state, } - return &ComposerWorkerTaskInstanceFieldSet{ - TaskInstance: NewAirflowTaskInstance(dagid, taskid, runid, mapIndex, host, state), - }, nil } } - - matches := airflowWorkerMarkingStatusTemplate.FindStringSubmatch(textPayload) - if matches != nil { - if workerId == "" { - return nil, fmt.Errorf("worker_id not found") + if matches := airflowWorkerMarkingStatusTemplate.FindStringSubmatch(textPayload); matches != nil { + mapIndex := "" + if i := airflowWorkerMarkingStatusTemplate.SubexpIndex("mapIndex"); i >= 0 && matches[i] != "" { + mapIndex = matches[i] } - - dagid := matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("dagid")] - taskid := matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("taskid")] - runid := matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("runid")] - // Need strings.ToLower because it might be capitalized depending on the version - stateStr := strings.ToLower(matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("state")]) + stateStr := matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("state")] state, err := stringToTiState(stateStr) if err != nil { - return nil, err + slog.Warn(fmt.Sprintf("failed to parse task instance state %q: %v", stateStr, err)) } - mapIndex := "-1" - if i := airflowWorkerMarkingStatusTemplate.SubexpIndex("mapIndex"); i >= 0 && matches[i] != "" { - mapIndex = matches[i] + return workerTaskInstanceInfo{ + dagId: matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("dagid")], + taskId: matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("taskid")], + runId: matches[airflowWorkerMarkingStatusTemplate.SubexpIndex("runid")], + mapIndex: mapIndex, + state: state, } - return &ComposerWorkerTaskInstanceFieldSet{ - TaskInstance: NewAirflowTaskInstance(dagid, taskid, runid, mapIndex, workerId, state), - }, nil } - - return nil, fmt.Errorf("not an Airflow Worker TaskInstance log") -} - -// tryReadFromLabels reads task instance info from labels if available. -// This is effective only when airflow 3.x is used. -func (c *ComposerWorkerTaskInstanceFieldSetReader) tryReadFromLabels(reader *structured.NodeReader) (*AirflowTaskInstance, error) { - workerId, err := reader.ReadString("labels.worker_id") - if err != nil { - return nil, fmt.Errorf("worker_id not found") - } - runid, err := reader.ReadString("labels.run-id") - if err != nil { - return nil, fmt.Errorf("run-id not found") + if matches := airflowWorkerFinalStateExtractTemplate.FindStringSubmatch(textPayload); matches != nil { + stateStr := matches[airflowWorkerFinalStateExtractTemplate.SubexpIndex("state")] + state, err := stringToTiState(stateStr) + if err != nil { + slog.Warn(fmt.Sprintf("failed to parse task instance state %q: %v", stateStr, err)) + } + return workerTaskInstanceInfo{ + state: state, + } } + return workerTaskInstanceInfo{} +} - workflow, err := reader.ReadString("labels.workflow") - if err != nil { - return nil, fmt.Errorf("workflow not found") - } +// Read parses structured log data to extract worker TaskInstance information. +func (c *ComposerWorkerTaskInstanceFieldSetReader) Read(reader *structured.NodeReader) (log.FieldSet, error) { + labelInfo := c.readLabels(reader) - taskid, err := reader.ReadString("labels.task-id") + textPayload, err := reader.ReadString("textPayload") if err != nil { - return nil, fmt.Errorf("task-id not found") + return nil, fmt.Errorf("textPayload not found") } - mapIndex, err := reader.ReadString("labels.map-index") - if err != nil { - return nil, fmt.Errorf("map-index not found") - } + payloadInfo := c.parsePayload(textPayload) + info := labelInfo.merge(payloadInfo) - textPayload, err := reader.ReadString("textPayload") - if err != nil { - return nil, fmt.Errorf("textPayload not found") + if info.mapIndex == "" { + info.mapIndex = "-1" } - matches := airflowWorkerFinalStateExtractTemplate.FindStringSubmatch(textPayload) - state := TASKINSTANCE_NONE - if matches != nil { - stateStr := strings.ToLower(matches[airflowWorkerFinalStateExtractTemplate.SubexpIndex("state")]) - if finalState, err := stringToTiState(stateStr); err == nil { - state = finalState - } + if info.dagId == "" || info.taskId == "" || info.runId == "" || info.workerId == "" { + return nil, fmt.Errorf("not an Airflow Worker TaskInstance log") } - return NewAirflowTaskInstance(workflow, taskid, runid, mapIndex, workerId, state), nil + return &ComposerWorkerTaskInstanceFieldSet{ + TaskInstance: NewAirflowTaskInstance(info.dagId, info.taskId, info.runId, info.mapIndex, info.workerId, info.state), + }, nil } var _ log.FieldSetReader = &ComposerWorkerTaskInstanceFieldSetReader{} diff --git a/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset_test.go b/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset_test.go index aa74d771..b0de94ab 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset_test.go +++ b/pkg/task/inspection/googlecloudclustercomposer/contract/fieldset_test.go @@ -262,6 +262,113 @@ func TestComposerWorkerTaskInstanceFieldSetReader_Read(t *testing.T) { ), }, }, + { + name: "running with labels (Airflow 2/3 on Cloud Composer)", + textPayload: `Running on host airflow-worker-mwtx5`, + labels: map[string]string{ + "worker_id": "airflow-worker-mwtx5", + "run-id": "scheduled__2026-08-05T16:36:00+00:00", + "workflow": "sample_khi_multi_step_dag", + "task-id": "step2_process", + "map-index": "-1", + }, + want: &ComposerWorkerTaskInstanceFieldSet{ + TaskInstance: NewAirflowTaskInstance( + "sample_khi_multi_step_dag", + "step2_process", + "scheduled__2026-08-05T16:36:00+00:00", + "-1", + "airflow-worker-mwtx5", + TASKINSTANCE_RUNNING, + ), + }, + }, + { + name: "final state skipped with labels (Airflow 3)", + textPayload: `Task finished [task_instance_id=019d10e4-71f5-7016-b412-aa1fbcfd16fc] [exit_code=0] [duration=0.41656468300061533] [final_state=skipped]`, + labels: map[string]string{ + "worker_id": "airflow-worker-test", + "run-id": "scheduled__2025-04-14T01:30:00+00:00", + "workflow": "airflow_monitoring", + "task-id": "echo", + "map-index": "2", + }, + want: &ComposerWorkerTaskInstanceFieldSet{ + TaskInstance: NewAirflowTaskInstance( + "airflow_monitoring", + "echo", + "scheduled__2025-04-14T01:30:00+00:00", + "2", + "airflow-worker-test", + TASKINSTANCE_SKIPPED, + ), + }, + }, + { + name: "marking status with labels", + textPayload: `Marking task as SUCCESS. dag_id=airflow_monitoring, task_id=echo, run_id=scheduled__2025-04-14T01:30:00+00:00, map_index=2, execution_date=20250414T013000, start_date=20250414T014000, end_date=20250414T014001`, + labels: map[string]string{ + "worker_id": "airflow-worker-5fqxd", + "run-id": "scheduled__2025-04-14T01:30:00+00:00", + "workflow": "airflow_monitoring", + "task-id": "echo", + "map-index": "2", + }, + want: &ComposerWorkerTaskInstanceFieldSet{ + TaskInstance: NewAirflowTaskInstance( + "airflow_monitoring", + "echo", + "scheduled__2025-04-14T01:30:00+00:00", + "2", + "airflow-worker-5fqxd", + TASKINSTANCE_SUCCESS, + ), + }, + }, + { + name: "missing worker_id with generic payload", + textPayload: `Any text payload`, + labels: map[string]string{ + "run-id": "scheduled__2025-04-14T01:30:00+00:00", + "workflow": "airflow_monitoring", + "task-id": "echo", + "map-index": "2", + }, + wantErr: true, + }, + { + name: "missing run-id with generic payload", + textPayload: `Any text payload`, + labels: map[string]string{ + "worker_id": "airflow-worker-test", + "workflow": "airflow_monitoring", + "task-id": "echo", + "map-index": "2", + }, + wantErr: true, + }, + { + name: "missing workflow with generic payload", + textPayload: `Any text payload`, + labels: map[string]string{ + "worker_id": "airflow-worker-test", + "run-id": "scheduled__2025-04-14T01:30:00+00:00", + "task-id": "echo", + "map-index": "2", + }, + wantErr: true, + }, + { + name: "missing task-id with generic payload", + textPayload: `Any text payload`, + labels: map[string]string{ + "worker_id": "airflow-worker-test", + "run-id": "scheduled__2025-04-14T01:30:00+00:00", + "workflow": "airflow_monitoring", + "map-index": "2", + }, + wantErr: true, + }, } for _, tt := range tests { @@ -309,109 +416,80 @@ func TestComposerWorkerTaskInstanceFieldSetReader_Read(t *testing.T) { } } -func TestComposerWorkerTaskInstanceFieldSetReader_tryReadFromLabels(t *testing.T) { +func TestComposerWorkerTaskInstanceFieldSetReader_parsePayload(t *testing.T) { reader := &ComposerWorkerTaskInstanceFieldSetReader{} - tests := []struct { - name string - labels map[string]string - want *AirflowTaskInstance - wantErr bool + testCases := []struct { + name string + textPayload string + wantInfo workerTaskInstanceInfo }{ { - name: "valid labels", - labels: map[string]string{ - "worker_id": "airflow-worker-abc", - "run-id": "scheduled__2025-04-14T01:30:00+00:00", - "workflow": "airflow_monitoring", - "task-id": "echo", - "map-index": "2", + name: "running without mapIndex", + textPayload: `Running on host airflow-worker-dpvl7`, + wantInfo: workerTaskInstanceInfo{ + dagId: "example", + taskId: "query3", + runId: "scheduled__2024-04-22T05:30:00+00:00", + workerId: "airflow-worker-dpvl7", + mapIndex: "", + state: TASKINSTANCE_QUEUED, }, - want: NewAirflowTaskInstance( - "airflow_monitoring", - "echo", - "scheduled__2025-04-14T01:30:00+00:00", - "2", - "airflow-worker-abc", - TASKINSTANCE_NONE, - ), }, { - name: "missing worker_id", - labels: map[string]string{ - "run-id": "scheduled__2025-04-14T01:30:00+00:00", - "workflow": "airflow_monitoring", - "task-id": "echo", - "map-index": "2", + name: "running with mapIndex", + textPayload: `Running on host airflow-worker-dpvl7`, + wantInfo: workerTaskInstanceInfo{ + dagId: "example", + taskId: "query3", + runId: "scheduled__2024-04-22T05:30:00+00:00", + workerId: "airflow-worker-dpvl7", + mapIndex: "2", + state: TASKINSTANCE_RUNNING, }, - wantErr: true, }, { - name: "missing run-id", - labels: map[string]string{ - "worker_id": "airflow-worker-abc", - "workflow": "airflow_monitoring", - "task-id": "echo", - "map-index": "2", + name: "running with unknown state", + textPayload: `Running on host airflow-worker-dpvl7`, + wantInfo: workerTaskInstanceInfo{ + dagId: "example", + taskId: "query3", + runId: "scheduled__2024-04-22T05:30:00+00:00", + workerId: "airflow-worker-dpvl7", + mapIndex: "", + state: "", }, - wantErr: true, }, { - name: "missing workflow", - labels: map[string]string{ - "worker_id": "airflow-worker-abc", - "run-id": "scheduled__2025-04-14T01:30:00+00:00", - "task-id": "echo", - "map-index": "2", + name: "marking status success", + textPayload: `Marking task as SUCCESS. dag_id=airflow_monitoring, task_id=echo, run_id=scheduled__2025-04-14T01:30:00+00:00, execution_date=20250414T013000, start_date=20250414T014000, end_date=20250414T014001`, + wantInfo: workerTaskInstanceInfo{ + dagId: "airflow_monitoring", + taskId: "echo", + runId: "scheduled__2025-04-14T01:30:00+00:00", + mapIndex: "", + state: TASKINSTANCE_SUCCESS, }, - wantErr: true, }, { - name: "missing task-id", - labels: map[string]string{ - "worker_id": "airflow-worker-abc", - "run-id": "scheduled__2025-04-14T01:30:00+00:00", - "workflow": "airflow_monitoring", - "map-index": "2", + name: "final state skipped", + textPayload: `Task finished [task_instance_id=019d10e4-71f5-7016-b412-aa1fbcfd16fc] [exit_code=0] [duration=0.41656468300061533] [final_state=skipped]`, + wantInfo: workerTaskInstanceInfo{ + state: TASKINSTANCE_SKIPPED, }, - wantErr: true, }, { - name: "missing map-index", - labels: map[string]string{ - "worker_id": "airflow-worker-abc", - "run-id": "scheduled__2025-04-14T01:30:00+00:00", - "workflow": "airflow_monitoring", - "task-id": "echo", - }, - wantErr: true, + name: "unmatched payload", + textPayload: `Some generic log line without task instance state`, + wantInfo: workerTaskInstanceInfo{}, }, } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - yamlStr := `textPayload: "task is running" -labels:` - for k, v := range tt.labels { - yamlStr += fmt.Sprintf("\n \"%s\": '%s'", k, v) - } - if len(tt.labels) == 0 { - yamlStr = "{}" - } - yamlNode, err := structured.FromYAML(yamlStr) - if err != nil { - t.Fatalf("failed to parse yaml: %v", err) - } - nodeReader := structured.NewNodeReader(yamlNode) - - got, err := reader.tryReadFromLabels(nodeReader) - if (err != nil) != tt.wantErr { - t.Errorf("tryReadFromLabels() error = %v, wantErr %v", err, tt.wantErr) - } - if tt.want != nil { - if diff := cmp.Diff(tt.want, got, cmp.AllowUnexported(AirflowTaskInstance{})); diff != "" { - t.Errorf("tryReadFromLabels() mismatch (-want +got):\n%s", diff) - } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + gotInfo := reader.parsePayload(tc.textPayload) + if diff := cmp.Diff(tc.wantInfo, gotInfo, cmp.AllowUnexported(workerTaskInstanceInfo{})); diff != "" { + t.Errorf("parsePayload() mismatch (-want +got):\n%s", diff) } }) } diff --git a/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go b/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go index 623bc9eb..6082b1db 100644 --- a/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go +++ b/pkg/task/inspection/googlecloudclustercomposer/impl/worker_test.go @@ -96,6 +96,31 @@ func TestAirflowWorkerMapperTask_ProcessLogByGroup(t *testing.T) { HasEvent(workerPath) }, }, + { + name: "Worker TaskInstance with none status generates event", + input: log.NewLogWithFieldSetsForTest( + &log.CommonFieldSet{Timestamp: timestamp}, + &googlecloudcommon_contract.GCPMainMessageFieldSet{MainMessage: "Any user task log"}, + &googlecloudclustercomposer_contract.ComposerFieldSet{ + WorkerID: "airflow-worker-abc", + }, + &googlecloudclustercomposer_contract.ComposerWorkerTaskInstanceFieldSet{ + TaskInstance: googlecloudclustercomposer_contract.NewAirflowTaskInstance( + "my_dag", "task_id_1", "2023-01-01T00:00:00Z", "-1", "airflow-worker-abc", googlecloudclustercomposer_contract.TASKINSTANCE_NONE, + ), + }, + ), + assert: func(t *testing.T, ctx context.Context, cs *khifilev6.TimelineChangeSet) { + envPath := googlecloudclustercomposer_contract.MustAirflowTimeline(ctx, "test-environment") + workerPath := googlecloudclustercomposer_contract.MustAirflowComponentTimeline(ctx, envPath, "airflow-worker-abc") + runPath := googlecloudclustercomposer_contract.MustAirflowDAGRunTimeline(ctx, envPath, "my_dag", "2023-01-01T00:00:00Z") + tiPath := googlecloudclustercomposer_contract.MustAirflowTaskInstanceTimeline(ctx, runPath, "task_id_1") + + testchangeset.AssertTimeline(t, cs). + HasEvent(workerPath). + HasEvent(tiPath) + }, + }, } mapper := &workerLogToTimelineMapper{}