Maintain notebook trust status after saving. (#9127)

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

@@ -19,7 +19,7 @@ import { INotificationService, Severity } from 'vs/platform/notification/common/
 import { Schemas } from 'vs/base/common/network';
 import { INotebookService } from 'sql/workbench/services/notebook/browser/notebookService';
 import { optional } from 'vs/platform/instantiation/common/instantiation';
-import { getErrorMessage } from 'vs/base/common/errors';
+import { getErrorMessage, onUnexpectedError } from 'vs/base/common/errors';
 import { generateUuid } from 'vs/base/common/uuid';
 import { IModelContentChangedEvent } from 'vs/editor/common/model/textModelEvents';
 import { firstIndex, find } from 'vs/base/common/arrays';
@@ -285,7 +285,8 @@ export class CellModel implements ICellModel {
 
 	private notifyExecutionComplete(): void {
 		if (this._notebookService) {
-			this._notebookService.serializeNotebookStateChange(this.notebookModel.notebookUri, NotebookChangeType.CellExecuted, this);
+			this._notebookService.serializeNotebookStateChange(this.notebookModel.notebookUri, NotebookChangeType.CellExecuted, this)
+				.catch(e => onUnexpectedError(e));
 		}
 	}
 

src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts

@@ -32,6 +32,7 @@ import { ResourceEditorInput } from 'vs/workbench/common/editor/resourceEditorIn
 import { FileEditorInput } from 'vs/workbench/contrib/files/common/editors/fileEditorInput';
 import { BinaryEditorModel } from 'vs/workbench/common/editor/binaryEditorModel';
 import { NotebookFindModel } from 'sql/workbench/contrib/notebook/find/notebookFindModel';
+import { onUnexpectedError } from 'vs/base/common/errors';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
 
@@ -160,7 +161,8 @@ export class NotebookEditorModel extends EditorModel {
 	private sendNotebookSerializationStateChange(): void {
 		let notebookModel = this.getNotebookModel();
 		if (notebookModel) {
-			this.notebookService.serializeNotebookStateChange(this.notebookUri, NotebookChangeType.Saved);
+			this.notebookService.serializeNotebookStateChange(this.notebookUri, NotebookChangeType.Saved)
+				.catch(e => onUnexpectedError(e));
 		}
 	}
 
@@ -283,12 +285,23 @@ export abstract class NotebookInput extends EditorInput {
 		return this._standardKernels;
 	}
 
-	save(groupId: number, options?: ITextFileSaveOptions): Promise<IEditorInput | undefined> {
+	async save(groupId: number, options?: ITextFileSaveOptions): Promise<IEditorInput | undefined> {
-		return this.textInput.save(groupId, options);
+		let input = await this.textInput.save(groupId, options);
+		await this.setTrustForNewEditor(input);
+		return input;
 	}
 
-	saveAs(group: number, options?: ITextFileSaveOptions): Promise<IEditorInput | undefined> {
+	async saveAs(group: number, options?: ITextFileSaveOptions): Promise<IEditorInput | undefined> {
-		return this.textInput.saveAs(group, options);
+		let input = await this.textInput.saveAs(group, options);
+		await this.setTrustForNewEditor(input);
+		return input;
+	}
+
+	private async setTrustForNewEditor(newInput: IEditorInput | undefined): Promise<void> {
+		let isTrusted = this._model.getNotebookModel().trustedMode;
+		if (isTrusted && newInput && newInput.getResource() !== this.getResource()) {
+			await this.notebookService.serializeNotebookStateChange(newInput.getResource(), NotebookChangeType.Saved, undefined, true);
+		}
 	}
 
 	public set standardKernels(value: IStandardKernelWithProvider[]) {

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

@@ -249,7 +249,7 @@ export class NotebookServiceStub implements INotebookService {
 	isNotebookTrustCached(notebookUri: URI, isDirty: boolean): Promise<boolean> {
 		throw new Error('Method not implemented.');
 	}
-	serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel): void {
+	serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel, isTrusted?: boolean): Promise<void> {
 		throw new Error('Method not implemented.');
 	}
 	navigateTo(notebookUri: URI, sectionId: string): void {

src/sql/workbench/services/notebook/browser/notebookService.ts

@@ -99,9 +99,12 @@ export interface INotebookService {
 	 * Serializes an impactful Notebook state change. This will result
 	 * in trusted state being serialized if needed, and notifications being
 	 * sent to listeners that can act on the point-in-time notebook state
-	 * @param notebookUri the URI identifying a notebook
+	 * @param notebookUri The URI identifying a notebook.
+	 * @param changeType The type of notebook state change to serialize.
+	 * @param cell (Optional) The notebook cell associated with the state change.
+	 * @param isTrusted (Optional) A manual override for the notebook's trusted state.
 	 */
-	serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel): void;
+	serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel, isTrusted?: boolean): Promise<void>;
 
 	/**
 	 *

src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts

@@ -36,7 +36,6 @@ import { SqlNotebookProvider } from 'sql/workbench/services/notebook/browser/sql
 import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { keys } from 'vs/base/common/map';
 import { IFileService, IFileStatWithMetadata } from 'vs/platform/files/common/files';
-import { RunOnceScheduler } from 'vs/base/common/async';
 import { Schemas } from 'vs/base/common/network';
 import { ILogService } from 'vs/platform/log/common/log';
 import { toErrorMessage } from 'vs/base/common/errorMessage';
@@ -116,7 +115,6 @@ export class NotebookService extends Disposable implements INotebookService {
 	private _overrideEditorThemeSetting: boolean;
 	private _trustedCacheQueue: URI[] = [];
 	private _unTrustedCacheQueue: URI[] = [];
-	private _updateTrustCacheScheduler: RunOnceScheduler;
 
 	constructor(
 		@ILifecycleService lifecycleService: ILifecycleService,
@@ -137,7 +135,6 @@ export class NotebookService extends Disposable implements INotebookService {
 		this._providersMemento = new Memento('notebookProviders', this._storageService);
 		this._trustedNotebooksMemento = new Memento('notebooks.trusted', this._storageService);
 
-		this._updateTrustCacheScheduler = new RunOnceScheduler(() => this.updateTrustedCache(), 250);
 		this._register(notebookRegistry.onNewRegistration(this.updateRegisteredProviders, this));
 		this.registerBuiltInProvider();
 		// If a provider has been already registered, the onNewRegistration event will not have a listener attached yet
@@ -579,7 +576,7 @@ export class NotebookService extends Disposable implements INotebookService {
 		}
 	}
 
-	serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel): void {
+	async serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel, isTrusted?: boolean): Promise<void> {
 		if (notebookUri.scheme !== Schemas.untitled) {
 			// Conditions for saving:
 			// 1. Not untitled. They're always trusted as we open them
@@ -588,7 +585,14 @@ export class NotebookService extends Disposable implements INotebookService {
 			// 4. Notebook is trusted. Don't need to save state of untrusted notebooks
 			let notebookUriString = notebookUri.toString();
 			if (changeType === NotebookChangeType.Saved && firstIndex(this._trustedCacheQueue, uri => uri.toString() === notebookUriString) < 0) {
-				// Only save if it's trusted
+				if (isTrusted) {
+					this._trustedCacheQueue.push(notebookUri);
+					await this.updateTrustedCache();
+				} else if (isTrusted === false) {
+					this._unTrustedCacheQueue.push(notebookUri);
+					await this.updateTrustedCache();
+				} else {
+					// Only save as trusted if the associated notebook model is trusted
 					let notebook = find(this.listNotebookEditors(), n => n.id === notebookUriString);
 					if (notebook && notebook.model) {
 						if (notebook.model.trustedMode) {
@@ -596,7 +600,8 @@ export class NotebookService extends Disposable implements INotebookService {
 						} else {
 							this._unTrustedCacheQueue.push(notebookUri);
 						}
-					this._updateTrustCacheScheduler.schedule();
+						await this.updateTrustedCache();
+					}
 				}
 			}
 		}