diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 30a64c5410..fa68022925 100644 --- a/NEXT_CHANGELOG.md +++ b/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 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 0000000000..7419eb67c8 --- /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 0000000000..091ed6b672 --- /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 0000000000..f784a18325 --- /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 0000000000..7df869a15e --- /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 0000000000..e160a45625 --- /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 0000000000..16cfc1589c --- /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 . 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 0000000000..6ee12cc669 --- /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/bundle/generate/downloader.go b/bundle/generate/downloader.go index 87aa4ef575..2a03cda042 100644 --- a/bundle/generate/downloader.go +++ b/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) diff --git a/bundle/generate/downloader_test.go b/bundle/generate/downloader_test.go index ed8c9cbf04..2bb887cb8e 100644 --- a/bundle/generate/downloader_test.go +++ b/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) {