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

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;