Normalize notebook cell source before wrapping jobs

packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.test.ts

@@ -432,6 +432,43 @@ describe(__filename, async () => {
             });
         });
 
+        it("should create wrapper for jupyter notebook with string cell source", async () => {
+            await withFile(async (localFilePath) => {
+                const originalData = {
+                    ...notebookMetadata,
+                    cells: [{...cellMetadata, source: "b = 2"}],
+                };
+                await writeFile(
+                    localFilePath.path,
+                    JSON.stringify(originalData),
+                    "utf-8"
+                );
+
+                await new WorkspaceFsWorkflowWrapper(
+                    instance(mockConnectionManager),
+                    instance(mockExtensionContext)
+                ).createNotebookWrapper(
+                    new LocalUri(localFilePath.path),
+                    new RemoteUri(originalFilePath),
+                    new RemoteUri(testDirPath),
+                    "IPYNB"
+                );
+
+                const wrapperData = await getWrapperData();
+                const expected = {
+                    ...originalData,
+                    cells: [wrapperData].concat([
+                        {...cellMetadata, source: ["b = 2"]},
+                    ]),
+                };
+
+                verifyMockServiceCalledWithExpectedData(
+                    JSON.stringify(expected),
+                    wrappedFilePath
+                );
+            });
+        });
+
         it("should rearrange kernel restart commands to the beginning", async () => {
             await withFile(async (localFilePath) => {
                 const originalData = {

packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.ts

@@ -144,7 +144,7 @@ export class WorkspaceFsWorkflowWrapper {
         @context ctx?: Context
     ) {
         // eslint-disable-next-line @typescript-eslint/naming-convention
-        type JupyterCell = {source: string[]; cell_type: string};
+        type JupyterCell = {source: string[] | string; cell_type: string};
         const data = await readFile(localFilePath.path, "utf-8");
         const originalJson: {cells: JupyterCell[] | undefined} =
             JSON.parse(data);
@@ -163,14 +163,20 @@ export class WorkspaceFsWorkflowWrapper {
         const cells = [bootstrapJson].concat(originalJson["cells"] ?? []).map(
             // Since each cell.source is a string array where each string can be
             // multiple lines, we need to split each string by \n and then flatten
-            (cell) =>
-                ({
-                    source: cell.source?.flatMap((line) =>
+            (cell) => {
+                const sourceLines = Array.isArray(cell.source)
+                    ? cell.source
+                    : typeof cell.source === "string"
+                      ? [cell.source]
+                      : undefined;
+                return {
+                    source: sourceLines?.flatMap((line) =>
                         line.trimEnd().split(/\r?\n/)
                     ),
                     type: cell.cell_type === "code" ? "code" : "not_code",
                     originalCell: cell,
-                }) as Cell
+                } as Cell;
+            }
         );
         originalJson["cells"] = rearrangeCells(cells).map((cell) => {
             if (cell.type === "not_code") {