github.com/databricks/cli@v0.203.0/bundle/run/output/job.go (about)

     1  package output
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sort"
     7  	"strings"
     8  
     9  	"github.com/databricks/databricks-sdk-go"
    10  	"github.com/databricks/databricks-sdk-go/service/jobs"
    11  )
    12  
    13  type TaskOutput struct {
    14  	TaskKey string
    15  	Output  RunOutput
    16  	EndTime int64
    17  }
    18  
    19  type JobOutput struct {
    20  	// output for tasks with a non empty output
    21  	TaskOutputs []TaskOutput `json:"task_outputs"`
    22  }
    23  
    24  // Returns tasks output in text form sorted in execution order based on task end time
    25  func (out *JobOutput) String() (string, error) {
    26  	if len(out.TaskOutputs) == 0 {
    27  		return "", nil
    28  	}
    29  	// When only one task, just return that output without any formatting
    30  	if len(out.TaskOutputs) == 1 {
    31  		for _, v := range out.TaskOutputs {
    32  			return v.Output.String()
    33  		}
    34  	}
    35  	result := strings.Builder{}
    36  	result.WriteString("Output:\n")
    37  	sort.Slice(out.TaskOutputs, func(i, j int) bool {
    38  		return out.TaskOutputs[i].EndTime < out.TaskOutputs[j].EndTime
    39  	})
    40  	for _, v := range out.TaskOutputs {
    41  		if v.Output == nil {
    42  			continue
    43  		}
    44  		taskString, err := v.Output.String()
    45  		if err != nil {
    46  			return "", nil
    47  		}
    48  		result.WriteString("=======\n")
    49  		result.WriteString(fmt.Sprintf("Task %s:\n", v.TaskKey))
    50  		result.WriteString(fmt.Sprintf("%s\n", taskString))
    51  	}
    52  	return result.String(), nil
    53  }
    54  
    55  func GetJobOutput(ctx context.Context, w *databricks.WorkspaceClient, runId int64) (*JobOutput, error) {
    56  	jobRun, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{
    57  		RunId: runId,
    58  	})
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  	result := &JobOutput{
    63  		TaskOutputs: make([]TaskOutput, len(jobRun.Tasks)),
    64  	}
    65  	for _, task := range jobRun.Tasks {
    66  		jobRunOutput, err := w.Jobs.GetRunOutput(ctx, jobs.GetRunOutputRequest{
    67  			RunId: task.RunId,
    68  		})
    69  		if err != nil {
    70  			return nil, err
    71  		}
    72  		task := TaskOutput{TaskKey: task.TaskKey, Output: toRunOutput(jobRunOutput), EndTime: task.EndTime}
    73  		result.TaskOutputs = append(result.TaskOutputs, task)
    74  	}
    75  	return result, nil
    76  }