Fix bundle generate job to preserve nested notebook directory structure

Signed-off-by: Varun Deep Saini <varun.23bcs10048@ms.sst.scaler.com>

Author: Varun Deep Saini

NEXT_CHANGELOG.md

@@ -6,6 +6,7 @@
 
 ### Bundles
 * Modify grants to use SDK types ([#4666](https://github.com/databricks/cli/pull/4666))
+* Fix `bundle generate` job to preserve nested notebook directory structure ([#4596](https://github.com/databricks/cli/pull/4596))
 
 ### Dependency updates
 * Bump databricks-sdk-go from v0.112.0 to v0.119.0 ([#4631](https://github.com/databricks/cli/pull/4631), [#4695](https://github.com/databricks/cli/pull/4695))

acceptance/bundle/generate/job_nested_notebooks/databricks.yml

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

acceptance/bundle/generate/job_nested_notebooks/out.job.yml

@@ -0,0 +1,11 @@
+resources:
+  jobs:
+    out:
+      name: dev.my_repo.my_job
+      tasks:
+        - task_key: my_notebook_task
+          notebook_task:
+            notebook_path: src/my_folder/my_notebook.py
+        - task_key: other_notebook_task
+          notebook_task:
+            notebook_path: src/other_folder/other_notebook.py

acceptance/bundle/generate/job_nested_notebooks/out.test.toml

@@ -0,0 +1,9 @@
+File successfully saved to src/my_folder/my_notebook.py
+File successfully saved to src/other_folder/other_notebook.py
+Job configuration successfully saved to out.job.yml
+=== old flattened files should be gone ===
+src/my_notebook.py removed
+src/other_notebook.py removed
+=== new nested files ===
+src/my_folder/my_notebook.py
+src/other_folder/other_notebook.py

acceptance/bundle/generate/job_nested_notebooks/output.txt

@@ -0,0 +1,12 @@
+mkdir -p src
+echo "old" > src/my_notebook.py
+echo "old" > src/other_notebook.py
+
+$CLI bundle generate job --existing-job-id 1234 --config-dir . --key out --force --source-dir src 2>&1 | sort
+
+echo "=== old flattened files should be gone ==="
+test ! -f src/my_notebook.py && echo "src/my_notebook.py removed" || echo "src/my_notebook.py still exists"
+test ! -f src/other_notebook.py && echo "src/other_notebook.py removed" || echo "src/other_notebook.py still exists"
+
+echo "=== new nested files ==="
+find src -type f | sort

acceptance/bundle/generate/job_nested_notebooks/script

@@ -0,0 +1,42 @@
+Ignore = ["src"]
+
+[[Server]]
+Pattern = "GET /api/2.2/jobs/get"
+Response.Body = '''
+{
+    "job_id": 11223344,
+    "settings": {
+        "name": "dev.my_repo.my_job",
+        "tasks": [
+            {
+                "task_key": "my_notebook_task",
+                "notebook_task": {
+                    "notebook_path": "/my_data_product/dev/my_folder/my_notebook"
+                }
+            },
+            {
+                "task_key": "other_notebook_task",
+                "notebook_task": {
+                    "notebook_path": "/my_data_product/dev/other_folder/other_notebook"
+                }
+            }
+        ]
+    }
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/workspace/get-status"
+Response.Body = '''
+{
+    "object_type": "NOTEBOOK",
+    "language": "PYTHON",
+    "repos_export_format": "SOURCE"
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/workspace/export"
+Response.Body = '''
+print("Hello, World!")
+'''

acceptance/bundle/generate/job_nested_notebooks/test.toml

@@ -11,6 +11,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/notebook"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
@@ -73,7 +74,7 @@ func (n *Downloader) markFileForDownload(ctx context.Context, filePath *string)
 		return err
 	}
 
-	*filePath = rel
+	*filePath = filepath.ToSlash(rel)
 	return nil
 }
 
@@ -109,7 +110,7 @@ func (n *Downloader) MarkDirectoryForDownload(ctx context.Context, dirPath *stri
 		return err
 	}
 
-	*dirPath = rel
+	*dirPath = filepath.ToSlash(rel)
 	return nil
 }
 
@@ -203,10 +204,75 @@ func (n *Downloader) markNotebookForDownload(ctx context.Context, notebookPath *
 		return err
 	}
 
-	*notebookPath = rel
+	*notebookPath = filepath.ToSlash(rel)
 	return nil
 }
 
+func (n *Downloader) MarkTasksForDownload(ctx context.Context, tasks []jobs.Task) error {
+	var paths []string
+	for _, task := range tasks {
+		if task.NotebookTask != nil {
+			paths = append(paths, task.NotebookTask.NotebookPath)
+		}
+	}
+	if len(paths) > 0 {
+		n.basePath = commonDirPrefix(paths)
+	}
+	for i := range tasks {
+		if err := n.MarkTaskForDownload(ctx, &tasks[i]); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (n *Downloader) CleanupOldFiles(ctx context.Context) {
+	for targetPath := range n.files {
+		rel, err := filepath.Rel(n.sourceDir, targetPath)
+		if err != nil {
+			continue
+		}
+		if filepath.Base(rel) == rel {
+			continue
+		}
+		oldPath := filepath.Join(n.sourceDir, filepath.Base(rel))
+		if _, isNewFile := n.files[oldPath]; isNewFile {
+			continue
+		}
+		if err := os.Remove(oldPath); err == nil {
+			log.Infof(ctx, "Removed previously generated file %s", filepath.ToSlash(oldPath))
+		}
+	}
+}
+
+// commonDirPrefix returns the longest common directory-aligned prefix of the given paths.
+func commonDirPrefix(paths []string) string {
+	if len(paths) == 0 {
+		return ""
+	}
+	if len(paths) == 1 {
+		return path.Dir(paths[0])
+	}
+
+	prefix := paths[0]
+	for _, p := range paths[1:] {
+		for !strings.HasPrefix(p, prefix) {
+			prefix = prefix[:len(prefix)-1]
+			if prefix == "" {
+				return ""
+			}
+		}
+	}
+
+	// Truncate to last '/' to ensure directory alignment.
+	if i := strings.LastIndex(prefix, "/"); i >= 0 {
+		prefix = prefix[:i]
+	} else {
+		prefix = ""
+	}
+	return prefix
+}
+
 func (n *Downloader) relativePath(fullPath string) string {
 	basePath := path.Dir(fullPath)
 	if n.basePath != "" {