Notebooks: add multi_connection field to notebook metadata (#14097)

* add multi_connection to nb metadata

* address PR comments

* pr comments

* adsd notebookmetadatakeys object

* pr comments

* use isUUID util method

src/sql/azdata.proposed.d.ts

@@ -75,6 +75,7 @@ declare module 'azdata' {
 
 		export interface INotebookMetadata {
 			connection_name?: string;
+			multi_connection_mode?: boolean;
 		}
 
 		export interface ICellMetadata {

src/sql/workbench/contrib/notebook/test/stubs.ts

@@ -76,6 +76,9 @@ export class NotebookModelStub implements INotebookModel {
 	get savedConnectionName(): string {
 		throw new Error('method not implemented.');
 	}
+	get multiConnectionMode(): boolean {
+		throw new Error('method not implemented.');
+	}
 	get providerId(): string {
 		throw new Error('method not implemented.');
 	}

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

@@ -308,6 +308,12 @@ export interface INotebookModel {
 	 * or undefined if none.
 	 */
 	readonly savedConnectionName: string | undefined;
+
+	/**
+	 * The connection mode of the notebook (single vs multiple connections)
+	 */
+	multiConnectionMode: boolean;
+
 	/**
 	 * Event fired on first initialization of the cells and
 	 * on subsequent change events

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

@@ -32,6 +32,7 @@ import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilit
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { values } from 'vs/base/common/collections';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { isUUID } from 'vs/base/common/uuid';
 
 /*
 * Used to control whether a message in a dialog/wizard is displayed as an error,
@@ -48,6 +49,19 @@ export class ErrorInfo {
 	}
 }
 
+interface INotebookMetadataInternal extends nb.INotebookMetadata {
+	azdata_notebook_guid?: string;
+}
+
+type NotebookMetadataKeys = Required<nb.INotebookMetadata>;
+const expectedMetadataKeys: NotebookMetadataKeys = {
+	kernelspec: undefined,
+	language_info: undefined,
+	tags: undefined,
+	connection_name: undefined,
+	multi_connection_mode: undefined
+};
+
 const saveConnectionNameConfigName = 'notebook.saveConnectionName';
 const injectedParametersMsg = localize('injectedParametersMsg', '# Injected-Parameters\n');
 
@@ -93,6 +107,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _kernelAliases: string[] = [];
 	private _currentKernelAlias: string | undefined;
 	private _selectedKernelDisplayName: string | undefined;
+	private _multiConnectionMode: boolean = false;
 
 	public requestConnectionHandler: (() => Promise<boolean>) | undefined;
 
@@ -214,6 +229,14 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._savedConnectionName;
 	}
 
+	public get multiConnectionMode(): boolean {
+		return this._multiConnectionMode;
+	}
+
+	public set multiConnectionMode(isMultiConnection: boolean) {
+		this._multiConnectionMode = isMultiConnection;
+	}
+
 	public get specs(): nb.IAllKernels | undefined {
 		let specs: nb.IAllKernels = {
 			defaultKernel: '',
@@ -373,33 +396,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			// if cells already exist, create them with language info (if it is saved)
 			this._cells = [];
 			if (contents) {
-				this._defaultLanguageInfo = contents.metadata?.language_info;
-				// If language info was serialized in the notebook, attempt to use that to decrease time
-				// required until colorization occurs
-				if (this._defaultLanguageInfo) {
-					this.updateLanguageInfo(this._defaultLanguageInfo);
-				}
-				this._savedKernelInfo = this.getSavedKernelInfo(contents);
-				this._savedConnectionName = this.getSavedConnectionName(contents);
 				if (contents.metadata) {
-					//Telemetry of loading notebook
+					this.loadContentMetadata(contents.metadata);
-					let metadata: any = contents.metadata;
-					if (metadata.azdata_notebook_guid && metadata.azdata_notebook_guid.length === 36) {
-						//Verify if it is actual GUID and then send it to the telemetry
-						let regex = new RegExp('(\{){0,1}[0-9a-fA-F]{8}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{12}(\}){0,1}');
-						if (regex.test(metadata.azdata_notebook_guid)) {
-							this.adstelemetryService.createActionEvent(TelemetryKeys.TelemetryView.Notebook, TelemetryKeys.TelemetryAction.Open)
-								.withAdditionalProperties({ azdata_notebook_guid: metadata.azdata_notebook_guid })
-								.send();
-						}
-					}
-					Object.keys(contents.metadata).forEach(key => {
-						let expectedKeys = ['kernelspec', 'language_info', 'tags', 'connection_name'];
-						// If custom metadata is defined, add to the _existingMetadata object
-						if (expectedKeys.indexOf(key) < 0) {
-							this._existingMetadata[key] = contents.metadata[key];
-						}
-					});
 				}
 				// Modify Notebook URI Params format from URI query to string space delimited format
 				let notebookUriParams: string = this.notebookUri.query;
@@ -449,6 +447,34 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
+	private loadContentMetadata(metadata: INotebookMetadataInternal): void {
+		this._savedKernelInfo = metadata.kernelspec;
+		this._defaultLanguageInfo = metadata.language_info;
+		// If language info was serialized in the notebook, attempt to use that to decrease time
+		// required until colorization occurs
+		if (this._defaultLanguageInfo) {
+			this.updateLanguageInfo(this._defaultLanguageInfo);
+		}
+		this._savedConnectionName = metadata.connection_name;
+		this._multiConnectionMode = !!metadata.multi_connection_mode;
+
+		//Telemetry of loading notebook
+		if (metadata.azdata_notebook_guid && metadata.azdata_notebook_guid.length === 36) {
+			//Verify if it is actual GUID and then send it to the telemetry
+			if (isUUID(metadata.azdata_notebook_guid)) {
+				this.adstelemetryService.createActionEvent(TelemetryKeys.TelemetryView.Notebook, TelemetryKeys.TelemetryAction.Open)
+					.withAdditionalProperties({ azdata_notebook_guid: metadata.azdata_notebook_guid })
+					.send();
+			}
+		}
+		Object.keys(metadata).forEach(key => {
+			// If custom metadata is defined, add to the _existingMetadata object
+			if (!Object.keys(expectedMetadataKeys).includes(key)) {
+				this._existingMetadata[key] = metadata[key];
+			}
+		});
+	}
+
 	public async requestModelLoad(): Promise<void> {
 		try {
 			this.setDefaultKernelAndProviderId();
@@ -1009,16 +1035,6 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return values(connections).find(connection => connection.connectionName === connectionName);
 	}
 
-	// Get saved connection name if saved in notebook file
-	private getSavedConnectionName(notebook: nb.INotebookContents): string | undefined {
-		return notebook?.metadata?.connection_name ? notebook.metadata.connection_name : undefined;
-	}
-
-	// Get default kernel info if saved in notebook file
-	private getSavedKernelInfo(notebook: nb.INotebookContents): nb.IKernelSpec | undefined {
-		return (notebook && notebook.metadata && notebook.metadata.kernelspec) ? notebook.metadata.kernelspec : undefined;
-	}
-
 	private getKernelSpecFromDisplayName(displayName: string): nb.IKernelSpec | undefined {
 		let kernel: nb.IKernelSpec = this.specs.kernels.find(k => k.display_name.toLowerCase() === displayName.toLowerCase());
 		if (!kernel) {
@@ -1251,6 +1267,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		metadata.kernelspec = this._savedKernelInfo;
 		metadata.language_info = this.languageInfo;
 		metadata.tags = this._tags;
+		metadata.multi_connection_mode = this._multiConnectionMode ? this._multiConnectionMode : undefined;
 		if (this.configurationService.getValue(saveConnectionNameConfigName)) {
 			metadata.connection_name = this._savedConnectionName;
 		}