Add ability to run ADS commands

src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts

@@ -193,7 +193,12 @@ configurationRegistry.registerConfiguration({
 			'type': 'boolean',
 			'default': true,
 			'description': localize('notebook.sqlStopOnError', "SQL kernel: stop Notebook execution when error occurs in a cell.")
-		}
+		},
+		'notebook.allowCommandExecution': {
+			'type': 'boolean',
+			'description': localize('notebook.allowCommandExecution', "Allow notebooks to execute Azure Data Studio Commands using the %%AZDATA_EXECUTE_COMMAND magic."),
+			'default': false
+		},
 	}
 });
 

src/sql/workbench/services/notebook/browser/models/cell.ts

@@ -25,14 +25,17 @@ 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 { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { Disposable } from 'vs/base/common/lifecycle';
 
 let modelId = 0;
 
-export class CellModel implements ICellModel {
+export class CellModel extends Disposable implements ICellModel {
 	public id: string;
 
 	private _cellType: nb.CellType;
-	private _source: string | string[];
+	private _source: string[];
 	private _language: string;
 	private _cellGuid: string;
 	private _future: FutureInternal;
@@ -56,18 +59,22 @@ export class CellModel implements ICellModel {
 	private _isCollapsed: boolean;
 	private _onCollapseStateChanged = new Emitter<boolean>();
 	private _modelContentChangedEvent: IModelContentChangedEvent;
+	private _isCommandExecutionSettingEnabled: boolean;
 
 	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
 			this.fromJSON(cellData);
 		} else {
 			this._cellType = CellTypes.Code;
-			this._source = '';
+			this._source = [''];
 		}
 
 		this._isEditMode = this._cellType !== CellTypes.Markdown;
@@ -80,6 +87,14 @@ 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();
+		let commandExecutionSettingValue: boolean = this._configurationService.getValue('notebook.allowCommandExecution');
+		this._isCommandExecutionSettingEnabled = commandExecutionSettingValue ? commandExecutionSettingValue : false;
+
+		this._register(this._configurationService.onDidChangeConfiguration(e => {
+			if (e.affectsConfiguration('notebook.allowCommandExecution')) {
+				this._isCommandExecutionSettingEnabled = this._configurationService.getValue('notebook.allowCommandExecution');
+			}
+		}));
 	}
 
 	public equals(other: ICellModel) {
@@ -336,19 +351,21 @@ export class CellModel implements ICellModel {
 
 					// requestExecute expects a string for the code parameter
 					content = Array.isArray(content) ? content.join('') : content;
-					const future = kernel.requestExecute({
+					if (!this.checkForAdsCommandMagic()) {
-						code: content,
+						const future = kernel.requestExecute({
-						stop_on_error: true
+							code: content,
-					}, false);
+							stop_on_error: true
-					this.setFuture(future as FutureInternal);
+						}, false);
-					this.fireExecutionStateChanged();
+						this.setFuture(future as FutureInternal);
-					// For now, await future completion. Later we should just track and handle cancellation based on model notifications
+						this.fireExecutionStateChanged();
-					let result: nb.IExecuteReplyMsg = <nb.IExecuteReplyMsg><any>await future.done;
+						// For now, await future completion. Later we should just track and handle cancellation based on model notifications
-					if (result && result.content) {
+						let result: nb.IExecuteReplyMsg = <nb.IExecuteReplyMsg><any>await future.done;
-						this.executionCount = result.content.execution_count;
+						if (result && result.content) {
-						if (result.content.status !== 'ok') {
+							this.executionCount = result.content.execution_count;
-							// TODO track error state
+							if (result.content.status !== 'ok') {
-							return false;
+								// TODO track error state
+								return false;
+							}
 						}
 					}
 				}
@@ -682,7 +699,7 @@ export class CellModel implements ICellModel {
 	}
 
 
-	private getMultilineSource(source: string | string[]): string | string[] {
+	private getMultilineSource(source: string | string[]): string[] {
 		if (source === undefined) {
 			return [];
 		}
@@ -707,6 +724,29 @@ export class CellModel implements ICellModel {
 		return source;
 	}
 
+	// Run ADS commands via a notebook. Configured by a setting (turned off by default).
+	private checkForAdsCommandMagic(): boolean {
+		if (!this._isCommandExecutionSettingEnabled) {
+			return false;
+		}
+
+		let executeCommandMagic = '%%AZDATA_EXECUTE_COMMAND';
+		if (this._source && this._source.length) {
+			if (this._source[0].startsWith(executeCommandMagic)) {
+				let commandNamePlusArgs = this._source[0].replace(executeCommandMagic + ' ', '');
+				if (commandNamePlusArgs) {
+					let commandName = commandNamePlusArgs.split(' ')[0];
+					if (commandName) {
+						let args = commandNamePlusArgs.replace(commandName + ' ', '');
+						this._commandService.executeCommand(commandName, [args]);
+					}
+					return true;
+				}
+			}
+		}
+		return false;
+	}
+
 	// Dispose and set current future to undefined
 	private disposeFuture() {
 		if (this._future) {