diff --git a/.nextchanges/bundles/bundle-generate-job-download-script.md b/.nextchanges/bundles/bundle-generate-job-download-script.md new file mode 100644 index 00000000000..14c29808030 --- /dev/null +++ b/.nextchanges/bundles/bundle-generate-job-download-script.md @@ -0,0 +1 @@ +* `bundle generate job` can now download workspace files referenced by `spark_python_task`, rewriting them to a relative path like it already does for notebooks. This is opt-in via the `--download-spark-python-files` flag ([#5799](https://github.com/databricks/cli/pull/5799)). diff --git a/acceptance/bundle/generate/spark_python_task_job/databricks.yml b/acceptance/bundle/generate/spark_python_task_job/databricks.yml new file mode 100644 index 00000000000..7419eb67c8a --- /dev/null +++ b/acceptance/bundle/generate/spark_python_task_job/databricks.yml @@ -0,0 +1,2 @@ +bundle: + name: spark_python_task_job diff --git a/acceptance/bundle/generate/spark_python_task_job/out.job.yml b/acceptance/bundle/generate/spark_python_task_job/out.job.yml new file mode 100644 index 00000000000..091ed6b6726 --- /dev/null +++ b/acceptance/bundle/generate/spark_python_task_job/out.job.yml @@ -0,0 +1,11 @@ +resources: + jobs: + out: + name: sparkpythonjob + tasks: + - task_key: workspace_file_task + spark_python_task: + python_file: outjob.py + - task_key: cloud_uri_task + spark_python_task: + python_file: dbfs:/FileStore/job.py diff --git a/acceptance/bundle/generate/spark_python_task_job/out.test.toml b/acceptance/bundle/generate/spark_python_task_job/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/bundle/generate/spark_python_task_job/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/generate/spark_python_task_job/outjob.py b/acceptance/bundle/generate/spark_python_task_job/outjob.py new file mode 100644 index 00000000000..7df869a15e7 --- /dev/null +++ b/acceptance/bundle/generate/spark_python_task_job/outjob.py @@ -0,0 +1 @@ +print("Hello, World!") diff --git a/acceptance/bundle/generate/spark_python_task_job/output.txt b/acceptance/bundle/generate/spark_python_task_job/output.txt new file mode 100644 index 00000000000..e160a45625f --- /dev/null +++ b/acceptance/bundle/generate/spark_python_task_job/output.txt @@ -0,0 +1,2 @@ +File successfully saved to outjob.py +Job configuration successfully saved to out.job.yml diff --git a/acceptance/bundle/generate/spark_python_task_job/script b/acceptance/bundle/generate/spark_python_task_job/script new file mode 100644 index 00000000000..a1131f4c5dd --- /dev/null +++ b/acceptance/bundle/generate/spark_python_task_job/script @@ -0,0 +1 @@ +$CLI bundle generate job --existing-job-id 1234 --config-dir . --key out --force --source-dir . --download-spark-python-files diff --git a/acceptance/bundle/generate/spark_python_task_job/test.toml b/acceptance/bundle/generate/spark_python_task_job/test.toml new file mode 100644 index 00000000000..6ee12cc669b --- /dev/null +++ b/acceptance/bundle/generate/spark_python_task_job/test.toml @@ -0,0 +1,41 @@ +[[Server]] +Pattern = "GET /api/2.2/jobs/get" +Response.Body = ''' +{ + "job_id": 11223344, + "settings": { + "name": "sparkpythonjob", + "tasks": [ + { + "task_key": "workspace_file_task", + "spark_python_task": { + "python_file": "/Workspace/Users/tester@databricks.com/outjob.py" + } + }, + { + "task_key": "cloud_uri_task", + "spark_python_task": { + "python_file": "dbfs:/FileStore/job.py" + } + } + ] + } +} +''' + +[[Server]] +Pattern = "GET /api/2.0/workspace/get-status" +Response.Body = ''' +{ + "path": "/Workspace/Users/tester@databricks.com/outjob.py", + "object_type": "FILE", + "language": "PYTHON", + "repos_export_format": "SOURCE" +} +''' + +[[Server]] +Pattern = "GET /api/2.0/workspace/export" +Response.Body = ''' +print("Hello, World!") +''' diff --git a/acceptance/bundle/help/bundle-generate-job/output.txt b/acceptance/bundle/help/bundle-generate-job/output.txt index f239691ab6a..092babc9af9 100644 --- a/acceptance/bundle/help/bundle-generate-job/output.txt +++ b/acceptance/bundle/help/bundle-generate-job/output.txt @@ -18,7 +18,8 @@ Examples: What gets generated: - Job configuration YAML file in the resources directory -- Any associated notebook or Python files in the source directory +- Any associated notebook files in the source directory +- Any associated python spark files in the source directory, if --download-spark-python-files is provided After generation, you can deploy this job to other targets using: databricks bundle deploy --target staging @@ -28,11 +29,12 @@ Usage: databricks bundle generate job [flags] Flags: - -d, --config-dir string Dir path where the output config will be stored (default "resources") - --existing-job-id int Job ID of the job to generate config for - -f, --force Force overwrite existing files in the output directory - -h, --help help for job - -s, --source-dir string Dir path where the downloaded files will be stored (default "src") + -d, --config-dir string Dir path where the output config will be stored (default "resources") + --download-spark-python-files download workspace files referenced by spark_python_task and rewrite them to a relative path + --existing-job-id int Job ID of the job to generate config for + -f, --force Force overwrite existing files in the output directory + -h, --help help for job + -s, --source-dir string Dir path where the downloaded files will be stored (default "src") Global Flags: --debug enable debug logging diff --git a/bundle/generate/downloader.go b/bundle/generate/downloader.go index 87aa4ef5751..334175496aa 100644 --- a/bundle/generate/downloader.go +++ b/bundle/generate/downloader.go @@ -34,14 +34,44 @@ type Downloader struct { sourceDir string configDir string basePath string + + // downloadSparkPythonFiles controls whether workspace files referenced by a + // spark_python_task are downloaded. It is opt-in because a Python file often + // imports sibling files that are not captured, so downloading only the entry + // point can produce a job that fails at runtime with missing imports. + downloadSparkPythonFiles bool +} + +// DownloaderOption configures a Downloader. +type DownloaderOption func(*Downloader) + +// WithSparkPythonFiles enables downloading workspace files referenced by a +// spark_python_task. +func WithSparkPythonFiles() DownloaderOption { + return func(n *Downloader) { + n.downloadSparkPythonFiles = true + } } func (n *Downloader) MarkTaskForDownload(ctx context.Context, task *jobs.Task) error { - if task.NotebookTask == nil { - return nil + if task.NotebookTask != nil { + return n.markNotebookForDownload(ctx, &task.NotebookTask.NotebookPath) + } + + if n.downloadSparkPythonFiles && task.SparkPythonTask != nil && isWorkspaceFileTask(task.SparkPythonTask.Source, task.SparkPythonTask.PythonFile) { + return n.markFileForDownload(ctx, &task.SparkPythonTask.PythonFile) } - return n.markNotebookForDownload(ctx, &task.NotebookTask.NotebookPath) + return nil +} + +// isWorkspaceFileTask reports whether a task file lives in the Databricks +// workspace and should be downloaded. Files sourced from Git (source: GIT) are +// left to the Git repository, and python_file also accepts cloud URIs +// (dbfs:/, s3:/, adls:/, gcs:/) which are not workspace paths; per the Jobs API, +// workspace files are absolute and begin with "/". +func isWorkspaceFileTask(source jobs.Source, filePath string) bool { + return source != jobs.SourceGit && strings.HasPrefix(filePath, "/") } func (n *Downloader) MarkPipelineLibraryForDownload(ctx context.Context, lib *pipelines.PipelineLibrary) error { @@ -218,6 +248,9 @@ func (n *Downloader) MarkTasksForDownload(ctx context.Context, tasks []jobs.Task if task.NotebookTask != nil { paths = append(paths, task.NotebookTask.NotebookPath) } + if n.downloadSparkPythonFiles && task.SparkPythonTask != nil && isWorkspaceFileTask(task.SparkPythonTask.Source, task.SparkPythonTask.PythonFile) { + paths = append(paths, task.SparkPythonTask.PythonFile) + } } if len(paths) > 0 { n.basePath = commonDirPrefix(paths) @@ -350,11 +383,15 @@ func writeAndClose(dst io.WriteCloser, src io.Reader) error { return err } -func NewDownloader(w *databricks.WorkspaceClient, sourceDir, configDir string) *Downloader { - return &Downloader{ +func NewDownloader(w *databricks.WorkspaceClient, sourceDir, configDir string, opts ...DownloaderOption) *Downloader { + n := &Downloader{ files: make(map[string]exportFile), w: w, sourceDir: sourceDir, configDir: configDir, } + for _, opt := range opts { + opt(n) + } + return n } diff --git a/bundle/generate/downloader_test.go b/bundle/generate/downloader_test.go index ed8c9cbf045..d3d6bca8aca 100644 --- a/bundle/generate/downloader_test.go +++ b/bundle/generate/downloader_test.go @@ -295,6 +295,98 @@ func TestDownloader_MarkTasksForDownload_SingleNotebook(t *testing.T) { assert.Len(t, downloader.files, 1) } +func TestDownloader_MarkTasksForDownload_SparkPythonFile(t *testing.T) { + ctx := t.Context() + m := mocks.NewMockWorkspaceClient(t) + + dir := "base/dir" + sourceDir := filepath.Join(dir, "source") + configDir := filepath.Join(dir, "config") + downloader := NewDownloader(m.WorkspaceClient, sourceDir, configDir, WithSparkPythonFiles()) + + pythonFile := "/Users/user/project/etl/job.py" + m.GetMockWorkspaceAPI().EXPECT().GetStatusByPath(ctx, pythonFile).Return(&workspace.ObjectInfo{ + Path: pythonFile, + }, nil) + + tasks := []jobs.Task{ + { + TaskKey: "spark_python_task", + SparkPythonTask: &jobs.SparkPythonTask{ + PythonFile: pythonFile, + }, + }, + } + + err := downloader.MarkTasksForDownload(ctx, tasks) + require.NoError(t, err) + + assert.Equal(t, "../source/job.py", tasks[0].SparkPythonTask.PythonFile) + require.Len(t, downloader.files, 1) + f := downloader.files[filepath.Join(sourceDir, "job.py")] + assert.Equal(t, pythonFile, f.path) + assert.Equal(t, workspace.ExportFormatSource, f.format) +} + +func TestDownloader_MarkTasksForDownload_SparkPythonFileSkipped(t *testing.T) { + ctx := t.Context() + // Cloud URIs and Git-sourced files are not workspace paths, so no + // get-status/download requests should be made for them. + w := newTestWorkspaceClient(t, func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path) + }) + + downloader := NewDownloader(w, "source", "config", WithSparkPythonFiles()) + + tasks := []jobs.Task{ + { + TaskKey: "cloud_uri_task", + SparkPythonTask: &jobs.SparkPythonTask{ + PythonFile: "dbfs:/FileStore/job.py", + }, + }, + { + TaskKey: "git_task", + SparkPythonTask: &jobs.SparkPythonTask{ + PythonFile: "etl/job.py", + Source: jobs.SourceGit, + }, + }, + } + + err := downloader.MarkTasksForDownload(ctx, tasks) + require.NoError(t, err) + assert.Empty(t, downloader.files) + assert.Equal(t, "dbfs:/FileStore/job.py", tasks[0].SparkPythonTask.PythonFile) + assert.Equal(t, "etl/job.py", tasks[1].SparkPythonTask.PythonFile) +} + +func TestDownloader_MarkTasksForDownload_SparkPythonFileDisabledByDefault(t *testing.T) { + ctx := t.Context() + // Without WithSparkPythonFiles, workspace files referenced by a + // spark_python_task are left untouched, so no requests are made. + w := newTestWorkspaceClient(t, func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path) + }) + + downloader := NewDownloader(w, "source", "config") + + pythonFile := "/Users/user/project/etl/job.py" + tasks := []jobs.Task{ + { + TaskKey: "spark_python_task", + SparkPythonTask: &jobs.SparkPythonTask{ + PythonFile: pythonFile, + }, + }, + } + + err := downloader.MarkTasksForDownload(ctx, tasks) + require.NoError(t, err) + assert.Empty(t, downloader.files) + assert.Equal(t, pythonFile, tasks[0].SparkPythonTask.PythonFile) +} + func TestDownloader_MarkTasksForDownload_NoNotebooks(t *testing.T) { ctx := t.Context() w := newTestWorkspaceClient(t, func(w http.ResponseWriter, r *http.Request) { diff --git a/cmd/bundle/generate/job.go b/cmd/bundle/generate/job.go index c3aba49c5f2..705d35bcff7 100644 --- a/cmd/bundle/generate/job.go +++ b/cmd/bundle/generate/job.go @@ -27,6 +27,7 @@ func NewGenerateJobCommand() *cobra.Command { var jobId int64 var force bool var bind bool + var downloadSparkPythonFiles bool cmd := &cobra.Command{ Use: "job", @@ -49,7 +50,8 @@ Examples: What gets generated: - Job configuration YAML file in the resources directory -- Any associated notebook or Python files in the source directory +- Any associated notebook files in the source directory +- Any associated python spark files in the source directory, if --download-spark-python-files is provided After generation, you can deploy this job to other targets using: databricks bundle deploy --target staging @@ -64,6 +66,7 @@ After generation, you can deploy this job to other targets using: cmd.Flags().BoolVarP(&force, "force", "f", false, `Force overwrite existing files in the output directory`) cmd.Flags().BoolVarP(&bind, "bind", "b", false, `automatically bind the generated resource to the existing resource`) cmd.Flags().MarkHidden("bind") + cmd.Flags().BoolVar(&downloadSparkPythonFiles, "download-spark-python-files", false, `download workspace files referenced by spark_python_task and rewrite them to a relative path`) cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := logdiag.InitContext(cmd.Context()) @@ -80,7 +83,11 @@ After generation, you can deploy this job to other targets using: return err } - downloader := generate.NewDownloader(w, sourceDir, configDir) + var opts []generate.DownloaderOption + if downloadSparkPythonFiles { + opts = append(opts, generate.WithSparkPythonFiles()) + } + downloader := generate.NewDownloader(w, sourceDir, configDir, opts...) // Don't download files if the job is using Git source // When Git source is used, the job will be using the files from the Git repository