diff --git a/CHANGELOG.md b/CHANGELOG.md index 7441692..b4accce 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Release Notes +**v2.1.0:** +- introduced a new notebook type for `Scala`, `SQL` and `R` +- removed interactive notebook kernel as it never really worked + - replaced it with new Databricks kernel to work with the new notebook type +- minor fix to handle `pip install` more generic + **v2.0.3:** - fixed bug with `Buffer` introduced in `v2.0.2` diff --git a/package.json b/package.json index 05a2dbf..a165162 100644 --- a/package.json +++ b/package.json @@ -2,7 +2,7 @@ "name": "databricks-vscode", "displayName": "Databricks Power Tools", "description": "Run notebooks cell-by-cell, browse and edit your Databricks Workspace, DBFS, Clusters, Jobs, Secrets, Repos and SQL. Supports Azure Databricks, Databricks on AWS and Databricks on GCP.", - "version": "2.0.3", + "version": "2.1.0", "publisher": "paiqo", "icon": "resources/databricks_extension.png", "author": { @@ -1158,7 +1158,25 @@ "when": "notebookKernel =~ /^paiqo.databricks-vscode\\// && isWorkspaceTrusted || isWorkspaceTrusted && !notebookKernel" } ] - } + }, + "notebooks": [ + { + "type": "databricks-notebook", + "displayName": "Databricks Notebook", + "selector": [ + { + "filenamePattern": "*.sql" + }, + { + "filenamePattern": "*.scala" + }, + { + "filenamePattern": "*.r" + } + ], + "priority": "default" + } + ] }, "scripts": { "vscode:prepublish": "npm run package", @@ -1189,4 +1207,4 @@ "webpack-cli": "^4.10.0", "buffer": "^6.0.3" } -} +} \ No newline at end of file diff --git a/src/databricksApi/_types.ts b/src/databricksApi/_types.ts index 51314be..11c2f31 100644 --- a/src/databricksApi/_types.ts +++ b/src/databricksApi/_types.ts @@ -41,6 +41,7 @@ export interface iDatabricksApiJobsRunsListResponse { export interface iDatabricksApiContextsCreateResponse { id: string; + error?: string; } export interface iDatabricksApiContextsStatusResponse { @@ -49,34 +50,40 @@ export interface iDatabricksApiContextsStatusResponse { } export interface iDatabricksApiContextsDestroyResponse { - id: string; + id?: string; + error?: string; } export interface iDatabricksApiCommandsExecuteResponse { - id: string; + id?: string; + error?: string; } export interface iDatabricksApiCommandsCancelResponse { - id: string; + id?: string; + error?: string; } export interface iDatabricksApiCommandsStatusResponse { - id: string; - status: string; - results: any; + id?: string; + status?: string; + results?: any; + error?: string; } export class ExecutionContext { - context_id: string; - cluster_id: string; - language: string; + context_id?: string; + cluster_id?: string; + language?: string; + error?: string; } export class ExecutionCommand { - command_id: string; - context_id: string; - cluster_id: string; - language: string; + command_id?: string; + context_id?: string; + cluster_id?: string; + language?: string; + error?: string; } export interface iDatabricksApiClustersSparkVersionsResponse { diff --git a/src/databricksApi/databricksApiService.ts b/src/databricksApi/databricksApiService.ts index fdd1d1b..e0c32ef 100644 --- a/src/databricksApi/databricksApiService.ts +++ b/src/databricksApi/databricksApiService.ts @@ -412,6 +412,10 @@ export abstract class DatabricksApiService { let response: iDatabricksApiContextsCreateResponse = await this.post(endpoint, body); + if (!response || response.error) { + return response; + } + let ret = { "cluster_id": cluster_id, "language": language, @@ -427,6 +431,10 @@ export abstract class DatabricksApiService { let response: iDatabricksApiContextsDestroyResponse = await this.post(endpoint, body); + if (!response || response.error) { + return response; + } + return response; } @@ -441,6 +449,10 @@ export abstract class DatabricksApiService { let response: iDatabricksApiCommandsExecuteResponse = await this.post(endpoint, body); + if (!response || response.error) { + return response; + } + let ret = { "command_id": response.id, "cluster_id": context.cluster_id, @@ -461,6 +473,10 @@ export abstract class DatabricksApiService { let response: iDatabricksApiCommandsStatusResponse = await this.get(endpoint, body); + if (!response || response.error) { + return response; + } + return response; } @@ -486,6 +502,10 @@ export abstract class DatabricksApiService { let response: iDatabricksApiCommandsCancelResponse = await this.post(endpoint, body); + if (!response || response.error) { + return response; + } + let ret = { "command_id": response.id, "cluster_id": command.cluster_id, diff --git a/src/extension.ts b/src/extension.ts index 8dadcef..80e8861 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -34,6 +34,7 @@ import { FSHelper } from './helpers/FSHelper'; import { DatabricksKernelManager } from './vscode/notebook/DatabricksKernelManager'; import { DatabricksSQLTreeItem } from './vscode/treeviews/sql/DatabricksSQLTreeItem'; import { DatabricksSecretTreeItem } from './vscode/treeviews/secrets/DatabricksSecretTreeItem'; +import { DatabricksNotebookSerializer } from './vscode/notebook/DatabricksNotebookSerializer'; export async function activate(context: vscode.ExtensionContext) { @@ -64,10 +65,12 @@ export async function activate(context: vscode.ExtensionContext) { ThisExtension.setStatusBar("Initialized!"); + let notebookSerializer = new DatabricksNotebookSerializer(context); + ThisExtension.setStatusBar("Initializing Kernels ...", true); DatabricksKernelManager.initialize(); vscode.commands.registerCommand('databricksKernel.restart', - (notebook: { notebookEditor: { notebookUri: vscode.Uri } } | undefined | vscode.Uri) => DatabricksKernelManager.restartNotebookKernel(notebook) + (notebook: { notebookEditor: { notebookUri: vscode.Uri } } | undefined | vscode.Uri) => DatabricksKernelManager.restartJupyterKernel(notebook) ); vscode.commands.registerCommand('databricksKernel.updateWidgets', (notebook: { notebookEditor: { notebookUri: vscode.Uri } } | undefined | vscode.Uri) => DatabricksKernelManager.updateWidgets(notebook) diff --git a/src/vscode/notebook/DatabricksKernel.ts b/src/vscode/notebook/DatabricksKernel.ts index 6cb4efc..145bf29 100644 --- a/src/vscode/notebook/DatabricksKernel.ts +++ b/src/vscode/notebook/DatabricksKernel.ts @@ -25,7 +25,7 @@ export type NotebookMagic = export type KernelType = "jupyter-notebook" - | "interactive" + | "databricks-notebook" // https://code.visualstudio.com/blogs/2021/11/08/custom-notebooks export class DatabricksKernel implements vscode.NotebookController { @@ -96,7 +96,7 @@ export class DatabricksKernel implements vscode.NotebookController { let execContext: ExecutionContext = this.getNotebookContext(notebook.uri); if (!execContext) { ThisExtension.setStatusBar("Initializing Kernel ...", true); - execContext = await this.initializeExecutionContext() + execContext = await this.initializeExecutionContext(notebook.metadata?.notebookLanguage?.databricksLanguage); this.setNotebookContext(notebook.uri, execContext); await this.updateRepoContext(notebook.uri); ThisExtension.setStatusBar("Kernel initialized!"); @@ -205,8 +205,8 @@ export class DatabricksKernel implements vscode.NotebookController { } } - private async initializeExecutionContext(): Promise { - return await DatabricksApiService.getExecutionContext(this.ClusterID, this.Language); + private async initializeExecutionContext(language?: ContextLanguage): Promise { + return await DatabricksApiService.getExecutionContext(this.ClusterID, language ?? this.Language); } setNotebookContext(notebookUri: vscode.Uri, context: ExecutionContext): void { @@ -282,13 +282,11 @@ export class DatabricksKernel implements vscode.NotebookController { } } - - private async parseCommand(cell: vscode.NotebookCell, executionContext: ExecutionContext): Promise<[ContextLanguage, string, NotebookMagic]> { let cmd: string = cell.document.getText(); let magicText: string = undefined; let commandText: string = cmd; - let language: ContextLanguage = this.Language; + let language: ContextLanguage = executionContext.language as ContextLanguage; if (cmd[0] == "%") { let lines = cmd.split('\n'); magicText = lines[0].split(" ")[0].slice(1).trim(); @@ -484,6 +482,19 @@ export class DatabricksKernel implements vscode.NotebookController { return; }); + if (command.error) { + execution.appendOutput(new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.text(command.error, 'text/html') + ])); + + execution.appendOutput(new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.error(new Error(command.error)), + ])); + + execution.end(false, Date.now()); + return; + } + let result = await DatabricksApiService.getCommandResult(command, true); if (result.results.resultType == "table") { @@ -570,7 +581,7 @@ export class DatabricksKernel implements vscode.NotebookController { } - if (["pip"].includes(magic)) { + if (["pip"].includes(magic) || commandText.includes("pip install")) { await this.updateRepoContext(cell.notebook.uri); } diff --git a/src/vscode/notebook/DatabricksKernelManager.ts b/src/vscode/notebook/DatabricksKernelManager.ts index 5e45c78..555a24e 100644 --- a/src/vscode/notebook/DatabricksKernelManager.ts +++ b/src/vscode/notebook/DatabricksKernelManager.ts @@ -7,8 +7,8 @@ import { DatabricksKernel } from './DatabricksKernel'; export abstract class DatabricksKernelManager { - private static NotebookKernelSuffix: string = "-jupyter-notebook"; - private static InteractiveKernelSuffix: string = "-interactive"; + private static JupyterKernelSuffix: string = "-jupyter-notebook"; + private static DatabricksKernelSuffix: string = "-databricks-notebook"; private static _kernels: Map = new Map(); @@ -57,40 +57,40 @@ export abstract class DatabricksKernelManager { return this._kernels.get(kernelName); } - static getNotebookKernelName(cluster: iDatabricksCluster): string { - return (cluster.kernel_id ?? cluster.cluster_id) + DatabricksKernelManager.NotebookKernelSuffix; + static getJupyterKernelName(cluster: iDatabricksCluster): string { + return (cluster.kernel_id ?? cluster.cluster_id) + DatabricksKernelManager.JupyterKernelSuffix; } - static getNotebookKernel(cluster: iDatabricksCluster): DatabricksKernel { - return this.getKernel(this.getNotebookKernelName(cluster)); + static getJupyterKernel(cluster: iDatabricksCluster): DatabricksKernel { + return this.getKernel(this.getJupyterKernelName(cluster)); } - static notebookKernelExists(cluster: iDatabricksCluster): boolean { - if (this.getKernel(this.getNotebookKernelName(cluster))) { + static jupyterKernelExists(cluster: iDatabricksCluster): boolean { + if (this.getKernel(this.getJupyterKernelName(cluster))) { return true; } return false; } - static getInteractiveKernelName(cluster: iDatabricksCluster): string { - return cluster.cluster_id + DatabricksKernelManager.InteractiveKernelSuffix + static getDatabricksKernelName(cluster: iDatabricksCluster): string { + return cluster.cluster_id + DatabricksKernelManager.DatabricksKernelSuffix } - static getInteractiveKernel(cluster: iDatabricksCluster): DatabricksKernel { - return this.getKernel(this.getInteractiveKernelName(cluster)); + static getDatabricksKernel(cluster: iDatabricksCluster): DatabricksKernel { + return this.getKernel(this.getDatabricksKernelName(cluster)); } - static interactiveKernelExists(cluster: iDatabricksCluster): boolean { - if (this.getInteractiveKernel(cluster)) { + static databricksKernelExists(cluster: iDatabricksCluster): boolean { + if (this.getDatabricksKernel(cluster)) { return true; } return false; } static async createKernels(cluster: iDatabricksCluster, logMessages: boolean = true): Promise { - if (!this.notebookKernelExists(cluster)) { + if (!this.jupyterKernelExists(cluster)) { let notebookKernel: DatabricksKernel = new DatabricksKernel(cluster); - this.setKernel(this.getNotebookKernelName(cluster), notebookKernel); + this.setKernel(this.getJupyterKernelName(cluster), notebookKernel); if (logMessages) { ThisExtension.log(`Notebook Kernel for Databricks cluster '${cluster.cluster_id}' created!`) } @@ -101,9 +101,9 @@ export abstract class DatabricksKernelManager { } } - if (!this.interactiveKernelExists(cluster)) { - let interactiveKernel: DatabricksKernel = new DatabricksKernel(cluster, "interactive"); - this.setKernel(this.getInteractiveKernelName(cluster), interactiveKernel); + if (!this.databricksKernelExists(cluster)) { + let interactiveKernel: DatabricksKernel = new DatabricksKernel(cluster, "databricks-notebook"); + this.setKernel(this.getDatabricksKernelName(cluster), interactiveKernel); if (logMessages) { ThisExtension.log(`Interactive Kernel for Databricks cluster '${cluster.cluster_id}' created!`) } @@ -116,8 +116,8 @@ export abstract class DatabricksKernelManager { } static async removeKernels(cluster: iDatabricksCluster, logMessages: boolean = true): Promise { - if (this.notebookKernelExists(cluster)) { - this.removeKernel(this.getNotebookKernelName(cluster)); + if (this.jupyterKernelExists(cluster)) { + this.removeKernel(this.getJupyterKernelName(cluster)); if (logMessages) { ThisExtension.log(`Notebook Kernel for Databricks cluster '${cluster.cluster_id}' removed!`) } @@ -128,8 +128,8 @@ export abstract class DatabricksKernelManager { } } - if (this.interactiveKernelExists(cluster)) { - this.removeKernel(this.getInteractiveKernelName(cluster)); + if (this.databricksKernelExists(cluster)) { + this.removeKernel(this.getDatabricksKernelName(cluster)); if (logMessages) { ThisExtension.log(`Interactive Kernel for Databricks cluster '${cluster.cluster_id}' removed!`) } @@ -142,16 +142,16 @@ export abstract class DatabricksKernelManager { } static async restartClusterKernel(cluster: iDatabricksCluster): Promise { - let kernel: DatabricksKernel = this.getNotebookKernel(cluster) + let kernel: DatabricksKernel = this.getJupyterKernel(cluster) if (kernel) { kernel.restart(); } } - static async restartNotebookKernel(notebook: { notebookEditor: { notebookUri: vscode.Uri } } | undefined | vscode.Uri): Promise { + static async restartJupyterKernel(notebook: { notebookEditor: { notebookUri: vscode.Uri } } | undefined | vscode.Uri): Promise { let notebookUri: vscode.Uri = undefined; - ThisExtension.setStatusBar("Restarting Kernel ...", true); + ThisExtension.setStatusBar("Restarting Jupyter Kernel ...", true); if (notebook instanceof vscode.Uri) { notebookUri = notebook; @@ -164,7 +164,7 @@ export abstract class DatabricksKernelManager { kernel.restart(notebookUri); } - ThisExtension.setStatusBar("Kernel restarted!"); + ThisExtension.setStatusBar("Jupyter Kernel restarted!"); } static async updateWidgets(notebook: { notebookEditor: { notebookUri: vscode.Uri } } | undefined | vscode.Uri): Promise { diff --git a/src/vscode/notebook/DatabricksNotebook.ts b/src/vscode/notebook/DatabricksNotebook.ts new file mode 100644 index 0000000..3887de2 --- /dev/null +++ b/src/vscode/notebook/DatabricksNotebook.ts @@ -0,0 +1,30 @@ +import * as vscode from 'vscode'; + + +export const DatabricksNotebookType: string = 'Databricks-notebook'; + + + +export class DatabricksNotebook extends vscode.NotebookData { + // empty for now, might be extended in the future if new features are added +} + +export class DatabricksNotebookCell extends vscode.NotebookCellData { + + get magic(): string { + if(this.value.startsWith("%")){ + return this.value.split(" ")[0]; + } + } + + set magic(newMagic: string) { + if(this.value.startsWith("%")){ + this.value = this.value.replace(this.magic, newMagic); + } + else + { + this.value = newMagic + "\n" + this.value; + } + } +} + diff --git a/src/vscode/notebook/DatabricksNotebookSerializer.ts b/src/vscode/notebook/DatabricksNotebookSerializer.ts new file mode 100644 index 0000000..234ab3b --- /dev/null +++ b/src/vscode/notebook/DatabricksNotebookSerializer.ts @@ -0,0 +1,157 @@ +import * as vscode from 'vscode'; +import { ThisExtension } from '../../ThisExtension'; +import { DatabricksNotebook, DatabricksNotebookCell } from './DatabricksNotebook'; +import { DatabricksLanguageMapping } from './_types'; + +export class DatabricksNotebookSerializer implements vscode.NotebookSerializer { + public readonly label: string = 'Databricks Notebook Serializer'; + + // language-indepenent header + private readonly HEADER_SUFFIX: string = `Databricks notebook source`; + private readonly MAGIC_PREFIX: string = `MAGIC`; + private readonly CELL_SEPARATOR_SUFFIX: string = "COMMAND ----------"; + + constructor(context: vscode.ExtensionContext) { + context.subscriptions.push( + vscode.workspace.registerNotebookSerializer( + 'databricks-notebook', this, { transientOutputs: true } + ) + ); + } + + private readonly LANGUAGE_MAPPING: DatabricksLanguageMapping[] = [ + { + "databricksLanguage": "python", + "vscodeLanguage": "python", + "magic": "%python", + "commentCharacters": "#", + "fileExtension": ".py" + }, + { + "databricksLanguage": "r", + "vscodeLanguage": "r", + "magic": "%r", + "commentCharacters": "#", + "fileExtension": ".r" + }, + { + "databricksLanguage": "scala", + "vscodeLanguage": "scala", + "magic": "%scala", + "commentCharacters": "//", + "fileExtension": ".scala" + }, + { + "databricksLanguage": "sql", + "vscodeLanguage": "sql", + "magic": "%sql", + "commentCharacters": "--", + "fileExtension": ".sql" + } + ] + ; + + public async deserializeNotebook(data: Uint8Array, token: vscode.CancellationToken): Promise { + var contents = Buffer.from(data).toString(); + + const lines: string[] = contents.split("\n"); + if (lines.length == 0) { + ThisExtension.log("Not a Databricks Notebook source file. Creating new Notebook."); + return { cells: [] }; + } + if (!lines[0].endsWith(this.HEADER_SUFFIX)) { + ThisExtension.log("File is not a valid Databricks Notebook source file."); + throw new Error("File is not a valid Databricks Notebook source file."); + } + + const commentChars = lines[0].split(" ")[0]; + + let notebookLanguage: DatabricksLanguageMapping = this.LANGUAGE_MAPPING.find(x => x.databricksLanguage == "python"); + let cellLanguage: DatabricksLanguageMapping = undefined; + let languages: DatabricksLanguageMapping[] = this.LANGUAGE_MAPPING.filter(x => x.commentCharacters == commentChars); + + if (languages != undefined || languages.length == 1) { + notebookLanguage = languages[0]; + } + else + { + // its Python or R + const rAssignments = contents.split("<-").length; + const pythonAssignments = contents.split("=").length; + + if(rAssignments > pythonAssignments) + { + notebookLanguage = this.LANGUAGE_MAPPING.find(x => x.databricksLanguage == "r"); + } + // else its Python which is the default anyway + } + let notebook: DatabricksNotebook = new DatabricksNotebook([]); + + const splitRegex = new RegExp(`\n\n${commentChars} COMMAND ----------\n\n`, "gm"); + + let rawCells: string[] = lines.slice(1).join("\n").split(splitRegex); + + for (const rawCell of rawCells) { + let cell = new DatabricksNotebookCell(vscode.NotebookCellKind.Code, rawCell, notebookLanguage.vscodeLanguage); + cell.metadata = { "cellLanguage": notebookLanguage }; + + // check for magic + if (rawCell.startsWith(`${commentChars} ${this.MAGIC_PREFIX}`)) { + let firstLine = rawCell.split("\n")[0]; + let firstLineValues = firstLine.split(/\s+/gm); + let magic = firstLineValues[2]; + + if(magic == "%md") + { + cell.kind = vscode.NotebookCellKind.Markup; + cell.value = cell.value.replace(new RegExp(`^${commentChars} ${this.MAGIC_PREFIX} ${magic}\n`, "gm"), ""); + } + else + { + cellLanguage = this.LANGUAGE_MAPPING.find(x => x.magic == magic); + cell.metadata = { "cellLanguage": cellLanguage }; + } + + cell.value = cell.value.replace(new RegExp(`^${commentChars} ${this.MAGIC_PREFIX} `, "gm"), ""); + } + cell.languageId = cell.metadata.cellLanguage.vscodeLanguage; + notebook.cells.push(cell); + } + + + // set metadata ? + // notebook.metadata. ... + notebook.metadata = { "notebookLanguage": notebookLanguage }; + + return notebook; + } + + public async serializeNotebook(data: DatabricksNotebook, token: vscode.CancellationToken): Promise { + // Map the Notebook data into the format we want to save the Notebook data as + let notebook: DatabricksNotebook = data; + + let notebookLanguage: DatabricksLanguageMapping = notebook.metadata.notebookLanguage; + + let finalCells: vscode.NotebookCellData[] = []; + + + for (const cell of notebook.cells as DatabricksNotebookCell[]) { + if (cell.kind == vscode.NotebookCellKind.Markup) { + cell.magic = "%md"; + cell.value = `${notebookLanguage.commentCharacters} ${this.MAGIC_PREFIX} ${cell.magic}\n${cell.value}`; + } + + if(cell.magic) + { + cell.value = cell.value.replace("\n", `\n${notebookLanguage.commentCharacters} ${this.MAGIC_PREFIX} `); + } + } + + let finalLines: string[] = []; + finalLines.push(`${notebookLanguage.commentCharacters} ${this.HEADER_SUFFIX}`); + finalLines = finalLines.concat(notebook.cells.flatMap(x => x.value)) + + // Give a string of all the data to save and VS Code will handle the rest + return await Buffer.from(finalLines.join(`\n\n${notebookLanguage.commentCharacters} COMMAND ----------\n\n`)); + } +} \ No newline at end of file diff --git a/src/vscode/notebook/_types.ts b/src/vscode/notebook/_types.ts new file mode 100644 index 0000000..c8e656b --- /dev/null +++ b/src/vscode/notebook/_types.ts @@ -0,0 +1,7 @@ +export interface DatabricksLanguageMapping { + databricksLanguage: string, + vscodeLanguage: string, + magic: string, + commentCharacters: string, + fileExtension: string +} \ No newline at end of file diff --git a/src/vscode/treeviews/clusters/DatabricksCluster.ts b/src/vscode/treeviews/clusters/DatabricksCluster.ts index f0eb9d3..7e8d701 100644 --- a/src/vscode/treeviews/clusters/DatabricksCluster.ts +++ b/src/vscode/treeviews/clusters/DatabricksCluster.ts @@ -83,7 +83,7 @@ export class DatabricksCluster extends DatabricksClusterTreeItem { states.push("STOPPED"); } - if (this.NotebookKernelExists) { + if (this.JupyterKernelExists) { states.push("KERNEL"); } else { @@ -139,23 +139,23 @@ export class DatabricksCluster extends DatabricksClusterTreeItem { return this._source; } - private get NotebookKernel(): DatabricksKernel { - return DatabricksKernelManager.getNotebookKernel(this.definition); + private get JupyterKernel(): DatabricksKernel { + return DatabricksKernelManager.getJupyterKernel(this.definition); } - public get NotebookKernelExists(): boolean { - if (this.NotebookKernel) { + public get JupyterKernelExists(): boolean { + if (this.JupyterKernel) { return true; } return false; } - private get InteractiveKernel(): DatabricksKernel { - return DatabricksKernelManager.getNotebookKernel(this.definition); + private get DatabricksKernel(): DatabricksKernel { + return DatabricksKernelManager.getJupyterKernel(this.definition); } - public get InteractiveKernelExists(): boolean { - if (this.InteractiveKernel) { + public get DatabricksKernelExists(): boolean { + if (this.DatabricksKernel) { return true; } return false; @@ -191,7 +191,7 @@ export class DatabricksCluster extends DatabricksClusterTreeItem { vscode.window.showErrorMessage(`ERROR: ${error}`); }); - let kernel = this.NotebookKernel; + let kernel = this.JupyterKernel; if (kernel) { kernel.disposeController(); }