Notebooks: Save connection information in metadata (#13060)

* save connection info in notebook metadata

* update attachTo dropdown based on saved alias

* add setting for saving connection (default=false)

* dont show saved conn if seting off + added test

* show conn dialog if save conn name setting off

* address PR comments and fix unit test

* change connectionName to connection_name

src/sql/azdata.proposed.d.ts

@@ -77,6 +77,10 @@ declare module 'azdata' {
 		export interface ICellOutputMetadata {
 			resultSet?: ResultSetSummary;
 		}
+
+		export interface INotebookMetadata {
+			connection_name?: string;
+		}
 	}
 
 	export type SqlDbType = 'BigInt' | 'Binary' | 'Bit' | 'Char' | 'DateTime' | 'Decimal'

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

@@ -331,7 +331,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			layoutChanged: this._notebookParams.input.layoutChanged,
 			capabilitiesService: this.capabilitiesService,
 			editorLoadedTimestamp: this._notebookParams.input.editorOpenedTimestamp
-		}, this.profile, this.logService, this.notificationService, this.adstelemetryService, this.capabilitiesService);
+		}, this.profile, this.logService, this.notificationService, this.adstelemetryService, this.connectionManagementService, this._configurationService, this.capabilitiesService);
 		let trusted = await this.notebookService.isNotebookTrustCached(this._notebookParams.notebookUri, this.isDirty());
 		this._register(model.onError((errInfo: INotification) => this.handleModelError(errInfo)));
 		this._register(model.contentChanged((change) => this.handleContentChanged(change)));
@@ -429,7 +429,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 			let attachToContainer = document.createElement('li');
 			let attachToDropdown = new AttachToDropdown(attachToContainer, this.contextViewService, this.modelReady,
-				this.connectionManagementService, this.connectionDialogService, this.notificationService, this.capabilitiesService);
+				this.connectionManagementService, this.connectionDialogService, this.notificationService, this.capabilitiesService, this._configurationService);
 			attachToDropdown.render(attachToContainer);
 			attachSelectBoxStyler(attachToDropdown, this.themeService);
 
@@ -493,7 +493,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 			let attachToContainer = document.createElement('div');
 			let attachToDropdown = new AttachToDropdown(attachToContainer, this.contextViewService, this.modelReady,
-				this.connectionManagementService, this.connectionDialogService, this.notificationService, this.capabilitiesService);
+				this.connectionManagementService, this.connectionDialogService, this.notificationService, this.capabilitiesService, this._configurationService);
 			attachToDropdown.render(attachToContainer);
 			attachSelectBoxStyler(attachToDropdown, this.themeService);
 

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

@@ -230,7 +230,12 @@ configurationRegistry.registerConfiguration({
 			'type': 'boolean',
 			'default': true,
 			'description': localize('notebook.setRichTextViewByDefault', "Set Rich Text View mode by default for text cells")
-		}
+		},
+		'notebook.saveConnectionName': {
+			'type': 'boolean',
+			'default': false,
+			'description': localize('notebook.saveConnectionName', "(Preview) Save connection name in notebook metadata.")
+		},
 	}
 });
 

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

@@ -355,6 +355,7 @@ export class KernelsDropdown extends SelectBox {
 }
 
 const attachToDropdownElementId = 'attach-to-dropdown';
+const saveConnectionNameConfigName = 'notebook.saveConnectionName';
 
 export class AttachToDropdown extends SelectBox {
 	private model: NotebookModel;
@@ -365,6 +366,7 @@ export class AttachToDropdown extends SelectBox {
 		@IConnectionDialogService private _connectionDialogService: IConnectionDialogService,
 		@INotificationService private _notificationService: INotificationService,
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
+		@IConfigurationService private _configurationService: IConfigurationService
 	) {
 		super([msgLoadingContexts], msgLoadingContexts, contextViewProvider, container, { labelText: attachToLabel, labelOnTop: false, ariaLabel: attachToLabel, id: attachToDropdownElementId } as ISelectBoxOptionsWithLabel);
 		if (modelReady) {
@@ -428,7 +430,14 @@ export class AttachToDropdown extends SelectBox {
 		if ((connProviderIds && connProviderIds.length === 0) || currentKernel === noKernel) {
 			this.setOptions([msgLocalHost]);
 		} else {
-			let connections: string[] = model.context && model.context.title && (connProviderIds.includes(this.model.context.providerName)) ? [model.context.title] : [msgSelectConnection];
+			let connections: string[] = [];
+			if (model.context && model.context.title && (connProviderIds.includes(this.model.context.providerName))) {
+				connections.push(model.context.title);
+			} else if (this._configurationService.getValue(saveConnectionNameConfigName) && model.savedConnectionName) {
+				connections.push(model.savedConnectionName);
+			} else {
+				connections.push(msgSelectConnection);
+			}
 			if (!connections.find(x => x === msgChangeConnection)) {
 				connections.push(msgChangeConnection);
 			}

src/sql/workbench/contrib/notebook/test/browser/cellToolbarActions.test.ts

@@ -213,5 +213,5 @@ export async function createandLoadNotebookModel(codeContent?: nb.INotebookConte
 		layoutChanged: undefined,
 		capabilitiesService: undefined
 	};
-	return new NotebookModel(defaultModelOptions, undefined, undefined, undefined, undefined);
+	return new NotebookModel(defaultModelOptions, undefined, undefined, undefined, undefined, undefined, undefined);
 }

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

@@ -971,7 +971,7 @@ suite('Notebook Editor Model', function (): void {
 		let options: INotebookModelOptions = assign({}, defaultModelOptions, <Partial<INotebookModelOptions>><unknown>{
 			factory: mockModelFactory.object
 		});
-		notebookModel = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService());
+		notebookModel = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await notebookModel.loadContents();
 	}
 

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

@@ -30,6 +30,8 @@ import { NotebookEditorContentManager } from 'sql/workbench/contrib/notebook/bro
 import { NotebookRange } from 'sql/workbench/services/notebook/browser/notebookService';
 import { NotebookMarkdownRenderer } from 'sql/workbench/contrib/notebook/browser/outputs/notebookMarkdown';
 import { NullAdsTelemetryService } from 'sql/platform/telemetry/common/adsTelemetryService';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { TestConfigurationService } from 'vs/platform/configuration/test/common/testConfigurationService';
 
 let expectedNotebookContent: nb.INotebookContents = {
 	cells: [{
@@ -73,6 +75,7 @@ suite('Notebook Find Model', function (): void {
 	const logService = new NullLogService();
 	let model: NotebookModel;
 	let markdownRenderer: NotebookMarkdownRenderer = new NotebookMarkdownRenderer();
+	let configurationService: IConfigurationService;
 
 	setup(async () => {
 		sessionReady = new Deferred<void>();
@@ -84,6 +87,7 @@ suite('Notebook Find Model', function (): void {
 		queryConnectionService.callBase = true;
 
 		instantiationService = new InstantiationService(serviceCollection, true);
+		configurationService = new TestConfigurationService();
 		defaultModelOptions = {
 			notebookUri: defaultUri,
 			factory: new ModelFactory(instantiationService),
@@ -433,7 +437,7 @@ suite('Notebook Find Model', function (): void {
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(contents));
 		defaultModelOptions.contentManager = mockContentManager.object;
 		// Initialize the model
-		model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents();
 		await model.requestModelLoad();
 	}

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

@@ -36,6 +36,9 @@ import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { uriPrefixes } from 'sql/platform/connection/common/utils';
 import { NullAdsTelemetryService } from 'sql/platform/telemetry/common/adsTelemetryService';
+import { TestConfigurationService } from 'sql/platform/connection/test/common/testConfigurationService';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 
 let expectedNotebookContent: nb.INotebookContents = {
 	cells: [{
@@ -87,6 +90,7 @@ let mockModelFactory: TypeMoq.Mock<ModelFactory>;
 let notificationService: TypeMoq.Mock<INotificationService>;
 let capabilitiesService: ICapabilitiesService;
 let instantiationService: IInstantiationService;
+let configurationService: IConfigurationService;
 
 suite('notebook model', function (): void {
 	let notebookManagers = [new NotebookManagerStub()];
@@ -107,6 +111,7 @@ suite('notebook model', function (): void {
 		queryConnectionService.callBase = true;
 		let serviceCollection = new ServiceCollection();
 		instantiationService = new InstantiationService(serviceCollection, true);
+		configurationService = new TestConfigurationService();
 		defaultModelOptions = {
 			notebookUri: defaultUri,
 			factory: new ModelFactory(instantiationService),
@@ -154,7 +159,7 @@ suite('notebook model', function (): void {
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(emptyNotebook));
 		defaultModelOptions.contentManager = mockContentManager.object;
 		// When I initialize the model
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents();
 
 		// Then I expect to have 0 code cell as the contents
@@ -170,7 +175,7 @@ suite('notebook model', function (): void {
 		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());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents(true);
 		await model.requestModelLoad();
 
@@ -187,7 +192,7 @@ suite('notebook model', function (): void {
 
 		// When I initalize the model
 		// Then it should throw
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		assert.equal(model.inErrorState, false);
 		await assert.rejects(async () => { await model.loadContents(); });
 		assert.equal(model.inErrorState, true);
@@ -200,7 +205,7 @@ suite('notebook model', function (): void {
 		defaultModelOptions.contentManager = mockContentManager.object;
 
 		// When I initalize the model
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents();
 
 		// Then I expect all cells to be in the model
@@ -228,7 +233,7 @@ suite('notebook model', function (): void {
 		defaultModelOptions.providerId = 'jupyter';
 
 		// When I initalize the model
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents();
 
 		// I expect the default provider to be jupyter
@@ -238,7 +243,7 @@ suite('notebook model', function (): void {
 		defaultModelOptions.providerId = 'SQL';
 
 		// When I initalize the model
-		model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents();
 
 		// I expect the default provider to be SQL
@@ -263,7 +268,7 @@ suite('notebook model', function (): void {
 		defaultModelOptions.contentManager = mockContentManager.object;
 
 		// When I initalize the model
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents();
 
 		let activeCellChangeCount = 0;
@@ -320,7 +325,7 @@ suite('notebook model', function (): void {
 		defaultModelOptions.contentManager = mockContentManager.object;
 
 		// When I initalize the model
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.loadContents();
 
 		// Count number of times onError event is fired
@@ -375,7 +380,7 @@ suite('notebook model', function (): void {
 		sessionReady.resolve();
 		let sessionFired = false;
 
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		model.onClientSessionReady((session) => sessionFired = true);
 		await model.loadContents();
 		await model.requestModelLoad();
@@ -405,7 +410,7 @@ suite('notebook model', function (): void {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
 		defaultModelOptions.contentManager = mockContentManager.object;
-		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService());
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
 		await model.requestModelLoad();
 
 		let actualChanged: NotebookContentChange;
@@ -474,7 +479,7 @@ suite('notebook model', function (): void {
 		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, undefined);
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, undefined, queryConnectionService.object, configurationService);
 		await model.loadContents();
 
 		let output = model.toJSON();
@@ -570,6 +575,57 @@ suite('notebook model', function (): void {
 		queryConnectionService.verify((c) => c.disconnect(TypeMoq.It.isAny()), TypeMoq.Times.once());
 	});
 
+	test('Should read connection name from notebook metadata and use its corresponding connection profile', async function () {
+		const connectionName = 'connectionName';
+		// Given a notebook with a connection name in metadata
+		let notebook: nb.INotebookContents = {
+			cells: [],
+			metadata: {
+				connection_name: connectionName
+			},
+			nbformat: 4,
+			nbformat_minor: 5
+		};
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(notebook));
+		defaultModelOptions.contentManager = mockContentManager.object;
+
+		// And a matching connection profile
+		let expectedConnectionProfile: IConnectionProfile = {
+			connectionName: connectionName,
+			serverName: '',
+			databaseName: '',
+			userName: '',
+			password: '',
+			authenticationType: '',
+			savePassword: true,
+			groupFullName: '',
+			groupId: '',
+			getOptionsKey: () => '',
+			matches: undefined,
+			providerName: '',
+			saveProfile: true,
+			id: '',
+			options: {}
+		};
+		sinon.stub(queryConnectionService.object, 'getConnections').returns([expectedConnectionProfile]);
+		sinon.stub(configurationService, 'getValue').returns(true);
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService);
+		await model.loadContents();
+
+		// I expect the saved connection name to be read
+		assert.equal(model.savedConnectionName, connectionName);
+
+		// When I request a connection
+		let spy = sinon.stub(model, 'changeContext').returns(Promise.resolve());
+		model.requestConnection();
+
+		// I expect the connection profile matching the saved connection name to be used
+		assert.ok(spy.calledWith(connectionName, expectedConnectionProfile));
+	});
+
 	async function loadModelAndStartClientSession(): Promise<NotebookModel> {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));
@@ -583,7 +639,7 @@ suite('notebook model', function (): void {
 		let options: INotebookModelOptions = assign({}, defaultModelOptions, <Partial<INotebookModelOptions>>{
 			factory: mockModelFactory.object
 		});
-		let model = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService(), capabilitiesService);
+		let model = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, capabilitiesService);
 		model.onClientSessionReady((session) => actualSession = session);
 		await model.requestModelLoad();
 

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

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

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

@@ -297,6 +297,11 @@ export interface INotebookModel {
 	 */
 	readonly context: ConnectionProfile | undefined;
 
+	/**
+	 * The connection name (alias) saved in the notebook metadata,
+	 * or undefined if none.
+	 */
+	readonly savedConnectionName: string | undefined;
 	/**
 	 * Event fired on first initialization of the cells and
 	 * on subsequent change events

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

@@ -29,6 +29,9 @@ import { notebookConstants } from 'sql/workbench/services/notebook/browser/inter
 import { IAdsTelemetryService } from 'sql/platform/telemetry/common/telemetry';
 import { Deferred } from 'sql/base/common/promise';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
+import { values } from 'vs/base/common/collections';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 
 /*
 * Used to control whether a message in a dialog/wizard is displayed as an error,
@@ -45,6 +48,8 @@ export class ErrorInfo {
 	}
 }
 
+const saveConnectionNameConfigName = 'notebook.saveConnectionName';
+
 export class NotebookModel extends Disposable implements INotebookModel {
 	private _contextsChangedEmitter = new Emitter<void>();
 	private _contextsLoadingEmitter = new Emitter<void>();
@@ -68,6 +73,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _language: string = '';
 	private _onErrorEmitter = new Emitter<INotification>();
 	private _savedKernelInfo: nb.IKernelSpec | undefined;
+	private _savedConnectionName: string | undefined;
 	private readonly _nbformat: number = nbversion.MAJOR_VERSION;
 	private readonly _nbformatMinor: number = nbversion.MINOR_VERSION;
 	private _activeConnection: ConnectionProfile | undefined;
@@ -93,6 +99,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		@ILogService private readonly logService: ILogService,
 		@INotificationService private readonly notificationService: INotificationService,
 		@IAdsTelemetryService private readonly adstelemetryService: IAdsTelemetryService,
+		@IConnectionManagementService private connectionManagementService: IConnectionManagementService,
+		@IConfigurationService private configurationService: IConfigurationService,
 		@ICapabilitiesService private _capabilitiesService?: ICapabilitiesService
 
 	) {
@@ -199,6 +207,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._activeConnection;
 	}
 
+	public get savedConnectionName(): string | undefined {
+		return this._savedConnectionName;
+	}
+
 	public get specs(): nb.IAllKernels | undefined {
 		let specs: nb.IAllKernels = {
 			defaultKernel: '',
@@ -325,6 +337,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			if (contents) {
 				this._defaultLanguageInfo = contents.metadata?.language_info;
 				this._savedKernelInfo = this.getSavedKernelInfo(contents);
+				this._savedConnectionName = this.getSavedConnectionName(contents);
 				if (contents.metadata) {
 					//Telemetry of loading notebook
 					let metadata: any = contents.metadata;
@@ -338,7 +351,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 						}
 					}
 					Object.keys(contents.metadata).forEach(key => {
-						let expectedKeys = ['kernelspec', 'language_info', 'tags'];
+						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];
@@ -387,6 +400,15 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public async requestConnection(): Promise<boolean> {
+		// If there is a saved connection name with a corresponding connection profile, use that one,
+		// otherwise show connection dialog
+		if (this.configurationService.getValue(saveConnectionNameConfigName) && this._savedConnectionName) {
+			let profile: ConnectionProfile | undefined = this.getConnectionProfileFromName(this._savedConnectionName);
+			if (profile) {
+				await this.changeContext(this._savedConnectionName, profile);
+				return true;
+			}
+		}
 		if (this.requestConnectionHandler) {
 			return this.requestConnectionHandler();
 		} else if (this.notificationService) {
@@ -568,6 +590,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 			if (this.isValidConnection(profile)) {
 				this._activeConnection = profile;
+				this._savedConnectionName = profile.connectionName;
 			}
 
 			clientSession.onKernelChanging(async (e) => {
@@ -856,6 +879,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 					this._kernelDisplayNameToConnectionProviderIds.set(newConnection.serverCapabilities.notebookKernelAlias, [newConnection.providerName]);
 				}
 				this._activeConnection = newConnection;
+				this._savedConnectionName = newConnection.connectionName;
 				this.setActiveConnectionIfDifferent(newConnection);
 				this._activeClientSession.updateConnection(newConnection.toIConnectionProfile()).then(
 					result => {
@@ -887,9 +911,19 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			this._activeConnection.id !== newConnection.id) {
 			// Change the active connection to newConnection
 			this._activeConnection = newConnection;
+			this._savedConnectionName = newConnection.connectionName;
 		}
 	}
 
+	private getConnectionProfileFromName(connectionName: string): ConnectionProfile | undefined {
+		let connections: ConnectionProfile[] = this.connectionManagementService.getConnections();
+		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 {
@@ -1128,6 +1162,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		metadata.kernelspec = this._savedKernelInfo;
 		metadata.language_info = this.languageInfo;
 		metadata.tags = this._tags;
+		if (this.configurationService.getValue(saveConnectionNameConfigName)) {
+			metadata.connection_name = this._savedConnectionName;
+		}
 		Object.keys(this._existingMetadata).forEach(key => {
 			metadata[key] = this._existingMetadata[key];
 		});