Download spark_python_task workspace files in bundle generate job

NEXT_CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bundles
 
+ * `bundle generate job` now downloads workspace files referenced by `spark_python_task`, rewriting them to a relative path like it already does for notebooks. Git-sourced files and cloud URIs are left untouched ([#5799](https://github.com/databricks/cli/pull/5799)).
+
 ### Dependency updates
 
 ### API Changes

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 .

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

bundle/generate/downloader.go

@@ -37,11 +37,24 @@ type Downloader struct {
 }
 
 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)
 	}
 
-	return n.markNotebookForDownload(ctx, &task.NotebookTask.NotebookPath)
+	if task.SparkPythonTask != nil && isWorkspaceFileTask(task.SparkPythonTask.Source, task.SparkPythonTask.PythonFile) {
+		return n.markFileForDownload(ctx, &task.SparkPythonTask.PythonFile)
+	}
+
+	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 +231,9 @@ func (n *Downloader) MarkTasksForDownload(ctx context.Context, tasks []jobs.Task
 		if task.NotebookTask != nil {
 			paths = append(paths, task.NotebookTask.NotebookPath)
 		}
+		if task.SparkPythonTask != nil && isWorkspaceFileTask(task.SparkPythonTask.Source, task.SparkPythonTask.PythonFile) {
+			paths = append(paths, task.SparkPythonTask.PythonFile)
+		}
 	}
 	if len(paths) > 0 {
 		n.basePath = commonDirPrefix(paths)

bundle/generate/downloader_test.go

@@ -295,6 +295,72 @@ 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)
+
+	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")
+
+	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_NoNotebooks(t *testing.T) {
 	ctx := t.Context()
 	w := newTestWorkspaceClient(t, func(w http.ResponseWriter, r *http.Request) {