Enable Parameterization via Notebook URI (#13869)

* Enable Parameterization via Notebook URI

* Add Parameterization and Move notebookModel tests

* minor typos

* parameter typo

* Multiple parameters through uri fix

* Address PR comments and tests fix

* add new injected parameters after original injected parameters

* Add new test to verify notebookURi parameter is injected after both parameter and injected cell

* fix tests

extensions/notebook/src/protocol/notebookUriHandler.ts

@@ -8,7 +8,6 @@ import * as azdata from 'azdata';
 
 import * as request from 'request';
 import * as path from 'path';
-import * as querystring from 'querystring';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
 import { IQuestion, QuestionTypes } from '../prompts/question';
@@ -33,15 +32,42 @@ export class NotebookUriHandler implements vscode.UriHandler {
 				vscode.window.showErrorMessage(localize('notebook.unsupportedAction', "Action {0} is not supported for this handler", uri.path));
 		}
 	}
-
+	/**
+	 * Our Azure Data Studio URIs follow the standard URI format, and we currently only support https and http URI schemes
+	 * azuredatastudio://microsoft.notebook/open?url=https://
+	 * azuredatastudio://microsoft.notebook/open?url=http://
+	 *
+	 * Example of URI (encoded):
+	 * azuredatastudio://microsoft.notebook/open?url=https%3A%2F%2Fraw.githubusercontent.com%2FVasuBhog%2FAzureDataStudio-Notebooks%2Fmain%2FDemo_Parameterization%2FInput.ipynb
+	 *
+	 * We also support parameters added to the URI for parameterization scenarios
+	 *
+	 * Parameters via the URI are formatted by adding the parameters after the .ipynb with a
+	 * query '?' and use '&' to distinguish a new parameter
+	 *
+	 * Example of Parameters query:
+	 * ...Input.ipynb?x=1&y=2'
+	 *
+	 * Encoded URI with parameters:
+	 * azuredatastudio://microsoft.notebook/open?url=https%3A%2F%2Fraw.githubusercontent.com%2FVasuBhog%2FAzureDataStudio-Notebooks%2Fmain%2FDemo_Parameterization%2FInput.ipynb%3Fx%3D1%26y%3D2
+	 * Decoded URI with parameters:
+	 * azuredatastudio://microsoft.notebook/open?url=https://raw.githubusercontent.com/VasuBhog/AzureDataStudio-Notebooks/main/Demo_Parameterization/Input.ipynb?x=1&y=2
+	 */
 	private open(uri: vscode.Uri): Promise<void> {
-		const data = querystring.parse(uri.query);
+		let data: string;
+		// We ensure that the URI is formatted properly
+		let urlIndex = uri.query.indexOf('url=');
+		if (urlIndex >= 0) {
+			// Querystring can not be used as it incorrectly turns parameters attached
+			// to the URI query into key/value pairs and would then fail to open the URI
+			data = uri.query.substr(urlIndex + 4);
+		}
 
-		if (!data.url) {
+		if (!data) {
 			console.warn('Failed to open URI:', uri);
 		}
 
-		return this.openNotebook(data.url);
+		return this.openNotebook(data);
 	}
 
 	private async openNotebook(url: string | string[]): Promise<void> {
@@ -70,7 +96,8 @@ export class NotebookUriHandler implements vscode.UriHandler {
 			}
 
 			let contents = await this.download(url);
-			let untitledUri = this.getUntitledUri(path.basename(uri.fsPath));
+			let untitledUriPath = this.getUntitledUriPath(path.basename(uri.fsPath));
+			let untitledUri = uri.with({ authority: '', scheme: 'untitled', path: untitledUriPath });
 			if (path.extname(uri.fsPath) === '.ipynb') {
 				await azdata.nb.showNotebookDocument(untitledUri, {
 					initialContent: contents,
@@ -112,7 +139,7 @@ export class NotebookUriHandler implements vscode.UriHandler {
 		});
 	}
 
-	private getUntitledUri(originalTitle: string): vscode.Uri {
+	private getUntitledUriPath(originalTitle: string): string {
 		let title = originalTitle;
 		let nextVal = 0;
 		let ext = path.extname(title);
@@ -126,6 +153,6 @@ export class NotebookUriHandler implements vscode.UriHandler {
 			}
 			nextVal++;
 		}
-		return vscode.Uri.parse(`untitled:${title}`);
+		return title;
 	}
 }

src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts

@@ -81,7 +81,31 @@ let expectedNotebookContentOneCell: nb.INotebookContents = {
 	nbformat_minor: 5
 };
 
+let expectedParameterizedNotebookContent: nb.INotebookContents = {
+	cells: [{
+		cell_type: CellTypes.Code,
+		source: ['x = 1 \ny = 2'],
+		metadata: { language: 'python', tags: ['parameters'] },
+		execution_count: 1
+	}, {
+		cell_type: CellTypes.Code,
+		source: ['x = 2 \ny = 5)'],
+		metadata: { language: 'python', tags: ['injected-parameters'] },
+		execution_count: 2
+	}],
+	metadata: {
+		kernelspec: {
+			name: 'python3',
+			language: 'python',
+			display_name: 'Python 3'
+		}
+	},
+	nbformat: 4,
+	nbformat_minor: 5
+};
+
 let defaultUri = URI.file('/some/path.ipynb');
+let notebookUriParams = URI.parse('https://hello.ipynb?x=1&y=2');
 
 let mockClientSession: IClientSession;
 let clientSessionOptions: IClientSessionOptions;
@@ -318,6 +342,139 @@ suite('notebook model', function (): void {
 		assert.equal(activeCellChangeCount, 5, 'Active cell change count is incorrect; should be 5');
 	});
 
+	test('Should set notebook parameter and injected parameter cell correctly', async function (): Promise<void> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedParameterizedNotebookContent));
+		defaultModelOptions.notebookUri = defaultUri;
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		assert.equal(model.notebookUri, defaultModelOptions.notebookUri, 'Notebook model has incorrect URI');
+		assert.equal(model.cells.length, 2, 'Cell count in notebook model is not correct');
+
+		// Set parameter cell and injected parameters cell
+		let notebookParamsCell = model.cells[0];
+		let notebookInjectedParamsCell = model.cells[1];
+
+		// Parameters Cell Validation
+		assert.equal(model.cells.indexOf(notebookParamsCell), 0, 'Notebook parameters cell should be first cell in notebook');
+		assert.equal(notebookParamsCell.isParameter, true, 'Notebook parameters cell should be tagged parameter');
+		assert.equal(notebookParamsCell.isInjectedParameter, false, 'Notebook parameters cell should not be tagged injected parameter');
+
+		// Injected Parameters Cell Validation
+		assert.equal(model.cells.indexOf(notebookInjectedParamsCell), 1, 'Notebook injected parameters cell should be second cell in notebook');
+		assert.equal(notebookInjectedParamsCell.isParameter, false, 'Notebook injected parameters cell should not be tagged parameter cell');
+		assert.equal(notebookInjectedParamsCell.isInjectedParameter, true, 'Notebook injected parameters cell should be tagged injected parameter');
+	});
+
+	test('Should set notebookUri parameters to new cell correctly', async function (): Promise<void> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContentOneCell));
+		defaultModelOptions.notebookUri = notebookUriParams;
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		assert.equal(model.notebookUri, defaultModelOptions.notebookUri, 'Notebook model has incorrect URI');
+		assert.equal(model.cells.length, 2, 'Cell count in notebook model is not correct');
+
+		// Validate notebookUri parameter cell is set as the only parameter cell
+		let notebookUriParamsCell = model.cells[0];
+		assert.equal(model.cells.indexOf(notebookUriParamsCell), 0, 'NotebookURI parameters cell should be first cell in notebook');
+		assert.equal(notebookUriParamsCell.isParameter, true, 'NotebookURI parameters cell should be tagged parameter');
+		assert.equal(notebookUriParamsCell.isInjectedParameter, false, 'NotebookURI parameters Cell should not be injected parameter');
+	});
+
+	test('Should set notebookUri parameters to new cell after parameters cell correctly', async function (): Promise<void> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		let expectedNotebookContentParameterCell = expectedNotebookContentOneCell;
+		//Set the cell to be tagged as parameter cell
+		expectedNotebookContentParameterCell.cells[0].metadata.tags = ['parameters'];
+
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContentParameterCell));
+		defaultModelOptions.notebookUri = notebookUriParams;
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		assert.equal(model.notebookUri, defaultModelOptions.notebookUri, 'Notebook model has incorrect URI');
+		assert.equal(model.cells.length, 2, 'Cell count in notebook model is not correct');
+
+		// Validate notebookUri parameter cell is set as injected parameter
+		let notebookUriParamsCell = model.cells[1];
+		assert.equal(model.cells.indexOf(notebookUriParamsCell), 1, 'NotebookURI parameters cell should be second cell in notebook');
+		assert.equal(notebookUriParamsCell.isParameter, false, 'NotebookURI parameters cell should not be tagged parameter cell');
+		assert.equal(notebookUriParamsCell.isInjectedParameter, true, 'NotebookURI parameters Cell should be injected parameter');
+	});
+
+	test('Should set notebookUri parameters to new cell after injected parameters cell correctly', async function (): Promise<void> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedParameterizedNotebookContent));
+		defaultModelOptions.notebookUri = notebookUriParams;
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		assert.equal(model.notebookUri, defaultModelOptions.notebookUri, 'Notebook model has incorrect URI');
+		assert.equal(model.cells.length, 3, 'Cell count in notebook model is not correct');
+
+		// Validate notebookUri parameter cell is set as an injected parameter after parameter and injected parameter cells
+		let notebookUriParamsCell = model.cells[2];
+		assert.equal(model.cells.indexOf(notebookUriParamsCell), 2, 'NotebookURI parameters cell should be third cell in notebook');
+		assert.equal(notebookUriParamsCell.isParameter, false, 'NotebookURI parameters cell should not be tagged parameter cell');
+		assert.equal(notebookUriParamsCell.isInjectedParameter, true, 'NotebookURI parameters Cell should be injected parameter');
+	});
+
+	test('Should move first cell below second cell correctly', async function (): Promise<void> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		assert.equal(model.notebookUri, defaultModelOptions.notebookUri, 'Notebook model has incorrect URI');
+		assert.equal(model.cells.length, 2, 'Cell count in notebook model is not correct');
+
+		let firstCell = model.cells[0];
+		let secondCell = model.cells[1];
+		// Move First Cell down
+		model.moveCell(firstCell, 1);
+		assert.equal(model.cells.indexOf(firstCell), 1, 'First Cell did not move down correctly');
+		assert.equal(model.cells.indexOf(secondCell), 0, 'Second Cell did not move up correctly');
+	});
+
+	test('Should move second cell up above the first cell correctly', async function (): Promise<void> {
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		assert.equal(model.notebookUri, defaultModelOptions.notebookUri, 'Notebook model has incorrect URI');
+		assert.equal(model.cells.length, 2, 'Cell count in notebook model is not correct');
+
+		let firstCell = model.cells[0];
+		let secondCell = model.cells[1];
+		// Move Second Cell up
+		model.moveCell(secondCell, 0);
+		assert.equal(model.cells.indexOf(firstCell), 1, 'First Cell did not move down correctly');
+		assert.equal(model.cells.indexOf(secondCell), 0, 'Second Cell did not move up correctly');
+	});
+
 	test('Should delete cells correctly', async function (): Promise<void> {
 		// Given a notebook with 2 cells
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);

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

@@ -49,6 +49,7 @@ export class ErrorInfo {
 }
 
 const saveConnectionNameConfigName = 'notebook.saveConnectionName';
+const injectedParametersMsg = localize('injectedParametersMsg', '# Injected-Parameters\n');
 
 export class NotebookModel extends Disposable implements INotebookModel {
 	private _contextsChangedEmitter = new Emitter<void>();
@@ -400,9 +401,20 @@ export class NotebookModel extends Disposable implements INotebookModel {
 						}
 					});
 				}
+				// Modify Notebook URI Params format from URI query to string space delimited format
+				let notebookUriParams: string = this.notebookUri.query;
+				notebookUriParams = notebookUriParams.replace(/&/g, '\n').replace(/=/g, ' = ');
+				// Get parameter cell and index to place new notebookUri parameters accordingly
+				let parameterCellIndex = 0;
+				let hasParameterCell = false;
+				let hasInjectedCell = false;
 				if (contents.cells && contents.cells.length > 0) {
 					this._cells = contents.cells.map(c => {
 						let cellModel = factory.createCell(c, { notebook: this, isTrusted: isTrusted });
+						if (cellModel.isParameter) {
+							parameterCellIndex = contents.cells.indexOf(c);
+							hasParameterCell = true;
+						}
 						/*
 						In a parameterized notebook there will be an injected parameter cell.
 						Papermill originally inserts the injected parameter with the comment "# Parameters"
@@ -410,13 +422,18 @@ export class NotebookModel extends Disposable implements INotebookModel {
 						So to make it clear we edit the injected parameters comment to indicate it is the Injected-Parameters cell.
 						*/
 						if (cellModel.isInjectedParameter) {
+							hasInjectedCell = true;
 							cellModel.source = cellModel.source.slice(1);
-							cellModel.source = ['# Injected-Parameters\n'].concat(cellModel.source);
+							cellModel.source = [injectedParametersMsg].concat(cellModel.source);
 						}
 						this.trackMarkdownTelemetry(<nb.ICellContents>c, cellModel);
 						return cellModel;
 					});
 				}
+				// Only add new parameter cell if notebookUri Parameters are found
+				if (notebookUriParams) {
+					this.addUriParameterCell(notebookUriParams, hasParameterCell, parameterCellIndex, hasInjectedCell);
+				}
 			}
 
 			// Trust notebook by default if there are no code cells
@@ -431,6 +448,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			throw error;
 		}
 	}
+
 	public async requestModelLoad(): Promise<void> {
 		try {
 			this.setDefaultKernelAndProviderId();
@@ -487,6 +505,33 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return cell;
 	}
 
+	/**
+	 * Adds Parameters cell based on Notebook URI parameters
+	 * @param notebookUriParams contains the parameters from Notebook URI
+	 * @param hasParameterCell notebook contains a parameter cell
+	 * @param parameterCellIndex index of the parameter cell in notebook
+	 * @param hasInjectedCell notebook contains a injected parameter cell
+	 */
+	private addUriParameterCell(notebookUriParams: string, hasParameterCell: boolean, parameterCellIndex: number, hasInjectedCell: boolean): void {
+		let uriParamsIndex = parameterCellIndex;
+		// Set new uri parameters as a Injected Parameters cell after original parameter cell
+		if (hasParameterCell) {
+			uriParamsIndex = parameterCellIndex + 1;
+			// Set the uri parameters after the injected parameter cell
+			if (hasInjectedCell) {
+				uriParamsIndex = uriParamsIndex + 1;
+			}
+			this.addCell('code', uriParamsIndex);
+			this.cells[uriParamsIndex].isInjectedParameter = true;
+			this.cells[uriParamsIndex].source = [injectedParametersMsg].concat(notebookUriParams);
+		} else {
+			// Set new parameters as the parameters cell as the first cell in the notebook
+			this.addCell('code', uriParamsIndex);
+			this.cells[uriParamsIndex].isParameter = true;
+			this.cells[uriParamsIndex].source = [notebookUriParams];
+		}
+	}
+
 	moveCell(cell: ICellModel, direction: MoveDirection): void {
 		if (this.inErrorState) {
 			return;