mirror of https://github.com/databricks/cli.git
77 lines
1.9 KiB
Go
77 lines
1.9 KiB
Go
package output
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
|
|
"github.com/databricks/databricks-sdk-go"
|
|
"github.com/databricks/databricks-sdk-go/service/jobs"
|
|
)
|
|
|
|
type TaskOutput struct {
|
|
TaskKey string
|
|
Output RunOutput
|
|
EndTime int64
|
|
}
|
|
|
|
type JobOutput struct {
|
|
// output for tasks with a non empty output
|
|
TaskOutputs []TaskOutput `json:"task_outputs"`
|
|
}
|
|
|
|
// Returns tasks output in text form sorted in execution order based on task end time
|
|
func (out *JobOutput) String() (string, error) {
|
|
if len(out.TaskOutputs) == 0 {
|
|
return "", nil
|
|
}
|
|
// When only one task, just return that output without any formatting
|
|
if len(out.TaskOutputs) == 1 {
|
|
for _, v := range out.TaskOutputs {
|
|
return v.Output.String()
|
|
}
|
|
}
|
|
result := strings.Builder{}
|
|
result.WriteString("Output:\n")
|
|
sort.Slice(out.TaskOutputs, func(i, j int) bool {
|
|
return out.TaskOutputs[i].EndTime < out.TaskOutputs[j].EndTime
|
|
})
|
|
for _, v := range out.TaskOutputs {
|
|
if v.Output == nil {
|
|
continue
|
|
}
|
|
taskString, err := v.Output.String()
|
|
if err != nil {
|
|
return "", nil
|
|
}
|
|
result.WriteString("=======\n")
|
|
result.WriteString(fmt.Sprintf("Task %s:\n", v.TaskKey))
|
|
result.WriteString(fmt.Sprintf("%s\n", taskString))
|
|
}
|
|
return result.String(), nil
|
|
}
|
|
|
|
func GetJobOutput(ctx context.Context, w *databricks.WorkspaceClient, runId int64) (*JobOutput, error) {
|
|
jobRun, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{
|
|
RunId: runId,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
result := &JobOutput{
|
|
TaskOutputs: make([]TaskOutput, len(jobRun.Tasks)),
|
|
}
|
|
for _, task := range jobRun.Tasks {
|
|
jobRunOutput, err := w.Jobs.GetRunOutput(ctx, jobs.GetRunOutputRequest{
|
|
RunId: task.RunId,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
task := TaskOutput{TaskKey: task.TaskKey, Output: toRunOutput(jobRunOutput), EndTime: task.EndTime}
|
|
result.TaskOutputs = append(result.TaskOutputs, task)
|
|
}
|
|
return result, nil
|
|
}
|