From 061d3510b9195db4f3674d835864c751a9ab38aa Mon Sep 17 00:00:00 2001 From: Chris LaFreniere <40371649+chlafreniere@users.noreply.github.com> Date: Sun, 28 Jun 2020 21:17:50 -0700 Subject: [PATCH] Run ADS Commands from Notebooks (#11069) * Run commands * cleanup * Plumbed through error and message output * undo unnecessary changes * Remove change to fix tests * pr feedback --- .../notebook/browser/notebook.contribution.ts | 13 +++ .../electron-browser/notebookUtils.test.ts | 47 ++++++++- .../services/notebook/browser/models/cell.ts | 97 ++++++++++++++++--- .../services/notebook/browser/utils.ts | 34 +++++++ 4 files changed, 175 insertions(+), 16 deletions(-) diff --git a/src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts b/src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts index af34a86d4d..aa10b2c88b 100644 --- a/src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts +++ b/src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts @@ -213,6 +213,19 @@ configurationRegistry.registerConfiguration({ } }); +configurationRegistry.registerConfiguration({ + 'id': 'notebook', + 'title': 'Notebook', + 'type': 'object', + 'properties': { + 'notebook.allowAzureDataStudioCommands': { + 'type': 'boolean', + 'default': false, + 'description': localize('notebook.allowADSCommands', "Allow notebooks to run Azure Data Studio commands.") + } + } +}); + /* *************** Output components *************** */ // Note: most existing types use the same component to render. In order to // preserve correct rank order, we register it once for each different rank of diff --git a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookUtils.test.ts b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookUtils.test.ts index c206162048..12cb3a7730 100644 --- a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookUtils.test.ts +++ b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookUtils.test.ts @@ -10,7 +10,7 @@ import { nb, ServerInfo } from 'azdata'; import { getHostAndPortFromEndpoint, isStream, getProvidersForFileName, asyncForEach, clusterEndpointsProperty, getClusterEndpoints, RawEndpoint, IEndpoint, getStandardKernelsForProvider, IStandardKernelWithProvider, rewriteUrlUsingRegex } from 'sql/workbench/services/notebook/browser/models/notebookUtils'; import { INotebookService, DEFAULT_NOTEBOOK_FILETYPE, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/browser/notebookService'; import { NotebookServiceStub } from 'sql/workbench/contrib/notebook/test/stubs'; -import { tryMatchCellMagic } from 'sql/workbench/services/notebook/browser/utils'; +import { tryMatchCellMagic, extractCellMagicCommandPlusArgs } from 'sql/workbench/services/notebook/browser/utils'; suite('notebookUtils', function (): void { const mockNotebookService = TypeMoq.Mock.ofType(NotebookServiceStub); @@ -121,6 +121,51 @@ suite('notebookUtils', function (): void { assert.strictEqual(result, null); }); + test('extractCellMagicCommandPlusArgs Test', async function (): Promise { + let result = extractCellMagicCommandPlusArgs(undefined, undefined); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('test', undefined); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs(undefined, 'test'); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('%%magic', 'magic'); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('%%magic ', 'magic'); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('magic', 'magic'); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('magic ', 'magic'); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('%%magic command', 'otherMagic'); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('%%magiccommand', 'magic'); + assert.strictEqual(result, undefined); + + result = extractCellMagicCommandPlusArgs('%%magic command', 'magic'); + assert.equal(result.commandId, 'command'); + assert.equal(result.args, ''); + + result = extractCellMagicCommandPlusArgs('%%magic command arg1', 'magic'); + assert.equal(result.commandId, 'command'); + assert.equal(result.args, 'arg1'); + + result = extractCellMagicCommandPlusArgs('%%magic command arg1 arg2', 'magic'); + assert.equal(result.commandId, 'command'); + assert.equal(result.args, 'arg1 arg2'); + + result = extractCellMagicCommandPlusArgs('%%magic command.id arg1 arg2 arg3', 'magic'); + assert.equal(result.commandId, 'command.id'); + assert.equal(result.args, 'arg1 arg2 arg3'); + }); + test('asyncForEach Test', async function (): Promise { let totalResult = 0; await asyncForEach([1, 2, 3, 4], async (value) => { diff --git a/src/sql/workbench/services/notebook/browser/models/cell.ts b/src/sql/workbench/services/notebook/browser/models/cell.ts index 528eeb529d..fdb2ddff7e 100644 --- a/src/sql/workbench/services/notebook/browser/models/cell.ts +++ b/src/sql/workbench/services/notebook/browser/models/cell.ts @@ -25,10 +25,15 @@ import { IModelContentChangedEvent } from 'vs/editor/common/model/textModelEvent import { firstIndex, find } from 'vs/base/common/arrays'; import { HideInputTag } from 'sql/platform/notebooks/common/outputRegistry'; import { FutureInternal, notebookConstants } from 'sql/workbench/services/notebook/browser/interfaces'; +import { ICommandService } from 'vs/platform/commands/common/commands'; +import { tryMatchCellMagic, extractCellMagicCommandPlusArgs } from 'sql/workbench/services/notebook/browser/utils'; +import { IConfigurationService } from 'vs/platform/configuration/common/configuration'; +import { Disposable } from 'vs/base/common/lifecycle'; let modelId = 0; +const ads_execute_command = 'ads_execute_command'; -export class CellModel implements ICellModel { +export class CellModel extends Disposable implements ICellModel { public id: string; private _cellType: nb.CellType; @@ -58,11 +63,15 @@ export class CellModel implements ICellModel { private _modelContentChangedEvent: IModelContentChangedEvent; private _showPreview: boolean = true; private _onCellPreviewChanged = new Emitter(); + private _isCommandExecutionSettingEnabled: boolean = false; constructor(cellData: nb.ICellContents, private _options: ICellModelOptions, - @optional(INotebookService) private _notebookService?: INotebookService + @optional(INotebookService) private _notebookService?: INotebookService, + @optional(ICommandService) private _commandService?: ICommandService, + @optional(IConfigurationService) private _configurationService?: IConfigurationService ) { + super(); this.id = `${modelId++}`; if (cellData) { // Read in contents if available @@ -82,6 +91,15 @@ export class CellModel implements ICellModel { // if the fromJson() method was already called and _cellGuid was previously set, don't generate another UUID unnecessarily this._cellGuid = this._cellGuid || generateUuid(); this.createUri(); + if (this._configurationService) { + const allowADSCommandsKey = 'notebook.allowAzureDataStudioCommands'; + this._isCommandExecutionSettingEnabled = this._configurationService.getValue(allowADSCommandsKey); + this._register(this._configurationService.onDidChangeConfiguration(e => { + if (e.affectsConfiguration(allowADSCommandsKey)) { + this._isCommandExecutionSettingEnabled = this._configurationService.getValue(allowADSCommandsKey); + } + })); + } } public equals(other: ICellModel) { @@ -353,19 +371,42 @@ export class CellModel implements ICellModel { // requestExecute expects a string for the code parameter content = Array.isArray(content) ? content.join('') : content; - const future = kernel.requestExecute({ - code: content, - stop_on_error: true - }, false); - this.setFuture(future as FutureInternal); - this.fireExecutionStateChanged(); - // For now, await future completion. Later we should just track and handle cancellation based on model notifications - let result: nb.IExecuteReplyMsg = await future.done; - if (result && result.content) { - this.executionCount = result.content.execution_count; - if (result.content.status !== 'ok') { - // TODO track error state - return false; + if (tryMatchCellMagic(this.source[0]) !== ads_execute_command || !this._isCommandExecutionSettingEnabled) { + const future = kernel.requestExecute({ + code: content, + stop_on_error: true + }, false); + this.setFuture(future as FutureInternal); + this.fireExecutionStateChanged(); + // For now, await future completion. Later we should just track and handle cancellation based on model notifications + let result: nb.IExecuteReplyMsg = await future.done; + if (result && result.content) { + this.executionCount = result.content.execution_count; + if (result.content.status !== 'ok') { + // TODO track error state + return false; + } + } + } else { + let result = extractCellMagicCommandPlusArgs(this._source[0], ads_execute_command); + // Similar to the markdown renderer, we should not allow downloadResource here + if (result?.commandId !== '_workbench.downloadResource') { + try { + // Need to reset outputs here (kernels do this on their own) + this._outputs = []; + let commandExecuted = this._commandService?.executeCommand(result.commandId, result.args); + // This will ensure that the run button turns into a stop button + this.fireExecutionStateChanged(); + await commandExecuted; + // For save files, if we output a message after saving the file, the file becomes dirty again. + // Special casing this to avoid this particular issue. + if (result?.commandId !== 'workbench.action.files.saveFiles') { + this.handleIOPub(this.toIOPubMessage(false)); + } + } catch (error) { + this.handleIOPub(this.toIOPubMessage(true, error?.message)); + return false; + } } } } @@ -711,6 +752,32 @@ export class CellModel implements ICellModel { return source; } + // Create an iopub message to display either a display result or an error result, + // in order to be displayed as part of a cell's outputs + private toIOPubMessage(isError: boolean, message?: string): nb.IIOPubMessage { + return { + channel: 'iopub', + type: 'iopub', + header: { + msg_id: undefined, + msg_type: isError ? 'error' : 'display_data' + }, + content: isError ? { + output_type: 'error', + evalue: message, + ename: '', + traceback: [] + } : { + output_type: 'execute_result', + data: { + 'text/html': localize('commandSuccessful', "Command executed successfully"), + } + }, + metadata: undefined, + parent_header: undefined + }; + } + // Dispose and set current future to undefined private disposeFuture() { if (this._future) { diff --git a/src/sql/workbench/services/notebook/browser/utils.ts b/src/sql/workbench/services/notebook/browser/utils.ts index e9d83db898..14cf97ac11 100644 --- a/src/sql/workbench/services/notebook/browser/utils.ts +++ b/src/sql/workbench/services/notebook/browser/utils.ts @@ -13,3 +13,37 @@ export function tryMatchCellMagic(input: string): string { let magicName = match && match[1]; return magicName; } + +/** + * When a cell is formatted in the following way, extract the commandId and args: + * %%ads_execute_command commandId arg1 arg2 + * Extract the commandId and the two args + * @param input cell source + * @param magicName magic name + */ +export function extractCellMagicCommandPlusArgs(input: string, magicName: string): ICommandPlusArgs | undefined { + if (input && magicName && input.startsWith(`%%${magicName}`)) { + let commandNamePlusArgs = input.replace(`%%${magicName}`, ''); + if (commandNamePlusArgs?.startsWith(' ')) { + // There needs to be a space between the magic name and the command id + commandNamePlusArgs = commandNamePlusArgs.slice(1); + let commandName = commandNamePlusArgs.split(' ')[0]; + if (commandName) { + let args = commandNamePlusArgs.replace(commandName, ''); + if (args?.startsWith(' ')) { + args = args.slice(1); + } + return { + commandId: commandName, + args: args + }; + } + } + } + return undefined; +} + +export interface ICommandPlusArgs { + commandId: string; + args: string; +}