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

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

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<INotebookService>(NotebookServiceStub);
@@ -121,6 +121,51 @@ suite('notebookUtils', function (): void {
 		assert.strictEqual(result, null);
 	});
 
+	test('extractCellMagicCommandPlusArgs Test', async function (): Promise<void> {
+		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<void> {
 		let totalResult = 0;
 		await asyncForEach([1, 2, 3, 4], async (value) => {

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<boolean>();
+	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 = <nb.IExecuteReplyMsg><any>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 = <nb.IExecuteReplyMsg><any>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: <nb.IHeader>{
+				msg_id: undefined,
+				msg_type: isError ? 'error' : 'display_data'
+			},
+			content: isError ? <nb.IErrorResult>{
+				output_type: 'error',
+				evalue: message,
+				ename: '',
+				traceback: []
+			} : <nb.IDisplayResult>{
+				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) {

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;
+}