diff --git a/acceptance/pipelines/datasets/flag-conflict/databricks.yml b/acceptance/pipelines/datasets/flag-conflict/databricks.yml new file mode 100644 index 00000000000..7e18f370db9 --- /dev/null +++ b/acceptance/pipelines/datasets/flag-conflict/databricks.yml @@ -0,0 +1,10 @@ +bundle: + name: test-pipeline-datasets-flag-conflict + +resources: + pipelines: + my_pipeline: + name: test-pipeline-datasets-flag-conflict + libraries: + - file: + path: pipeline_file.py diff --git a/acceptance/pipelines/datasets/flag-conflict/out.test.toml b/acceptance/pipelines/datasets/flag-conflict/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/pipelines/datasets/flag-conflict/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/pipelines/datasets/flag-conflict/output.txt b/acceptance/pipelines/datasets/flag-conflict/output.txt new file mode 100644 index 00000000000..3484054d224 --- /dev/null +++ b/acceptance/pipelines/datasets/flag-conflict/output.txt @@ -0,0 +1,3 @@ + +=== datasets rejects --force-dry-run and --no-dry-run together (before any bundle work) +Error: --force-dry-run and --no-dry-run cannot be used together diff --git a/acceptance/pipelines/datasets/flag-conflict/pipeline_file.py b/acceptance/pipelines/datasets/flag-conflict/pipeline_file.py new file mode 100644 index 00000000000..38d6db7748a --- /dev/null +++ b/acceptance/pipelines/datasets/flag-conflict/pipeline_file.py @@ -0,0 +1,6 @@ +import dlt + + +@dlt.table +def my_table(): + return spark.range(10) diff --git a/acceptance/pipelines/datasets/flag-conflict/script b/acceptance/pipelines/datasets/flag-conflict/script new file mode 100644 index 00000000000..189b9c30ae3 --- /dev/null +++ b/acceptance/pipelines/datasets/flag-conflict/script @@ -0,0 +1,2 @@ +title "datasets rejects --force-dry-run and --no-dry-run together (before any bundle work)\n" +musterr $CLI pipelines datasets --force-dry-run --no-dry-run diff --git a/acceptance/pipelines/datasets/force/databricks.yml b/acceptance/pipelines/datasets/force/databricks.yml new file mode 100644 index 00000000000..0cef6c7a781 --- /dev/null +++ b/acceptance/pipelines/datasets/force/databricks.yml @@ -0,0 +1,10 @@ +bundle: + name: test-pipeline-datasets-force + +resources: + pipelines: + my_pipeline: + name: test-pipeline-datasets-force + libraries: + - file: + path: pipeline_file.py diff --git a/acceptance/pipelines/datasets/force/out.test.toml b/acceptance/pipelines/datasets/force/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/pipelines/datasets/force/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/pipelines/datasets/force/output.txt b/acceptance/pipelines/datasets/force/output.txt new file mode 100644 index 00000000000..ec014fb2599 --- /dev/null +++ b/acceptance/pipelines/datasets/force/output.txt @@ -0,0 +1,24 @@ + +>>> [CLI] pipelines deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-pipeline-datasets-force/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! +View your pipeline my_pipeline here: [DATABRICKS_URL]/pipelines/[UUID]?w=[NUMID] + +=== datasets --force-dry-run always triggers a fresh dry-run (announces the billed cluster) +>>> [CLI] pipelines datasets --force-dry-run +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +Triggering a fresh dry-run (this starts a billed cluster). Press Ctrl-C to cancel. +Update URL: [DATABRICKS_URL]/#joblist/pipelines/[UUID]/updates/[UUID] + +Update ID: [UUID] +Name Type +main.demo.source MATERIALIZED_VIEW +main.demo.filtered MATERIALIZED_VIEW +main.demo.aggregated MATERIALIZED_VIEW +recent VIEW +main.demo.archive SINK diff --git a/acceptance/pipelines/datasets/force/pipeline_file.py b/acceptance/pipelines/datasets/force/pipeline_file.py new file mode 100644 index 00000000000..38d6db7748a --- /dev/null +++ b/acceptance/pipelines/datasets/force/pipeline_file.py @@ -0,0 +1,6 @@ +import dlt + + +@dlt.table +def my_table(): + return spark.range(10) diff --git a/acceptance/pipelines/datasets/force/script b/acceptance/pipelines/datasets/force/script new file mode 100644 index 00000000000..b82dc38d9c2 --- /dev/null +++ b/acceptance/pipelines/datasets/force/script @@ -0,0 +1,4 @@ +trace $CLI pipelines deploy + +title "datasets --force-dry-run always triggers a fresh dry-run (announces the billed cluster)" +trace $CLI pipelines datasets --force-dry-run diff --git a/acceptance/pipelines/datasets/happy/databricks.yml b/acceptance/pipelines/datasets/happy/databricks.yml new file mode 100644 index 00000000000..8fe7475f9fb --- /dev/null +++ b/acceptance/pipelines/datasets/happy/databricks.yml @@ -0,0 +1,10 @@ +bundle: + name: test-pipeline-datasets + +resources: + pipelines: + my_pipeline: + name: test-pipeline-datasets + libraries: + - file: + path: pipeline_file.py diff --git a/acceptance/pipelines/datasets/happy/out.test.toml b/acceptance/pipelines/datasets/happy/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/pipelines/datasets/happy/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/pipelines/datasets/happy/output.txt b/acceptance/pipelines/datasets/happy/output.txt new file mode 100644 index 00000000000..4df9fc9ac93 --- /dev/null +++ b/acceptance/pipelines/datasets/happy/output.txt @@ -0,0 +1,89 @@ + +>>> [CLI] pipelines deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-pipeline-datasets/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! +View your pipeline my_pipeline here: [DATABRICKS_URL]/pipelines/[UUID]?w=[NUMID] + +=== datasets: default triggers a dry-run, then lists the pipeline's datasets +>>> [CLI] pipelines datasets +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +No readable dry-run found; triggering one (this starts a billed cluster). Press Ctrl-C to cancel. +Update URL: [DATABRICKS_URL]/#joblist/pipelines/[UUID]/updates/[UUID] + +Update ID: [UUID] +Name Type +main.demo.source MATERIALIZED_VIEW +main.demo.filtered MATERIALIZED_VIEW +main.demo.aggregated MATERIALIZED_VIEW +recent VIEW +main.demo.archive SINK + +=== datasets with an explicit pipeline KEY (same result) +>>> [CLI] pipelines datasets my_pipeline +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +Name Type +main.demo.source MATERIALIZED_VIEW +main.demo.filtered MATERIALIZED_VIEW +main.demo.aggregated MATERIALIZED_VIEW +recent VIEW +main.demo.archive SINK + +=== datasets as JSON +>>> [CLI] pipelines datasets --output json +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +[ + { + "dataset_ref": "n1", + "name": "main.demo.source", + "full_name": "`main`.`demo`.`source`", + "dataset_type": "MATERIALIZED_VIEW" + }, + { + "dataset_ref": "n2", + "name": "main.demo.filtered", + "full_name": "`main`.`demo`.`filtered`", + "dataset_type": "MATERIALIZED_VIEW" + }, + { + "dataset_ref": "n3", + "name": "main.demo.aggregated", + "full_name": "`main`.`demo`.`aggregated`", + "dataset_type": "MATERIALIZED_VIEW" + }, + { + "dataset_ref": "n4", + "name": "recent", + "full_name": "", + "dataset_type": "VIEW" + }, + { + "dataset_ref": "s1", + "name": "", + "full_name": "main.demo.archive", + "dataset_type": "SINK" + } +] + +=== datasets --no-dry-run reads the existing dry-run (no re-trigger) +>>> [CLI] pipelines datasets --no-dry-run +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +Name Type +main.demo.source MATERIALIZED_VIEW +main.demo.filtered MATERIALIZED_VIEW +main.demo.aggregated MATERIALIZED_VIEW +recent VIEW +main.demo.archive SINK diff --git a/acceptance/pipelines/datasets/happy/pipeline_file.py b/acceptance/pipelines/datasets/happy/pipeline_file.py new file mode 100644 index 00000000000..38d6db7748a --- /dev/null +++ b/acceptance/pipelines/datasets/happy/pipeline_file.py @@ -0,0 +1,6 @@ +import dlt + + +@dlt.table +def my_table(): + return spark.range(10) diff --git a/acceptance/pipelines/datasets/happy/script b/acceptance/pipelines/datasets/happy/script new file mode 100644 index 00000000000..a3a83c3f7bc --- /dev/null +++ b/acceptance/pipelines/datasets/happy/script @@ -0,0 +1,13 @@ +trace $CLI pipelines deploy + +title "datasets: default triggers a dry-run, then lists the pipeline's datasets" +trace $CLI pipelines datasets + +title "datasets with an explicit pipeline KEY (same result)" +trace $CLI pipelines datasets my_pipeline + +title "datasets as JSON" +trace $CLI pipelines datasets --output json + +title "datasets --no-dry-run reads the existing dry-run (no re-trigger)" +trace $CLI pipelines datasets --no-dry-run diff --git a/acceptance/pipelines/datasets/no-dry-run-missing/databricks.yml b/acceptance/pipelines/datasets/no-dry-run-missing/databricks.yml new file mode 100644 index 00000000000..5c1f03a8a6e --- /dev/null +++ b/acceptance/pipelines/datasets/no-dry-run-missing/databricks.yml @@ -0,0 +1,10 @@ +bundle: + name: test-pipeline-datasets-no-dry-run-missing + +resources: + pipelines: + my_pipeline: + name: test-pipeline-datasets-no-dry-run-missing + libraries: + - file: + path: pipeline_file.py diff --git a/acceptance/pipelines/datasets/no-dry-run-missing/out.test.toml b/acceptance/pipelines/datasets/no-dry-run-missing/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/pipelines/datasets/no-dry-run-missing/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/pipelines/datasets/no-dry-run-missing/output.txt b/acceptance/pipelines/datasets/no-dry-run-missing/output.txt new file mode 100644 index 00000000000..7a751b253f2 --- /dev/null +++ b/acceptance/pipelines/datasets/no-dry-run-missing/output.txt @@ -0,0 +1,14 @@ + +>>> [CLI] pipelines deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-pipeline-datasets-no-dry-run-missing/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! +View your pipeline my_pipeline here: [DATABRICKS_URL]/pipelines/[UUID]?w=[NUMID] + +=== datasets --no-dry-run with no prior dry-run reports the graph as unavailable +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +Error: datasets for my_pipeline: no dry-run found; re-run without --no-dry-run to trigger one diff --git a/acceptance/pipelines/datasets/no-dry-run-missing/pipeline_file.py b/acceptance/pipelines/datasets/no-dry-run-missing/pipeline_file.py new file mode 100644 index 00000000000..38d6db7748a --- /dev/null +++ b/acceptance/pipelines/datasets/no-dry-run-missing/pipeline_file.py @@ -0,0 +1,6 @@ +import dlt + + +@dlt.table +def my_table(): + return spark.range(10) diff --git a/acceptance/pipelines/datasets/no-dry-run-missing/script b/acceptance/pipelines/datasets/no-dry-run-missing/script new file mode 100644 index 00000000000..75f0e5171b1 --- /dev/null +++ b/acceptance/pipelines/datasets/no-dry-run-missing/script @@ -0,0 +1,4 @@ +trace $CLI pipelines deploy + +title "datasets --no-dry-run with no prior dry-run reports the graph as unavailable\n" +musterr $CLI pipelines datasets --no-dry-run diff --git a/acceptance/pipelines/datasets/no-retrigger/databricks.yml b/acceptance/pipelines/datasets/no-retrigger/databricks.yml new file mode 100644 index 00000000000..dfdff1893de --- /dev/null +++ b/acceptance/pipelines/datasets/no-retrigger/databricks.yml @@ -0,0 +1,10 @@ +bundle: + name: test-pipeline-datasets-no-retrigger + +resources: + pipelines: + my_pipeline: + name: test-pipeline-datasets-no-retrigger + libraries: + - file: + path: pipeline_file.py diff --git a/acceptance/pipelines/datasets/no-retrigger/out.test.toml b/acceptance/pipelines/datasets/no-retrigger/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/pipelines/datasets/no-retrigger/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/pipelines/datasets/no-retrigger/output.txt b/acceptance/pipelines/datasets/no-retrigger/output.txt new file mode 100644 index 00000000000..9ccb09073d5 --- /dev/null +++ b/acceptance/pipelines/datasets/no-retrigger/output.txt @@ -0,0 +1,44 @@ + +>>> [CLI] pipelines deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-pipeline-datasets-no-retrigger/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! +View your pipeline my_pipeline here: [DATABRICKS_URL]/pipelines/[UUID]?w=[NUMID] + +>>> [CLI] pipelines datasets +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +No readable dry-run found; triggering one (this starts a billed cluster). Press Ctrl-C to cancel. +Update URL: [DATABRICKS_URL]/#joblist/pipelines/[UUID]/updates/[UUID] + +Update ID: [UUID] +Name Type +main.demo.source MATERIALIZED_VIEW +main.demo.filtered MATERIALIZED_VIEW +main.demo.aggregated MATERIALIZED_VIEW +recent VIEW +main.demo.archive SINK + +>>> [CLI] pipelines datasets +Recommendation: This command runs the last deployed version of the code + +If you've made local changes, run 'databricks pipelines deploy' first to ensure they are included. + +Name Type +main.demo.source MATERIALIZED_VIEW +main.demo.filtered MATERIALIZED_VIEW +main.demo.aggregated MATERIALIZED_VIEW +recent VIEW +main.demo.archive SINK + +=== only one dry-run was triggered across two reads +{ + "method": "POST", + "path": "/api/2.0/pipelines/[UUID]/updates", + "body": { + "validate_only": true + } +} diff --git a/acceptance/pipelines/datasets/no-retrigger/pipeline_file.py b/acceptance/pipelines/datasets/no-retrigger/pipeline_file.py new file mode 100644 index 00000000000..38d6db7748a --- /dev/null +++ b/acceptance/pipelines/datasets/no-retrigger/pipeline_file.py @@ -0,0 +1,6 @@ +import dlt + + +@dlt.table +def my_table(): + return spark.range(10) diff --git a/acceptance/pipelines/datasets/no-retrigger/script b/acceptance/pipelines/datasets/no-retrigger/script new file mode 100644 index 00000000000..61657a86302 --- /dev/null +++ b/acceptance/pipelines/datasets/no-retrigger/script @@ -0,0 +1,10 @@ +trace $CLI pipelines deploy +rm out.requests.txt + +# First call triggers a dry-run (POST .../updates); the second reuses it. print_requests.py +# excludes GETs, so the recorded POSTs show exactly one StartUpdate across both reads. +trace $CLI pipelines datasets +trace $CLI pipelines datasets + +title "only one dry-run was triggered across two reads\n" +print_requests.py //pipelines --sort diff --git a/acceptance/pipelines/datasets/no-retrigger/test.toml b/acceptance/pipelines/datasets/no-retrigger/test.toml new file mode 100644 index 00000000000..159efe02696 --- /dev/null +++ b/acceptance/pipelines/datasets/no-retrigger/test.toml @@ -0,0 +1 @@ +RecordRequests = true diff --git a/cmd/pipelines/commands.go b/cmd/pipelines/commands.go index da8bc06ace1..d2b375e863e 100644 --- a/cmd/pipelines/commands.go +++ b/cmd/pipelines/commands.go @@ -17,6 +17,7 @@ func Commands() []*cobra.Command { destroyCommand(), runCommand(), dryRunCommand(), + datasetsCommand(), historyCommand(), logsCommand(), openCommand(), diff --git a/cmd/pipelines/dataflow.go b/cmd/pipelines/dataflow.go new file mode 100644 index 00000000000..e9cbaea74d9 --- /dev/null +++ b/cmd/pipelines/dataflow.go @@ -0,0 +1,506 @@ +package pipelines + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "os/signal" + "strconv" + "strings" + "time" + + "github.com/databricks/cli/bundle/run/progress" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/pipelines" + "github.com/spf13/cobra" +) + +const ( + // gap between update-status polls while a triggered dry-run runs + pollInterval = time.Second + // grace for completed updated but no results + graceDelay = 2 * time.Second + // page size per entity endpoint + nodesPageSize = 20 + flowsPageSize = 20 + diagnosticsPageSize = 20 + // window for an existing dry-run + listUpdatesPageSize = 100 + + errCodeInvalidStateTransition = "INVALID_STATE_TRANSITION" + + severityError = "ERROR" + // dataset_type for a flattened sink + sinkNodeType = "SINK" + // messages for dry-run triggers + triggerNotice = "No readable dry-run found; triggering one (this starts a billed cluster). Press Ctrl-C to cancel." + forceTriggerNotice = "Triggering a fresh dry-run (this starts a billed cluster). Press Ctrl-C to cancel." +) + +var ( + errContinuous = errors.New("the dataflow graph is unavailable for continuous pipelines") + errNoDryRun = errors.New("no dry-run found; re-run without --no-dry-run to trigger one") + errActiveUpdate = errors.New("an update is already active for this pipeline") + errDryRunInProgress = errors.New("the latest dry-run is still in progress") +) + +// hand-rolled response shapes +type dagDataset struct { + Ref string `json:"dataset_ref"` + Name string `json:"name"` + FullName string `json:"full_name"` + DatasetType string `json:"dataset_type"` +} + +type dagSink struct { + Ref string `json:"sink_ref"` + Name string `json:"name"` + TableName string `json:"table_name"` +} + +// graph vertex: either Dataset or Sink +type dagNode struct { + Dataset *dagDataset `json:"dataset"` + Sink *dagSink `json:"sink"` +} + +type listNodesResponse struct { + Nodes []dagNode `json:"nodes"` + NextPageToken string `json:"next_page_token"` +} + +// node a diagnostic relates to; exactly one name field is set +type dagDiagnosticNode struct { + DatasetName string `json:"dataset_name"` + SinkName string `json:"sink_name"` + FlowName string `json:"flow_name"` +} + +type dagDiagnostic struct { + Severity string `json:"severity"` + Code string `json:"code"` + Message string `json:"message"` + DocumentURI string `json:"document_uri"` + Range struct { + Start struct { + Line int `json:"line"` + } `json:"start"` + } `json:"range"` + RelatedNodes []dagDiagnosticNode `json:"related_pipeline_nodes"` + Details struct { + // structured JVM error (error_class, sql_state) + Exception struct { + ErrorClass string `json:"error_class"` + SQLState string `json:"sql_state"` + } `json:"exception"` + } `json:"details"` +} + +type listDiagnosticsResponse struct { + Diagnostics []dagDiagnostic `json:"diagnostics"` + NextPageToken string `json:"next_page_token"` +} + +// dry-run mode flags shared by the pipeline-entities commands +type dagRunOpts struct { + forceDryRun bool + noDryRun bool +} + +// workspace operations resolveDryRun orchestrates +type updateDeps struct { + continuous func(ctx context.Context) (bool, error) + listUpdates func(ctx context.Context) ([]pipelines.UpdateInfo, error) + trigger func(ctx context.Context) (updateID string, err error) + poll func(ctx context.Context, updateID string) (pipelines.UpdateInfoState, error) + stop func(ctx context.Context) +} + +func isTerminal(s pipelines.UpdateInfoState) bool { + switch s { + case pipelines.UpdateInfoStateCompleted, pipelines.UpdateInfoStateFailed, pipelines.UpdateInfoStateCanceled: + return true + default: + return false + } +} + +// most recent update (by creation time) matching filter, if any +// nil filter matches all updates +func newestUpdate(updates []pipelines.UpdateInfo, filter func(pipelines.UpdateInfo) bool) (pipelines.UpdateInfo, bool) { + var best pipelines.UpdateInfo + found := false + for _, u := range updates { + if filter != nil && !filter(u) { + continue + } + if !found || u.CreationTime > best.CreationTime { + best, found = u, true + } + } + return best, found +} + +func activeUpdateError(u pipelines.UpdateInfo) error { + if u.ValidateOnly { + return fmt.Errorf("%w: a dry-run is in progress (update %s); re-run once it finishes", errActiveUpdate, u.UpdateId) + } + return fmt.Errorf("%w: an update is running (update %s); a dry-run cannot start while an update is active, retry once it completes", errActiveUpdate, u.UpdateId) +} + +// resolves validate-only update whose graph to read, triggers a fresh dry-run when needed +// returns the update id and terminal state. continuous/active-update return errors. +func resolveDryRun(ctx context.Context, d updateDeps, opts dagRunOpts) (id string, state pipelines.UpdateInfoState, err error) { + // stop a running dry-run if the command is interrupted/erorr etc. + signalCtx, stop := signal.NotifyContext(ctx, os.Interrupt) + defer stop() + triggered := false + defer func() { + if triggered && err != nil { + d.stop(ctx) + } + }() + + // Continuous pipelines cannot be dry-run --> we don't serve them + continuous, err := d.continuous(signalCtx) + if err != nil { + return "", "", err + } + if continuous { + return "", "", errContinuous + } + + triggerAndPoll := func(notice string) (string, pipelines.UpdateInfoState, error) { + cmdio.LogString(signalCtx, notice) + newID, err := d.trigger(signalCtx) + if err != nil { + return "", "", mapTriggerError(signalCtx, d, err) + } + triggered = true + st, err := d.poll(signalCtx, newID) + return newID, st, err + } + + if opts.forceDryRun { + return triggerAndPoll(forceTriggerNotice) + } + + updates, err := d.listUpdates(signalCtx) + if err != nil { + return "", "", err + } + + latest, found := newestUpdate(updates, func(u pipelines.UpdateInfo) bool { return u.ValidateOnly }) + + if found && latest.State == pipelines.UpdateInfoStateCompleted { + return latest.UpdateId, latest.State, nil + } + + if opts.noDryRun { + switch { + case found && !isTerminal(latest.State): + // in progress update + return "", "", fmt.Errorf("%w (update %s); re-run once it finishes", errDryRunInProgress, latest.UpdateId) + case found: + // failed/canceled dry-run: surface its diagnostics + return latest.UpdateId, latest.State, nil + default: + return "", "", errNoDryRun + } + } + + // active update + if u, ok := newestUpdate(updates, nil); ok && !isTerminal(u.State) { + return "", "", activeUpdateError(u) + } + // trigger fresh dry-run + return triggerAndPoll(triggerNotice) +} + +func mapTriggerError(ctx context.Context, d updateDeps, err error) error { + apiErr, ok := errors.AsType[*apierr.APIError](err) + if !ok || apiErr.ErrorCode != errCodeInvalidStateTransition { + return err + } + if updates, e := d.listUpdates(ctx); e == nil { + if u, ok := newestUpdate(updates, nil); ok && !isTerminal(u.State) { + return activeUpdateError(u) + } + } + return err +} + +// wires updateDeps to the workspace client for the given pipeline +func newUpdateDeps(w *databricks.WorkspaceClient, pipelineID string) updateDeps { + return updateDeps{ + continuous: func(ctx context.Context) (bool, error) { + resp, err := w.Pipelines.GetByPipelineId(ctx, pipelineID) + if err != nil { + return false, err + } + return resp.Spec != nil && resp.Spec.Continuous, nil + }, + listUpdates: func(ctx context.Context) ([]pipelines.UpdateInfo, error) { + resp, err := w.Pipelines.ListUpdates(ctx, pipelines.ListUpdatesRequest{ + PipelineId: pipelineID, + MaxResults: listUpdatesPageSize, + }) + if err != nil { + return nil, err + } + return resp.Updates, nil + }, + trigger: func(ctx context.Context) (string, error) { + res, err := w.Pipelines.StartUpdate(ctx, pipelines.StartUpdate{ + PipelineId: pipelineID, + ValidateOnly: true, + }) + if err != nil { + return "", err + } + return res.UpdateId, nil + }, + poll: func(ctx context.Context, updateID string) (pipelines.UpdateInfoState, error) { + return streamUpdateProgress(ctx, w, pipelineID, updateID) + }, + stop: func(ctx context.Context) { stopUpdate(ctx, w, pipelineID) }, + } +} + +// polls the update to a terminal state while streaming its progress events +// (update URL, then each update_progress/flow_progress line) to stderr +func streamUpdateProgress(ctx context.Context, w *databricks.WorkspaceClient, pipelineID, updateID string) (pipelines.UpdateInfoState, error) { + cmdio.Log(ctx, progress.NewPipelineUpdateUrlEvent(w.Config.Host, updateID, pipelineID)) + tracker := progress.NewUpdateTracker(pipelineID, updateID, w) + for { + + if events, err := tracker.Events(ctx); err != nil { + log.Warnf(ctx, "failed to fetch progress events for %s: %v", updateID, err) + } else { + for i := range events { + cmdio.Log(ctx, &events[i]) + } + } + + resp, err := w.Pipelines.GetUpdateByPipelineIdAndUpdateId(ctx, pipelineID, updateID) + if err != nil { + return "", err + } + if resp.Update != nil && isTerminal(resp.Update.State) { + cmdio.LogString(ctx, "Update ID: "+updateID) + return resp.Update.State, nil + } + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-time.After(pollInterval): + } + } +} + +// best-effort cancels a triggered dry-run +func stopUpdate(ctx context.Context, w *databricks.WorkspaceClient, pipelineID string) { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Minute) + defer cancel() + cmdio.LogString(cleanupCtx, "Stopping the dry-run that was triggered...") + if _, err := w.Pipelines.Stop(cleanupCtx, pipelines.StopRequest{PipelineId: pipelineID}); err != nil { + log.Warnf(cleanupCtx, "failed to stop the dry-run for %s (a billed cluster may still be running): %v", pipelineID, err) + } +} + +// issues GET path with page_size + page_token until the response has no next token +func fetchAllPages[R, T any](ctx context.Context, c *client.DatabricksClient, headers map[string]string, path, updateID string, pageSize int, items func(*R) []T, next func(*R) string) ([]T, error) { + var out []T + token := "" + for { + query := map[string]string{ + "update_id": updateID, + "page_size": strconv.Itoa(pageSize), + } + if token != "" { + query["page_token"] = token + } + var resp R + if err := c.Do(ctx, http.MethodGet, path, headers, nil, query, &resp); err != nil { + return nil, err + } + out = append(out, items(&resp)...) + token = next(&resp) + if token == "" { + return out, nil + } + } +} + +func graphPath(pipelineID, leaf string) string { + return fmt.Sprintf("/api/2.0/pipelines/%s/entities/%s", pipelineID, leaf) +} + +// returns the pipeline's datasets +func fetchDatasets(ctx context.Context, c *client.DatabricksClient, headers map[string]string, pipelineID, updateID string) ([]dagDataset, error) { + nodes, err := fetchNodes(ctx, c, headers, pipelineID, updateID) + if err != nil { + return nil, err + } + datasets := datasetsFromNodes(nodes) + if len(datasets) > 0 { + return datasets, nil + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(graceDelay): + } + nodes, err = fetchNodes(ctx, c, headers, pipelineID, updateID) + if err != nil { + return nil, err + } + return datasetsFromNodes(nodes), nil +} + +// returns the pipeline's graph nodes (datasets and sinks) for the update +func fetchNodes(ctx context.Context, c *client.DatabricksClient, headers map[string]string, pipelineID, updateID string) ([]dagNode, error) { + return fetchAllPages(ctx, c, headers, graphPath(pipelineID, "nodes"), updateID, nodesPageSize, + func(r *listNodesResponse) []dagNode { return r.Nodes }, + func(r *listNodesResponse) string { return r.NextPageToken }) +} + +// use name if present, otherwise full_name (sinks) +func displayName(d dagDataset) string { + if d.Name != "" { + return d.Name + } + return d.FullName +} + +// flattens a graph node into the dagDataset shape +func nodeToDataset(n dagNode) (dagDataset, bool) { + switch { + case n.Dataset != nil: + return *n.Dataset, true + case n.Sink != nil: + fullName := n.Sink.TableName + if fullName == "" { + fullName = n.Sink.Name + } + return dagDataset{Ref: n.Sink.Ref, FullName: fullName, DatasetType: sinkNodeType}, true + default: + return dagDataset{}, false + } +} + +// flattens every graph node into the dagDataset shape +func datasetsFromNodes(nodes []dagNode) []dagDataset { + datasets := make([]dagDataset, 0, len(nodes)) + for _, n := range nodes { + if d, ok := nodeToDataset(n); ok { + datasets = append(datasets, d) + } + } + return datasets +} + +func fetchDiagnostics(ctx context.Context, c *client.DatabricksClient, headers map[string]string, pipelineID, updateID string) ([]dagDiagnostic, error) { + return fetchAllPages(ctx, c, headers, graphPath(pipelineID, "diagnostics"), updateID, diagnosticsPageSize, + func(r *listDiagnosticsResponse) []dagDiagnostic { return r.Diagnostics }, + func(r *listDiagnosticsResponse) string { return r.NextPageToken }) +} + +// gets the diagnostics for a non-completed update +func nonCompletedError(ctx context.Context, c *client.DatabricksClient, headers map[string]string, pipelineID, updateID string, state pipelines.UpdateInfoState) error { + if state == pipelines.UpdateInfoStateCanceled { + return errors.New("dry-run was canceled") + } + diagnostics, err := fetchDiagnostics(ctx, c, headers, pipelineID, updateID) + if err != nil { + return fmt.Errorf("dry-run failed; could not fetch diagnostics: %w", err) + } + return fmt.Errorf("dry-run failed: %s", diagnosticsSummary(diagnostics)) +} + +// joins the error-severity diagnostics into a single line +func diagnosticsSummary(diagnostics []dagDiagnostic) string { + var parts []string + for _, d := range diagnostics { + if d.Severity != severityError { + continue + } + parts = append(parts, formatDiagnostic(d)) + } + if len(parts) == 0 { + return "no error diagnostics returned" + } + return strings.Join(parts, "; ") +} + +func formatDiagnostic(d dagDiagnostic) string { + msg := d.Message + if target := diagnosticTarget(d.RelatedNodes); target != "" { + msg = target + ": " + msg + } + // prefer the structured error identifiers (error_class [sql_state]) over the bare code + if id := exceptionID(d); id != "" { + msg = fmt.Sprintf("%s [%s]", msg, id) + } else if d.Code != "" { + msg = fmt.Sprintf("%s (%s)", msg, d.Code) + } + if d.DocumentURI != "" { + // the API reports zero-based lines; add 1 so it matches the user's 1-based editor + msg = fmt.Sprintf("%s at %s:%d", msg, d.DocumentURI, d.Range.Start.Line+1) + } + return msg +} + +// joins the structured exception's error_class and sql_state (either may be absent) into a single identifier +func exceptionID(d dagDiagnostic) string { + ex := d.Details.Exception + switch { + case ex.ErrorClass != "" && ex.SQLState != "": + return fmt.Sprintf("%s %s", ex.ErrorClass, ex.SQLState) + case ex.ErrorClass != "": + return ex.ErrorClass + case ex.SQLState != "": + return ex.SQLState + default: + return "" + } +} + +func diagnosticTarget(nodes []dagDiagnosticNode) string { + for _, n := range nodes { + switch { + case n.DatasetName != "": + return n.DatasetName + case n.SinkName != "": + return n.SinkName + case n.FlowName != "": + return n.FlowName + } + } + return "" +} + +// writes v as indented JSON to the command's stdout with a trailing newline +func renderJSON(cmd *cobra.Command, v any) error { + out, err := json.MarshalIndent(v, "", " ") + if err != nil { + return err + } + _, err = cmd.OutOrStdout().Write(append(out, '\n')) + return err +} + +// returns a non-nil slice so JSON output renders [] rather than null +func orEmpty[T any](s []T) []T { + if s == nil { + return []T{} + } + return s +} diff --git a/cmd/pipelines/dataflow_e2e_test.go b/cmd/pipelines/dataflow_e2e_test.go new file mode 100644 index 00000000000..3d2cf844ff8 --- /dev/null +++ b/cmd/pipelines/dataflow_e2e_test.go @@ -0,0 +1,242 @@ +package pipelines + +import ( + "bytes" + "context" + "encoding/json" + "strconv" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/pipelines" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// These tests drive the real datasets/lineage command flow over HTTP against the testserver fake. +// The client and the seeding accessor share e2eToken so they hit the same fake workspace. +const e2eToken = "e2etoken" + +func newDataflowServer(t *testing.T) (*testserver.Server, *databricks.WorkspaceClient) { + t.Helper() + server := testserver.New(t) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: e2eToken}) + require.NoError(t, err) + return server, w +} + +func createDataflowPipeline(ctx context.Context, t *testing.T, w *databricks.WorkspaceClient, continuous bool) string { + t.Helper() + resp, err := w.Pipelines.Create(ctx, pipelines.CreatePipeline{ + Name: "dag-test", + Serverless: true, + Continuous: continuous, + }) + require.NoError(t, err) + return resp.PipelineId +} + +// sampleGraph is a small graph: main.s.orders feeds both main.s.orders_by_date and main.s.big_orders. +func sampleGraph() *testserver.DataflowGraph { + ds := func(ref, fullName string) testserver.DataflowGraphNode { + return testserver.DataflowGraphNode{Dataset: &testserver.DataflowGraphDataset{ + DatasetRef: ref, + FullName: fullName, + DatasetType: "MATERIALIZED_VIEW", + }} + } + return &testserver.DataflowGraph{ + Nodes: []testserver.DataflowGraphNode{ + ds("n1", "main.s.orders"), + ds("n2", "main.s.orders_by_date"), + ds("n3", "main.s.big_orders"), + }, + Flows: []testserver.DataflowGraphFlow{ + {InputNodeRefs: []string{"n1"}, OutputNodeRef: "n2"}, + {InputNodeRefs: []string{"n1"}, OutputNodeRef: "n3"}, + }, + } +} + +func TestDatasetsE2EPaginates(t *testing.T) { + ctx, cmd, buf := renderCmd(t, flags.OutputJSON) + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, false) + + // More datasets than one page (nodesPageSize) so fetchAllPages must follow the page token. + nodes := make([]testserver.DataflowGraphNode, 60) + for i := range nodes { + nodes[i] = testserver.DataflowGraphNode{Dataset: &testserver.DataflowGraphDataset{ + DatasetRef: "n" + strconv.Itoa(i), + FullName: "main.s.t" + strconv.Itoa(i), + DatasetType: "MATERIALIZED_VIEW", + }} + } + server.Workspace(e2eToken).SetPipelineGraph(pipelineID, &testserver.DataflowGraph{Nodes: nodes}) + + require.NoError(t, runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{})) + + var got []dagDataset + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Len(t, got, 60) +} + +func TestDatasetsE2ERendersNamesCleanly(t *testing.T) { + ctx, cmd, buf := renderCmd(t, flags.OutputText) + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, false) + // real backend shapes: backticked full_name dataset, empty-full_name view, nameless sink + server.Workspace(e2eToken).SetPipelineGraph(pipelineID, &testserver.DataflowGraph{ + Nodes: []testserver.DataflowGraphNode{ + {Dataset: &testserver.DataflowGraphDataset{DatasetRef: "n1", Name: "main.s.orders", FullName: "`main`.`s`.`orders`", DatasetType: "MATERIALIZED_VIEW"}}, + {Dataset: &testserver.DataflowGraphDataset{DatasetRef: "n2", Name: "recent", DatasetType: "VIEW"}}, + {Sink: &testserver.DataflowGraphSink{SinkRef: "s1", TableName: "main.s.archive"}}, + }, + }) + + require.NoError(t, runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{})) + assert.Equal(t, "Name Type\nmain.s.orders MATERIALIZED_VIEW\nrecent VIEW\nmain.s.archive SINK\n", buf.String()) +} + +func TestDatasetsE2ETriggeredDryRunStreamsProgress(t *testing.T) { + // A triggered dry-run streams the update URL and terminal Update ID to stderr, leaving the + // dataset table on stdout. + stderr := &bytes.Buffer{} + cmd := &cobra.Command{} + out := flags.OutputText + cmd.Flags().Var(&out, "output", "") + stdout := &bytes.Buffer{} + cmd.SetOut(stdout) + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, stdout, stderr, "", "")) + + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, false) + server.Workspace(e2eToken).SetPipelineGraph(pipelineID, &testserver.DataflowGraph{ + Nodes: []testserver.DataflowGraphNode{ + {Dataset: &testserver.DataflowGraphDataset{DatasetRef: "n1", Name: "main.s.orders", DatasetType: "MATERIALIZED_VIEW"}}, + }, + }) + + require.NoError(t, runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{forceDryRun: true})) + progress := stderr.String() + assert.Contains(t, progress, "Update URL: ") + assert.Contains(t, progress, "/updates/") + assert.Contains(t, progress, "Update ID: ") + assert.Equal(t, "Name Type\nmain.s.orders MATERIALIZED_VIEW\n", stdout.String()) +} + +func TestDatasetsE2EFailedDryRunSurfacesDiagnostics(t *testing.T) { + ctx, cmd, _ := renderCmd(t, flags.OutputText) + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, false) + ws := server.Workspace(e2eToken) + ws.SeedPipelineUpdate(&testserver.PipelineUpdate{ + PipelineId: pipelineID, + UpdateId: "failed-dry-run", + State: pipelines.UpdateInfoStateFailed, + ValidateOnly: true, + CreationTime: 1, + }) + ws.SetPipelineGraph(pipelineID, &testserver.DataflowGraph{ + Diagnostics: []testserver.DataflowGraphDiagnostic{ + { + Severity: "ERROR", + Code: "TABLE_NOT_FOUND", + Message: "missing table", + RelatedNodes: []testserver.DataflowGraphDiagnosticNode{{DatasetName: "main.s.missing"}}, + }, + }, + }) + + err := runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{noDryRun: true}) + require.Error(t, err) + assert.Contains(t, err.Error(), "dry-run failed") + assert.Contains(t, err.Error(), "missing table") + assert.Contains(t, err.Error(), "TABLE_NOT_FOUND") +} + +func TestDatasetsE2EFailedDryRunSurfacesExceptionDetail(t *testing.T) { + ctx, cmd, _ := renderCmd(t, flags.OutputText) + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, false) + ws := server.Workspace(e2eToken) + ws.SeedPipelineUpdate(&testserver.PipelineUpdate{ + PipelineId: pipelineID, + UpdateId: "failed-dry-run", + State: pipelines.UpdateInfoStateFailed, + ValidateOnly: true, + CreationTime: 1, + }) + // The structured exception's error_class and sql_state must surface, since they pinpoint the + // failure where the top-level message is generic. + ws.SetPipelineGraph(pipelineID, &testserver.DataflowGraph{ + Diagnostics: []testserver.DataflowGraphDiagnostic{ + { + Severity: "ERROR", + Code: "TABLE_NOT_FOUND", + Message: "query failed", + Details: &testserver.DataflowGraphDiagnosticDetails{ + Exception: &testserver.DataflowGraphDiagnosticException{ + ErrorClass: "UNRESOLVED_COLUMN", + SQLState: "42703", + }, + }, + }, + }, + }) + + err := runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{noDryRun: true}) + require.Error(t, err) + assert.Contains(t, err.Error(), "UNRESOLVED_COLUMN 42703") +} + +func TestDatasetsE2EContinuousPipelineErrors(t *testing.T) { + ctx, cmd, _ := renderCmd(t, flags.OutputText) + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, true) + server.Workspace(e2eToken).SetPipelineGraph(pipelineID, sampleGraph()) + + err := runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{}) + assert.ErrorIs(t, err, errContinuous) +} + +func TestDatasetsE2EActiveUpdateErrors(t *testing.T) { + ctx, cmd, _ := renderCmd(t, flags.OutputText) + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, false) + server.Workspace(e2eToken).SeedPipelineUpdate(&testserver.PipelineUpdate{ + PipelineId: pipelineID, + UpdateId: "running", + State: pipelines.UpdateInfoStateRunning, + ValidateOnly: false, + CreationTime: 1, + }) + + err := runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{}) + require.ErrorIs(t, err, errActiveUpdate) + assert.Contains(t, err.Error(), "an update is running") +} + +func TestDatasetsE2ENoDryRunInProgressErrors(t *testing.T) { + ctx, cmd, _ := renderCmd(t, flags.OutputText) + server, w := newDataflowServer(t) + pipelineID := createDataflowPipeline(ctx, t, w, false) + // The newest validate-only update is still running; --no-dry-run must report it as in-progress, + // not fetch diagnostics and render "dry-run failed". + server.Workspace(e2eToken).SeedPipelineUpdate(&testserver.PipelineUpdate{ + PipelineId: pipelineID, + UpdateId: "running-dry-run", + State: pipelines.UpdateInfoStateRunning, + ValidateOnly: true, + CreationTime: 1, + }) + + err := runDatasets(ctx, cmd, w, pipelineID, "key", dagRunOpts{noDryRun: true}) + require.ErrorIs(t, err, errDryRunInProgress) + assert.NotContains(t, err.Error(), "failed") +} diff --git a/cmd/pipelines/dataflow_test.go b/cmd/pipelines/dataflow_test.go new file mode 100644 index 00000000000..f8b5eced9c5 --- /dev/null +++ b/cmd/pipelines/dataflow_test.go @@ -0,0 +1,367 @@ +package pipelines + +import ( + "context" + "errors" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/pipelines" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeDeps is an in-memory updateDeps for testing resolveDryRun without a workspace. +type fakeDeps struct { + continuous bool + continuousErr error + updates []pipelines.UpdateInfo + updatesErr error + triggerID string + triggerErr error + pollState pipelines.UpdateInfoState + pollErr error + continuousCalls int + listCalls int + triggered bool + stopped bool +} + +func (f *fakeDeps) deps() updateDeps { + return updateDeps{ + continuous: func(ctx context.Context) (bool, error) { + f.continuousCalls++ + return f.continuous, f.continuousErr + }, + listUpdates: func(ctx context.Context) ([]pipelines.UpdateInfo, error) { + f.listCalls++ + return f.updates, f.updatesErr + }, + trigger: func(ctx context.Context) (string, error) { + f.triggered = true + return f.triggerID, f.triggerErr + }, + poll: func(ctx context.Context, updateID string) (pipelines.UpdateInfoState, error) { + return f.pollState, f.pollErr + }, + stop: func(ctx context.Context) { f.stopped = true }, + } +} + +func upd(id string, state pipelines.UpdateInfoState, validateOnly bool, created int64) pipelines.UpdateInfo { + return pipelines.UpdateInfo{UpdateId: id, State: state, ValidateOnly: validateOnly, CreationTime: created} +} + +func TestResolveDryRun(t *testing.T) { + completed := pipelines.UpdateInfoStateCompleted + cases := []struct { + name string + opts dagRunOpts + fake fakeDeps + wantID string + wantState pipelines.UpdateInfoState + wantErrIs error + wantErrSub string + wantTriggered bool + }{ + { + name: "default reads the latest completed dry-run", + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u2", completed, true, 2), upd("u1", completed, false, 1)}}, + wantID: "u2", + wantState: completed, + }, + { + name: "default triggers when no dry-run exists", + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u1", completed, false, 1)}, triggerID: "u9", pollState: completed}, + wantID: "u9", + wantState: completed, + wantTriggered: true, + }, + { + name: "default triggers when the update list is empty", + fake: fakeDeps{updates: nil, triggerID: "u9", pollState: completed}, + wantID: "u9", + wantState: completed, + wantTriggered: true, + }, + { + name: "default collapses a failed dry-run into a fresh trigger", + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u2", pipelines.UpdateInfoStateFailed, true, 2)}, triggerID: "u9", pollState: completed}, + wantID: "u9", + wantState: completed, + wantTriggered: true, + }, + { + name: "default rejects an in-flight dry-run", + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, true, 3), upd("u1", completed, true, 1)}}, + wantErrIs: errActiveUpdate, + wantErrSub: "dry-run is in progress", + }, + { + name: "default rejects an in-flight real run", + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, false, 3)}}, + wantErrIs: errActiveUpdate, + wantErrSub: "an update is running", + }, + { + name: "default reads a completed dry-run even when a newer real run is active", + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, false, 3), upd("u2", completed, true, 2)}}, + wantID: "u2", + wantState: completed, + }, + { + name: "default rejects a continuous pipeline", + fake: fakeDeps{continuous: true}, + wantErrIs: errContinuous, + }, + { + // Continuous pipelines cannot be dry-run --> we don't serve them + name: "default rejects a continuous pipeline with a running update", + fake: fakeDeps{continuous: true, updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, false, 3)}}, + wantErrIs: errContinuous, + }, + { + name: "no-dry-run reads the latest completed dry-run", + opts: dagRunOpts{noDryRun: true}, + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u2", completed, true, 2)}}, + wantID: "u2", + wantState: completed, + }, + { + name: "no-dry-run returns a failed dry-run for diagnostics", + opts: dagRunOpts{noDryRun: true}, + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u2", pipelines.UpdateInfoStateFailed, true, 2)}}, + wantID: "u2", + wantState: pipelines.UpdateInfoStateFailed, + }, + { + name: "no-dry-run reports an in-progress dry-run instead of mislabeling it failed", + opts: dagRunOpts{noDryRun: true}, + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, true, 3)}}, + wantErrIs: errDryRunInProgress, + wantErrSub: "still in progress", + }, + { + name: "no-dry-run errors when no dry-run exists", + opts: dagRunOpts{noDryRun: true}, + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u1", completed, false, 1)}}, + wantErrIs: errNoDryRun, + }, + { + // Continuous pipelines cannot be dry-run --> we don't serve them + name: "no-dry-run rejects a continuous pipeline", + opts: dagRunOpts{noDryRun: true}, + fake: fakeDeps{continuous: true, updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, true, 3)}}, + wantErrIs: errContinuous, + }, + { + name: "no-dry-run is exempt from the active-update check", + opts: dagRunOpts{noDryRun: true}, + fake: fakeDeps{updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, false, 3), upd("u2", completed, true, 2)}}, + wantID: "u2", + wantState: completed, + }, + { + name: "force triggers a fresh dry-run", + opts: dagRunOpts{forceDryRun: true}, + fake: fakeDeps{triggerID: "u9", pollState: completed}, + wantID: "u9", + wantState: completed, + wantTriggered: true, + }, + { + name: "force rejects a continuous pipeline", + opts: dagRunOpts{forceDryRun: true}, + fake: fakeDeps{continuous: true}, + wantErrIs: errContinuous, + }, + { + name: "force returns a failed state for diagnostics", + opts: dagRunOpts{forceDryRun: true}, + fake: fakeDeps{triggerID: "u9", pollState: pipelines.UpdateInfoStateFailed}, + wantID: "u9", + wantState: pipelines.UpdateInfoStateFailed, + wantTriggered: true, + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + f := c.fake + id, state, err := resolveDryRun(ctx, f.deps(), c.opts) + + if c.wantErrIs != nil || c.wantErrSub != "" { + require.Error(t, err) + if c.wantErrIs != nil { + assert.ErrorIs(t, err, c.wantErrIs) + } + if c.wantErrSub != "" { + assert.Contains(t, err.Error(), c.wantErrSub) + } + } else { + require.NoError(t, err) + assert.Equal(t, c.wantID, id) + assert.Equal(t, c.wantState, state) + } + assert.Equal(t, c.wantTriggered, f.triggered) + }) + } +} + +func TestResolveDryRunForceSkipsListButChecksContinuous(t *testing.T) { + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + f := fakeDeps{triggerID: "u9", pollState: pipelines.UpdateInfoStateCompleted} + _, _, err := resolveDryRun(ctx, f.deps(), dagRunOpts{forceDryRun: true}) + require.NoError(t, err) + assert.Equal(t, 0, f.listCalls, "force must not list updates") + assert.Equal(t, 1, f.continuousCalls, "force must still reject continuous pipelines") +} + +func TestResolveDryRunChecksContinuousBeforeListing(t *testing.T) { + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + // Continuous pipelines cannot be dry-run --> we don't serve them + f := fakeDeps{continuous: true, updates: []pipelines.UpdateInfo{upd("u2", pipelines.UpdateInfoStateCompleted, true, 2)}} + _, _, err := resolveDryRun(ctx, f.deps(), dagRunOpts{noDryRun: true}) + require.ErrorIs(t, err, errContinuous) + assert.Equal(t, 1, f.continuousCalls) + assert.Equal(t, 0, f.listCalls, "continuous pipelines must be rejected before listing updates") +} + +func TestResolveDryRunStopsTriggeredDryRunOnInterrupt(t *testing.T) { + base, _ := cmdio.NewTestContextWithStdout(t.Context()) + ctx, cancel := context.WithCancel(base) + cancel() // simulate Ctrl-C / timeout + + f := fakeDeps{triggerID: "u9", pollErr: context.Canceled} + _, _, err := resolveDryRun(ctx, f.deps(), dagRunOpts{forceDryRun: true}) + require.Error(t, err) + assert.True(t, f.triggered) + assert.True(t, f.stopped, "an interrupted, triggered dry-run must be stopped") +} + +func TestResolveDryRunStopsTriggeredDryRunOnPollError(t *testing.T) { + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + // Polling fails with a non-cancellation error (e.g. a 500 or auth expiry) after we triggered: + // we never confirmed a terminal state, so the billed cluster must still be stopped. + f := fakeDeps{triggerID: "u9", pollErr: errors.New("boom")} + _, _, err := resolveDryRun(ctx, f.deps(), dagRunOpts{forceDryRun: true}) + require.Error(t, err) + assert.True(t, f.triggered) + assert.True(t, f.stopped, "a triggered dry-run must be stopped when polling errors out") +} + +func TestResolveDryRunDoesNotStopOnTerminalFailure(t *testing.T) { + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + // A FAILED dry-run is terminal: its cluster already stopped, so we must NOT issue a stop. + f := fakeDeps{triggerID: "u9", pollState: pipelines.UpdateInfoStateFailed} + _, state, err := resolveDryRun(ctx, f.deps(), dagRunOpts{forceDryRun: true}) + require.NoError(t, err) + assert.Equal(t, pipelines.UpdateInfoStateFailed, state) + assert.True(t, f.triggered) + assert.False(t, f.stopped, "a terminal (FAILED) dry-run must not be stopped") +} + +func TestResolveDryRunMapsActiveUpdateTriggerError(t *testing.T) { + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + // StartUpdate is rejected because a real run started in the race window. + f := fakeDeps{ + triggerErr: &apierr.APIError{ErrorCode: errCodeInvalidStateTransition, Message: "active update"}, + updates: []pipelines.UpdateInfo{upd("u3", pipelines.UpdateInfoStateRunning, false, 3)}, + } + _, _, err := resolveDryRun(ctx, f.deps(), dagRunOpts{forceDryRun: true}) + require.Error(t, err) + assert.ErrorIs(t, err, errActiveUpdate) + assert.Contains(t, err.Error(), "an update is running") +} + +func TestResolveDryRunPassesThroughNonActiveTriggerError(t *testing.T) { + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + boom := errors.New("boom") + f := fakeDeps{triggerErr: boom} + _, _, err := resolveDryRun(ctx, f.deps(), dagRunOpts{forceDryRun: true}) + assert.ErrorIs(t, err, boom) +} + +func TestIsTerminal(t *testing.T) { + for _, s := range []pipelines.UpdateInfoState{pipelines.UpdateInfoStateCompleted, pipelines.UpdateInfoStateFailed, pipelines.UpdateInfoStateCanceled} { + assert.True(t, isTerminal(s), string(s)) + } + for _, s := range []pipelines.UpdateInfoState{pipelines.UpdateInfoStateRunning, pipelines.UpdateInfoStateWaitingForResources, pipelines.UpdateInfoStateQueued, ""} { + assert.False(t, isTerminal(s), string(s)) + } +} + +func TestNewestUpdate(t *testing.T) { + updates := []pipelines.UpdateInfo{ + upd("u1", pipelines.UpdateInfoStateCompleted, false, 1), + upd("u3", pipelines.UpdateInfoStateCompleted, true, 3), + upd("u2", pipelines.UpdateInfoStateCompleted, true, 2), + } + newest, ok := newestUpdate(updates, nil) + require.True(t, ok) + assert.Equal(t, "u3", newest.UpdateId) + + dryRun, ok := newestUpdate(updates, func(u pipelines.UpdateInfo) bool { return u.ValidateOnly }) + require.True(t, ok) + assert.Equal(t, "u3", dryRun.UpdateId) + + _, ok = newestUpdate(nil, nil) + assert.False(t, ok) +} + +func TestActiveUpdateError(t *testing.T) { + assert.Contains(t, activeUpdateError(upd("u1", pipelines.UpdateInfoStateRunning, true, 1)).Error(), "dry-run is in progress") + assert.Contains(t, activeUpdateError(upd("u1", pipelines.UpdateInfoStateRunning, false, 1)).Error(), "an update is running") +} + +func TestDiagnosticsSummary(t *testing.T) { + assert.Equal(t, "no error diagnostics returned", diagnosticsSummary(nil)) + + // Only ERROR-severity diagnostics are summarized; WARNING and INFORMATION are dropped. + diags := []dagDiagnostic{ + {Severity: "WARNING", Message: "ignored warning"}, + {Severity: "INFORMATION", Message: "ignored insight"}, + {Severity: severityError, Code: "TABLE_NOT_FOUND", Message: "missing table", RelatedNodes: []dagDiagnosticNode{{DatasetName: "main.s.t"}}}, + } + diags[2].Range.Start.Line = 12 + diags[2].DocumentURI = "file:///a.py" + // The zero-based line 12 renders as 1-based 13. + assert.Equal(t, "main.s.t: missing table (TABLE_NOT_FOUND) at file:///a.py:13", diagnosticsSummary(diags)) +} + +func TestFormatDiagnosticExceptionDetail(t *testing.T) { + // error_class + sql_state take precedence over the bare code and pinpoint the failure. + var d dagDiagnostic + d.Message = "boom" + d.Code = "TABLE_NOT_FOUND" + d.Details.Exception.ErrorClass = "UNRESOLVED_COLUMN" + d.Details.Exception.SQLState = "42703" + assert.Equal(t, "boom [UNRESOLVED_COLUMN 42703]", formatDiagnostic(d)) + + // Either field alone still surfaces; empty when neither is set. + var classOnly dagDiagnostic + classOnly.Details.Exception.ErrorClass = "UNRESOLVED_COLUMN" + assert.Equal(t, "UNRESOLVED_COLUMN", exceptionID(classOnly)) + var stateOnly dagDiagnostic + stateOnly.Details.Exception.SQLState = "42703" + assert.Equal(t, "42703", exceptionID(stateOnly)) + assert.Empty(t, exceptionID(dagDiagnostic{})) +} + +func TestNonCompletedErrorCanceled(t *testing.T) { + // A canceled dry-run has no diagnostics, so it must report as canceled (not "failed"). It + // short-circuits before any diagnostics fetch, so a nil client is fine here. + err := nonCompletedError(t.Context(), nil, nil, "pid", "uid", pipelines.UpdateInfoStateCanceled) + require.Error(t, err) + assert.Contains(t, err.Error(), "canceled") + assert.NotContains(t, err.Error(), "failed") +} + +func TestDiagnosticTarget(t *testing.T) { + assert.Equal(t, "ds", diagnosticTarget([]dagDiagnosticNode{{DatasetName: "ds"}})) + assert.Equal(t, "sink", diagnosticTarget([]dagDiagnosticNode{{SinkName: "sink"}})) + assert.Equal(t, "flow", diagnosticTarget([]dagDiagnosticNode{{FlowName: "flow"}})) + assert.Empty(t, diagnosticTarget(nil)) +} diff --git a/cmd/pipelines/datasets.go b/cmd/pipelines/datasets.go new file mode 100644 index 00000000000..01484b1dd0c --- /dev/null +++ b/cmd/pipelines/datasets.go @@ -0,0 +1,122 @@ +package pipelines + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/databricks/cli/cmd/bundle/utils" + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/logdiag" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/pipelines" + "github.com/spf13/cobra" +) + +func datasetsCommand() *cobra.Command { + var forceDryRun bool + var noDryRun bool + var timeout time.Duration + cmd := &cobra.Command{ + Use: "datasets [KEY]", + Short: "List the datasets a pipeline defines", + Long: `List the datasets a pipeline defines. + +By default this reads the latest dry-run's already-computed graph; if no dry-run exists yet it +triggers one and waits for it. Use --force-dry-run to always +trigger a fresh dry-run, or --no-dry-run to never trigger one (reports the graph as unavailable +when none exists). + +KEY is the pipeline's key in the bundle; it is optional if the bundle defines a +single pipeline.`, + Args: root.MaximumNArgs(1), + Hidden: true, + } + cmd.Flags().BoolVar(&forceDryRun, "force-dry-run", false, "Always trigger a fresh dry-run (spins up a billed cluster) instead of reading the latest.") + cmd.Flags().BoolVar(&noDryRun, "no-dry-run", false, "Never trigger a dry-run; report the graph as unavailable if none exists.") + cmd.Flags().DurationVar(&timeout, "timeout", 20*time.Minute, "Maximum time to wait for the graph.") + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := logdiag.InitContext(cmd.Context()) + cmd.SetContext(ctx) + + if forceDryRun && noDryRun { + return errors.New("--force-dry-run and --no-dry-run cannot be used together") + } + + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + cmd.SetContext(ctx) + + b, err := utils.ProcessBundle(cmd, utils.ProcessOptions{ErrorOnEmptyState: true, SkipInitContext: true}) + if err != nil { + return err + } + suggestPipelineDeploy(ctx, cmd) + key, _, err := resolveRunArgument(ctx, b, args) + if err != nil { + return err + } + pipelineID, err := resolvePipelineIdFromKey(ctx, b, key) + if err != nil { + return err + } + + w := b.WorkspaceClient(ctx) + return runDatasets(ctx, cmd, w, pipelineID, key, dagRunOpts{forceDryRun: forceDryRun, noDryRun: noDryRun}) + } + return cmd +} + +// resolves the dry-run graph for the pipeline and renders its datasets +func runDatasets(ctx context.Context, cmd *cobra.Command, w *databricks.WorkspaceClient, pipelineID, key string, opts dagRunOpts) error { + apiClient, err := client.New(w.Config) + if err != nil { + return fmt.Errorf("create API client: %w", err) + } + headers := auth.WorkspaceIDHeaders(w.Config) + + updateID, state, err := resolveDryRun(ctx, newUpdateDeps(w, pipelineID), opts) + if err != nil { + return fmt.Errorf("datasets for %s: %w", key, err) + } + if state != pipelines.UpdateInfoStateCompleted { + return nonCompletedError(ctx, apiClient, headers, pipelineID, updateID, state) + } + + datasets, err := fetchDatasets(ctx, apiClient, headers, pipelineID, updateID) + if err != nil { + return fmt.Errorf("fetch datasets for %s: %w", key, err) + } + return renderDatasets(ctx, cmd, datasets) +} + +// datasetRow is the display projection of a node: its user-visible name and type. +type datasetRow struct { + Name string + Type string +} + +func datasetRows(datasets []dagDataset) []datasetRow { + rows := make([]datasetRow, len(datasets)) + for i, d := range datasets { + rows[i] = datasetRow{Name: displayName(d), Type: d.DatasetType} + } + return rows +} + +func renderDatasets(ctx context.Context, cmd *cobra.Command, datasets []dagDataset) error { + switch root.OutputType(cmd) { + case flags.OutputText: + return cmdio.RenderWithTemplate(ctx, datasetRows(datasets), "", datasetsTemplate) + case flags.OutputJSON: + return renderJSON(cmd, orEmpty(datasets)) + default: + return fmt.Errorf("unknown output type %s", root.OutputType(cmd)) + } +} diff --git a/cmd/pipelines/datasets_test.go b/cmd/pipelines/datasets_test.go new file mode 100644 index 00000000000..234e6623aaa --- /dev/null +++ b/cmd/pipelines/datasets_test.go @@ -0,0 +1,74 @@ +package pipelines + +import ( + "bytes" + "context" + "io" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// builds a command plus a cmdio context that captures stdout (results) to the returned +// buffer, mirroring how root wires cmdio to the command in production. +func renderCmd(t *testing.T, out flags.Output) (context.Context, *cobra.Command, *bytes.Buffer) { + t.Helper() + cmd := &cobra.Command{} + value := out + cmd.Flags().Var(&value, "output", "") + buf := &bytes.Buffer{} + cmd.SetOut(buf) + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, nil, buf, io.Discard, "", "")) + return ctx, cmd, buf +} + +func TestRenderDatasets(t *testing.T) { + datasets := []dagDataset{ + // dataset renders by name (not backticked full_name); view by name; sink by full_name + {Name: "main.s.a", FullName: "`main`.`s`.`a`", DatasetType: "MATERIALIZED_VIEW"}, + {Name: "recent", DatasetType: "VIEW"}, + {FullName: "main.s.archive", DatasetType: "SINK"}, + } + + t.Run("text", func(t *testing.T) { + ctx, cmd, buf := renderCmd(t, flags.OutputText) + require.NoError(t, renderDatasets(ctx, cmd, datasets)) + assert.Equal(t, "Name Type\nmain.s.a MATERIALIZED_VIEW\nrecent VIEW\nmain.s.archive SINK\n", buf.String()) + }) + t.Run("text empty", func(t *testing.T) { + ctx, cmd, buf := renderCmd(t, flags.OutputText) + require.NoError(t, renderDatasets(ctx, cmd, nil)) + assert.Equal(t, "(none)\n", buf.String()) + }) + t.Run("json empty renders an array", func(t *testing.T) { + ctx, cmd, buf := renderCmd(t, flags.OutputJSON) + require.NoError(t, renderDatasets(ctx, cmd, nil)) + assert.Equal(t, "[]\n", buf.String()) + }) +} + +func TestDatasetsFromNodes(t *testing.T) { + nodes := []dagNode{ + {Dataset: &dagDataset{Ref: "n1", Name: "main.s.a", FullName: "`main`.`s`.`a`", DatasetType: "STREAMING_TABLE"}}, + {Sink: &dagSink{Ref: "s1", TableName: "main.s.archive"}}, + // a Kafka sink has no table_name, so its identifier arrives in name + {Sink: &dagSink{Ref: "s2", Name: "events_kafka_sink"}}, + } + // datasets and sinks are both listed; a sink becomes a SINK-typed row + got := datasetsFromNodes(nodes) + require.Len(t, got, 3) + assert.Equal(t, "main.s.a", got[0].Name) + assert.Equal(t, dagDataset{Ref: "s1", FullName: "main.s.archive", DatasetType: sinkNodeType}, got[1]) + assert.Equal(t, dagDataset{Ref: "s2", FullName: "events_kafka_sink", DatasetType: sinkNodeType}, got[2]) +} + +func TestDisplayName(t *testing.T) { + assert.Equal(t, "main.s.a", displayName(dagDataset{Name: "main.s.a", FullName: "`main`.`s`.`a`"})) + assert.Equal(t, "recent", displayName(dagDataset{Name: "recent"})) + // a sink has no name, so it falls back to full_name + assert.Equal(t, "main.s.archive", displayName(dagDataset{FullName: "main.s.archive"})) +} diff --git a/cmd/pipelines/templates.go b/cmd/pipelines/templates.go index 3c471b366bb..cc8800159d4 100644 --- a/cmd/pipelines/templates.go +++ b/cmd/pipelines/templates.go @@ -72,6 +72,12 @@ Pipeline configurations for this update: {{- end }} ` +// for datasets command +const datasetsTemplate = `{{if .}}Name Type +{{range .}}{{.Name}} {{.Type}} +{{end}}{{else}}(none) +{{end}}` + // progressEventsTemplate is the template for displaying progress events const progressEventsTemplate = `{{- if .ProgressEvents }} {{ printf "%-25s %s\n" "Run Phase" "Duration" }} diff --git a/libs/testserver/fake_workspace.go b/libs/testserver/fake_workspace.go index df9499c7e91..1e3488673f3 100644 --- a/libs/testserver/fake_workspace.go +++ b/libs/testserver/fake_workspace.go @@ -170,7 +170,8 @@ type FakeWorkspace struct { JobRuns map[int64]jobs.Run JobRunOutputs map[int64]jobs.RunOutput Pipelines map[string]pipelines.GetPipelineResponse - PipelineUpdates map[string]bool + PipelineUpdates map[string]*PipelineUpdate + PipelineGraphs map[string]*DataflowGraph Monitors map[string]catalog.MonitorInfo Apps map[string]apps.App Schemas map[string]catalog.SchemaInfo @@ -336,7 +337,8 @@ func NewFakeWorkspace(url, token string) *FakeWorkspace { JobRunOutputs: map[int64]jobs.RunOutput{}, Grants: map[string][]catalog.PrivilegeAssignment{}, Pipelines: map[string]pipelines.GetPipelineResponse{}, - PipelineUpdates: map[string]bool{}, + PipelineUpdates: map[string]*PipelineUpdate{}, + PipelineGraphs: map[string]*DataflowGraph{}, Monitors: map[string]catalog.MonitorInfo{}, Apps: map[string]apps.App{}, Catalogs: map[string]catalog.CatalogInfo{}, diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index a51a13b4afe..846c1250d23 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -388,7 +388,23 @@ func AddDefaultHandlers(server *Server) { }) server.Handle("POST", "/api/2.0/pipelines/{pipeline_id}/updates", func(req Request) any { - return req.Workspace.PipelineStartUpdate(req.Vars["pipeline_id"]) + return req.Workspace.PipelineStartUpdate(req, req.Vars["pipeline_id"]) + }) + + server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}/updates", func(req Request) any { + return req.Workspace.PipelineListUpdates(req, req.Vars["pipeline_id"]) + }) + + server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}/entities/nodes", func(req Request) any { + return req.Workspace.PipelineDataflowGraph(req, req.Vars["pipeline_id"], "nodes") + }) + + server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}/entities/flows", func(req Request) any { + return req.Workspace.PipelineDataflowGraph(req, req.Vars["pipeline_id"], "flows") + }) + + server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}/entities/diagnostics", func(req Request) any { + return req.Workspace.PipelineDataflowGraph(req, req.Vars["pipeline_id"], "diagnostics") }) server.Handle("GET", "/api/2.0/pipelines/{pipeline_id}/events", func(req Request) any { diff --git a/libs/testserver/pipelines.go b/libs/testserver/pipelines.go index 3512c8e5f5f..69343f0cef2 100644 --- a/libs/testserver/pipelines.go +++ b/libs/testserver/pipelines.go @@ -1,13 +1,84 @@ package testserver import ( + "cmp" "encoding/json" "fmt" + "slices" + "strconv" "strings" "github.com/databricks/databricks-sdk-go/service/pipelines" ) +// dataflowGraphMaxPageSize is the per-leaf page-size cap the backend enforces: it rejects (does not +// clamp) a larger page_size with a 400, so the fake does the same. +var dataflowGraphMaxPageSize = map[string]int{"nodes": 50, "flows": 50, "diagnostics": 20} + +// PipelineUpdate is a stored pipeline update. StartUpdate seeds one in a terminal COMPLETED state; +// tests seed failed or in-flight updates via SeedPipelineUpdate. +type PipelineUpdate struct { + PipelineId string + UpdateId string + State pipelines.UpdateInfoState + ValidateOnly bool + CreationTime int64 +} + +// DataflowGraph is the dry-run graph a pipeline exposes via the pipeline entities endpoints. It has +// no create API, so tests seed it with SetPipelineGraph. +type DataflowGraph struct { + Nodes []DataflowGraphNode + Flows []DataflowGraphFlow + Diagnostics []DataflowGraphDiagnostic +} + +type DataflowGraphNode struct { + Dataset *DataflowGraphDataset `json:"dataset,omitempty"` + Sink *DataflowGraphSink `json:"sink,omitempty"` +} + +type DataflowGraphDataset struct { + DatasetRef string `json:"dataset_ref,omitempty"` + Name string `json:"name,omitempty"` + FullName string `json:"full_name,omitempty"` + DatasetType string `json:"dataset_type,omitempty"` +} + +type DataflowGraphSink struct { + SinkRef string `json:"sink_ref,omitempty"` + Name string `json:"name,omitempty"` + TableName string `json:"table_name,omitempty"` +} + +type DataflowGraphFlow struct { + InputNodeRefs []string `json:"input_node_refs,omitempty"` + OutputNodeRef string `json:"output_node_ref,omitempty"` +} + +type DataflowGraphDiagnostic struct { + Severity string `json:"severity,omitempty"` + Code string `json:"code,omitempty"` + Message string `json:"message,omitempty"` + RelatedNodes []DataflowGraphDiagnosticNode `json:"related_pipeline_nodes,omitempty"` + Details *DataflowGraphDiagnosticDetails `json:"details,omitempty"` +} + +type DataflowGraphDiagnosticDetails struct { + Exception *DataflowGraphDiagnosticException `json:"exception,omitempty"` +} + +type DataflowGraphDiagnosticException struct { + ErrorClass string `json:"error_class,omitempty"` + SQLState string `json:"sql_state,omitempty"` +} + +type DataflowGraphDiagnosticNode struct { + DatasetName string `json:"dataset_name,omitempty"` + SinkName string `json:"sink_name,omitempty"` + FlowName string `json:"flow_name,omitempty"` +} + func (s *FakeWorkspace) PipelineGet(pipelineId string) Response { defer s.LockUnlock()() @@ -130,7 +201,7 @@ func (s *FakeWorkspace) PipelineUpdate(req Request, pipelineId string) Response return Response{} } -func (s *FakeWorkspace) PipelineStartUpdate(pipelineId string) Response { +func (s *FakeWorkspace) PipelineStartUpdate(req Request, pipelineId string) Response { defer s.LockUnlock()() _, exists := s.Pipelines[pipelineId] @@ -141,8 +212,26 @@ func (s *FakeWorkspace) PipelineStartUpdate(pipelineId string) Response { } } + // The body is optional (a plain run omits validate_only), so ignore unmarshal errors. + var body pipelines.StartUpdate + _ = json.Unmarshal(req.Body, &body) + updateId := nextUUID() - s.PipelineUpdates[updateId] = true + // Default to a terminal COMPLETED state; tests seed other states via SeedPipelineUpdate. Stamp + // creation_time so newestUpdate has a real ordering key rather than tying at zero. + s.PipelineUpdates[updateId] = &PipelineUpdate{ + PipelineId: pipelineId, + UpdateId: updateId, + State: pipelines.UpdateInfoStateCompleted, + ValidateOnly: body.ValidateOnly, + CreationTime: nowMilli(), + } + + // Seed a deterministic graph so acceptance tests, which drive the API but can't call + // SetPipelineGraph, read a non-empty one. + if s.PipelineGraphs[pipelineId] == nil { + s.PipelineGraphs[pipelineId] = defaultDataflowGraph() + } return Response{ Body: pipelines.StartUpdateResponse{ @@ -181,7 +270,7 @@ func (s *FakeWorkspace) PipelineGetUpdate(pipelineId, updateId string) Response } // Check if the update exists - _, updateExists := s.PipelineUpdates[updateId] + update, updateExists := s.PipelineUpdates[updateId] if !updateExists { return Response{ StatusCode: 404, @@ -192,8 +281,10 @@ func (s *FakeWorkspace) PipelineGetUpdate(pipelineId, updateId string) Response return Response{ Body: pipelines.GetUpdateResponse{ Update: &pipelines.UpdateInfo{ - UpdateId: updateId, - State: pipelines.UpdateInfoStateCompleted, + UpdateId: update.UpdateId, + State: update.State, + ValidateOnly: update.ValidateOnly, + CreationTime: update.CreationTime, }, }, } @@ -217,3 +308,149 @@ func (s *FakeWorkspace) PipelineStop(pipelineId string) Response { }, } } + +// PipelineListUpdates lists a pipeline's updates, newest first, honoring max_results. +func (s *FakeWorkspace) PipelineListUpdates(req Request, pipelineId string) Response { + defer s.LockUnlock()() + + _, exists := s.Pipelines[pipelineId] + if !exists { + return Response{ + StatusCode: 404, + Body: map[string]string{"message": fmt.Sprintf("The specified pipeline %s was not found.", pipelineId)}, + } + } + + var updates []pipelines.UpdateInfo + for _, u := range s.PipelineUpdates { + if u.PipelineId != pipelineId { + continue + } + updates = append(updates, pipelines.UpdateInfo{ + UpdateId: u.UpdateId, + State: u.State, + ValidateOnly: u.ValidateOnly, + CreationTime: u.CreationTime, + }) + } + // Newest first: descending CreationTime, so flip the comparator operands. + slices.SortFunc(updates, func(a, b pipelines.UpdateInfo) int { return cmp.Compare(b.CreationTime, a.CreationTime) }) + + if raw := req.URL.Query().Get("max_results"); raw != "" { + if maxResults, err := strconv.Atoi(raw); err == nil && maxResults > 0 && maxResults < len(updates) { + updates = updates[:maxResults] + } + } + + return Response{ + Body: pipelines.ListUpdatesResponse{ + Updates: updates, + }, + } +} + +// PipelineDataflowGraph returns one page of the pipeline's seeded dataflow graph for leaf (nodes, +// flows, or diagnostics), driving the CLI's pagination loop. +func (s *FakeWorkspace) PipelineDataflowGraph(req Request, pipelineId, leaf string) Response { + defer s.LockUnlock()() + + _, exists := s.Pipelines[pipelineId] + if !exists { + return Response{ + StatusCode: 404, + Body: map[string]string{"message": fmt.Sprintf("The specified pipeline %s was not found.", pipelineId)}, + } + } + + if maxPageSize, ok := dataflowGraphMaxPageSize[leaf]; ok { + if size, _ := strconv.Atoi(req.URL.Query().Get("page_size")); size > maxPageSize { + return Response{StatusCode: 400, Body: map[string]string{ + "error_code": "INVALID_PARAMETER_VALUE", + "message": fmt.Sprintf("Invalid page size. The page size must be between 1 and %d", maxPageSize), + }} + } + } + + graph := s.PipelineGraphs[pipelineId] + if graph == nil { + graph = &DataflowGraph{} + } + + switch leaf { + case "nodes": + page, next := paginate(graph.Nodes, req) + return Response{Body: map[string]any{"nodes": page, "next_page_token": next}} + case "flows": + page, next := paginate(graph.Flows, req) + return Response{Body: map[string]any{"flows": page, "next_page_token": next}} + case "diagnostics": + page, next := paginate(graph.Diagnostics, req) + return Response{Body: map[string]any{"diagnostics": page, "next_page_token": next}} + default: + return Response{StatusCode: 404} + } +} + +// paginate returns the page of items selected by the page_size and page_token query params, plus the +// next-page token (empty once the slice is exhausted). +func paginate[T any](items []T, req Request) ([]T, string) { + start := 0 + if tok := req.URL.Query().Get("page_token"); tok != "" { + start, _ = strconv.Atoi(tok) + } + if start > len(items) { + start = len(items) + } + end := len(items) + if size, err := strconv.Atoi(req.URL.Query().Get("page_size")); err == nil && size > 0 && start+size < end { + end = start + size + } + next := "" + if end < len(items) { + next = strconv.Itoa(end) + } + return items[start:end], next +} + +// defaultDataflowGraph is the deterministic canned graph a dry-run produces for acceptance tests, +// which drive the API but can't call SetPipelineGraph. It mirrors the real backend's identifier +// shapes: a dataset's name is unquoted-dotted while its full_name is per-segment backtick-quoted; a +// view has a populated name but empty full_name; a sink carries its identifier in table_name. +func defaultDataflowGraph() *DataflowGraph { + ds := func(ref, name, fullName, datasetType string) DataflowGraphNode { + return DataflowGraphNode{Dataset: &DataflowGraphDataset{ + DatasetRef: ref, + Name: name, + FullName: fullName, + DatasetType: datasetType, + }} + } + return &DataflowGraph{ + Nodes: []DataflowGraphNode{ + ds("n1", "main.demo.source", "`main`.`demo`.`source`", "MATERIALIZED_VIEW"), + ds("n2", "main.demo.filtered", "`main`.`demo`.`filtered`", "MATERIALIZED_VIEW"), + ds("n3", "main.demo.aggregated", "`main`.`demo`.`aggregated`", "MATERIALIZED_VIEW"), + ds("n4", "recent", "", "VIEW"), + {Sink: &DataflowGraphSink{SinkRef: "s1", TableName: "main.demo.archive"}}, + }, + Flows: []DataflowGraphFlow{ + {InputNodeRefs: []string{"n1"}, OutputNodeRef: "n2"}, + {InputNodeRefs: []string{"n1"}, OutputNodeRef: "n3"}, + {InputNodeRefs: []string{"n2"}, OutputNodeRef: "n4"}, + {InputNodeRefs: []string{"n3"}, OutputNodeRef: "s1"}, + }, + } +} + +// SetPipelineGraph seeds the dataflow graph returned for a pipeline's dry-run. +func (s *FakeWorkspace) SetPipelineGraph(pipelineId string, graph *DataflowGraph) { + defer s.LockUnlock()() + s.PipelineGraphs[pipelineId] = graph +} + +// SeedPipelineUpdate stores a pre-existing update that StartUpdate would not otherwise produce (a +// failed or in-flight one). +func (s *FakeWorkspace) SeedPipelineUpdate(update *PipelineUpdate) { + defer s.LockUnlock()() + s.PipelineUpdates[update.UpdateId] = update +} diff --git a/libs/testserver/pipelines_test.go b/libs/testserver/pipelines_test.go index 39114714082..3b96c722766 100644 --- a/libs/testserver/pipelines_test.go +++ b/libs/testserver/pipelines_test.go @@ -1,6 +1,7 @@ package testserver import ( + "net/url" "testing" "github.com/databricks/databricks-sdk-go/service/pipelines" @@ -51,7 +52,7 @@ func TestPipelineCreate_AllowsSingleSegmentTargetSchemaName(t *testing.T) { func TestPipelineStartUpdate_HandlesNonExistentPipeline(t *testing.T) { workspace := NewFakeWorkspace("http://test", "dbapi123") - response := workspace.PipelineStartUpdate("non-existent-pipeline") + response := workspace.PipelineStartUpdate(Request{}, "non-existent-pipeline") assert.Equal(t, 404, response.StatusCode) assert.Contains(t, response.Body.(map[string]string)["message"], "The specified pipeline non-existent-pipeline was not found") } @@ -69,12 +70,31 @@ func TestPipelineGetUpdate_HandlesNonExistent(t *testing.T) { assert.Contains(t, response.Body.(map[string]string)["message"], "The specified update non-existent-update was not found") } +func TestPipelineDataflowGraph_RejectsPageSizeAboveLeafMax(t *testing.T) { + workspace := NewFakeWorkspace("http://test", "dbapi123") + pipelineId := createTestPipeline(t, workspace) + + graphReq := func(pageSize string) Request { + return Request{URL: &url.URL{RawQuery: url.Values{"page_size": {pageSize}}.Encode()}} + } + + // Diagnostics max is 20; nodes/flows max is 50. The backend rejects (does not clamp) an + // oversized page_size with a 400, so the fake must too. + resp := workspace.PipelineDataflowGraph(graphReq("50"), pipelineId, "diagnostics") + assert.Equal(t, 400, resp.StatusCode) + assert.Equal(t, "INVALID_PARAMETER_VALUE", resp.Body.(map[string]string)["error_code"]) + + assert.Equal(t, 400, workspace.PipelineDataflowGraph(graphReq("51"), pipelineId, "nodes").StatusCode) + assert.Equal(t, 0, workspace.PipelineDataflowGraph(graphReq("20"), pipelineId, "diagnostics").StatusCode) + assert.Equal(t, 0, workspace.PipelineDataflowGraph(graphReq("50"), pipelineId, "nodes").StatusCode) +} + func TestPipelineStop_AfterUpdate(t *testing.T) { workspace := NewFakeWorkspace("http://test", "dbapi123") pipelineId := createTestPipeline(t, workspace) - startResponse := workspace.PipelineStartUpdate(pipelineId) + startResponse := workspace.PipelineStartUpdate(Request{}, pipelineId) assert.Equal(t, 0, startResponse.StatusCode) stopResponse := workspace.PipelineStop(pipelineId) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 5ae3141bfd2..74fe05eadd6 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -355,6 +355,12 @@ func (s *Server) getWorkspaceForToken(token string) *FakeWorkspace { return s.fakeWorkspaces[key] } +// Workspace returns the fake workspace backing token, creating it if needed. Tests use it to seed +// read-only state (e.g. a pipeline's dataflow graph) that has no create API. +func (s *Server) Workspace(token string) *FakeWorkspace { + return s.getWorkspaceForToken(token) +} + func (s *Server) serve(w http.ResponseWriter, r *http.Request, handler HandlerFunc, vars map[string]string) { token := getToken(r)