diff --git a/packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.test.ts b/packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.test.ts index 1d2409cba..df06aab8f 100644 --- a/packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.test.ts +++ b/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 = { diff --git a/packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.ts b/packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.ts index 6c70448b6..31d0c1e2b 100644 --- a/packages/databricks-vscode/src/workspace-fs/WorkspaceFsWorkflowWrapper.ts +++ b/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") {