Download spark_python_task workspace files behind --download-spark-python-files flag

.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)).

acceptance/bundle/generate/spark_python_task_job/databricks.yml

@@ -0,0 +1,2 @@
+bundle:
+  name: spark_python_task_job

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

acceptance/bundle/generate/spark_python_task_job/outjob.py

@@ -0,0 +1 @@
+print("Hello, World!")

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

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

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!")
+'''

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

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
 }

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) {

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