Refactor output and progress into separate packages in run (#335)

Tested manually that output and progress logging still works
This commit is contained in:
shreyas-goenka 2023-04-14 14:40:34 +02:00 committed by GitHub
parent b9c68b4bd5
commit e8018a7209
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 27 additions and 24 deletions

View File

@ -8,6 +8,8 @@ import (
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config/resources"
"github.com/databricks/bricks/bundle/run/output"
"github.com/databricks/bricks/bundle/run/progress"
"github.com/databricks/bricks/libs/cmdio"
"github.com/databricks/bricks/libs/log"
"github.com/databricks/databricks-sdk-go/retries"
@ -197,7 +199,7 @@ func logProgressCallback(ctx context.Context, progressLogger *cmdio.Logger) func
prevState = state
}
event := &JobProgressEvent{
event := &progress.JobProgressEvent{
Timestamp: time.Now(),
JobId: i.JobId,
RunId: i.RunId,
@ -214,7 +216,7 @@ func logProgressCallback(ctx context.Context, progressLogger *cmdio.Logger) func
}
}
func (r *jobRunner) Run(ctx context.Context, opts *Options) (RunOutput, error) {
func (r *jobRunner) Run(ctx context.Context, opts *Options) (output.RunOutput, error) {
jobID, err := strconv.ParseInt(r.job.ID, 10, 64)
if err != nil {
return nil, fmt.Errorf("job ID is not an integer: %s", r.job.ID)
@ -274,7 +276,7 @@ func (r *jobRunner) Run(ctx context.Context, opts *Options) (RunOutput, error) {
// The task completed successfully.
case jobs.RunResultStateSuccess:
log.Infof(ctx, "Run has completed successfully!")
return getJobOutput(ctx, r.bundle.WorkspaceClient(), *runId)
return output.GetJobOutput(ctx, r.bundle.WorkspaceClient(), *runId)
// The run was stopped after reaching the timeout.
case jobs.RunResultStateTimedout:

View File

@ -1,4 +1,4 @@
package run
package output
import (
"context"
@ -50,7 +50,7 @@ func (out *JobOutput) String() (string, error) {
return result.String(), nil
}
func getJobOutput(ctx context.Context, w *databricks.WorkspaceClient, runId int64) (*JobOutput, error) {
func GetJobOutput(ctx context.Context, w *databricks.WorkspaceClient, runId int64) (*JobOutput, error) {
jobRun, err := w.Jobs.GetRun(ctx, jobs.GetRun{
RunId: runId,
})

View File

@ -1,4 +1,4 @@
package run
package output
import (
"testing"

5
bundle/run/output/run.go Normal file
View File

@ -0,0 +1,5 @@
package output
type RunOutput interface {
String() (string, error)
}

View File

@ -1,4 +1,4 @@
package run
package output
import (
"encoding/json"

View File

@ -1,4 +1,4 @@
package run
package output
import (
"fmt"

View File

@ -8,7 +8,8 @@ import (
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config/resources"
"github.com/databricks/bricks/bundle/run/pipeline"
"github.com/databricks/bricks/bundle/run/output"
"github.com/databricks/bricks/bundle/run/progress"
"github.com/databricks/bricks/libs/cmdio"
"github.com/databricks/bricks/libs/flags"
"github.com/databricks/bricks/libs/log"
@ -136,7 +137,7 @@ type pipelineRunner struct {
pipeline *resources.Pipeline
}
func (r *pipelineRunner) Run(ctx context.Context, opts *Options) (RunOutput, error) {
func (r *pipelineRunner) Run(ctx context.Context, opts *Options) (output.RunOutput, error) {
var pipelineID = r.pipeline.ID
// Include resource key in logger.
@ -161,7 +162,7 @@ func (r *pipelineRunner) Run(ctx context.Context, opts *Options) (RunOutput, err
updateID := res.UpdateId
// setup progress logger and tracker to query events
updateTracker := pipeline.NewUpdateTracker(pipelineID, updateID, w)
updateTracker := progress.NewUpdateTracker(pipelineID, updateID, w)
progressLogger, ok := cmdio.FromContext(ctx)
if !ok {
return nil, fmt.Errorf("no progress logger found")
@ -172,7 +173,7 @@ func (r *pipelineRunner) Run(ctx context.Context, opts *Options) (RunOutput, err
}
// Log the pipeline update URL as soon as it is available.
progressLogger.Log(pipeline.NewUpdateUrlEvent(w.Config.Host, updateID, pipelineID))
progressLogger.Log(progress.NewUpdateUrlEvent(w.Config.Host, updateID, pipelineID))
// Poll update for completion and post status.
// Note: there is no "StartUpdateAndWait" wrapper for this API.

View File

@ -1,4 +1,4 @@
package run
package progress
import (
"strings"

View File

@ -1,4 +1,4 @@
package run
package progress
import (
"testing"

View File

@ -1,4 +1,4 @@
package pipeline
package progress
import (
"context"

View File

@ -1,4 +1,4 @@
package pipeline
package progress
import "fmt"

View File

@ -1,4 +1,4 @@
package pipeline
package progress
import (
"testing"

View File

@ -6,20 +6,15 @@ import (
"strings"
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/run/output"
)
// TODO: refactor this package into service specific implementations. Its
// getting bloated. (https://github.com/databricks/bricks/issues/282)
type key string
func (k key) Key() string {
return string(k)
}
type RunOutput interface {
String() (string, error)
}
// Runner defines the interface for a runnable resource (or workload).
type Runner interface {
// Key returns the fully qualified (unique) identifier for this runnable resource.
@ -27,7 +22,7 @@ type Runner interface {
Key() string
// Run the underlying worklow.
Run(ctx context.Context, opts *Options) (RunOutput, error)
Run(ctx context.Context, opts *Options) (output.RunOutput, error)
}
// Find locates a runner matching the specified argument.