Fix #3439 Trusted doesn't get saved in Notebooks (#5507)

* Fix #3439 Trusted doesn't get saved in Notebooks
The main fix is to add a memento to store trust information. This is only needed for saved files - untitled files are always trusted as the user created them.
On clicking trusted or saving a file, the trusted state is cached. In the future, we will also handle code execution here too by sending notification on snapshot state.
I found issue #5506 during testing - existing issue where we should track trusted state changing on run. In the case all cells are ran, the whole notebook should become trusted.

Finally, I did a decent amount of refactoring to move more logic to the model - removing unnecessary calls from components which duplicated model behavior, moving trust notification to the model or at least the notebook service completely.

Added tests and logging for catch handling

src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts

@@ -399,7 +399,6 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		const fileInput = isUntitled ? this._untitledEditorService.createOrGet(uri, notebookModeId, options.initialContent) :
 			this._editorService.createInput({ resource: uri, language: notebookModeId });
 		let input = this._instantiationService.createInstance(NotebookInput, path.basename(uri.fsPath), uri, fileInput);
-		input.isTrusted = isUntitled;
 		input.defaultKernel = options.defaultKernel;
 		input.connectionProfile = new ConnectionProfile(this._capabilitiesService, options.connectionProfile);
 		if (isUntitled) {

src/sql/workbench/common/customInputConverter.ts

@@ -57,11 +57,7 @@ export function convertEditorInput(input: EditorInput, options: IQueryEditorOpti
 		//Notebook
 		uri = getNotebookEditorUri(input, instantiationService);
 		if (uri) {
-			let fileName: string = 'untitled';
-			if (input) {
-				fileName = input.getName();
-			}
-			let notebookInput: NotebookInput = instantiationService.createInstance(NotebookInput, fileName, uri, input);
+			let notebookInput: NotebookInput = instantiationService.createInstance(NotebookInput, input.getName(), uri, input);
 			return notebookInput;
 		}
 	}

src/sql/workbench/parts/notebook/models/contracts.ts

@@ -42,5 +42,6 @@ export enum NotebookChangeType {
 	CellSourceUpdated,
 	CellOutputUpdated,
 	DirtyStateChanged,
-	KernelChanged
+	KernelChanged,
+	TrustChanged
 }

src/sql/workbench/parts/notebook/models/notebookModel.ts

@@ -235,6 +235,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				c.trustedMode = this._trustedMode;
 			});
 		}
+		this._contentChangedEmitter.fire({
+			changeType: NotebookChangeType.TrustChanged
+		});
 	}
 
 	public get activeConnection(): IConnectionProfile {

src/sql/workbench/parts/notebook/notebook.component.ts

@@ -188,19 +188,6 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		this._model.addCell(cellType);
 	}
 
-	// Updates Notebook model's trust details
-	public updateModelTrustDetails(isTrusted: boolean) {
-		this._model.trustedMode = isTrusted;
-		this._model.cells.forEach(cell => {
-			cell.trustedMode = isTrusted;
-		});
-		//Updates dirty state
-		if (this._notebookParams.input) {
-			this._notebookParams.input.updateModel();
-		}
-		this.detectChanges();
-	}
-
 	public onKeyDown(event) {
 		switch (event.key) {
 			case 'ArrowDown':
@@ -294,7 +281,8 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			capabilitiesService: this.capabilitiesService
 		}, this.profile, this.logService);
 		model.onError((errInfo: INotification) => this.handleModelError(errInfo));
-		await model.requestModelLoad(this._notebookParams.isTrusted);
+		let trusted = await this.notebookService.isNotebookTrustCached(this._notebookParams.notebookUri, this.isDirty());
+		await model.requestModelLoad(trusted);
 		model.contentChanged((change) => this.handleContentChanged(change));
 		model.onProviderIdChange((provider) => this.handleProviderIdChanged(provider));
 		this._model = this._register(model);

src/sql/workbench/parts/notebook/notebookActions.ts

@@ -238,7 +238,7 @@ export class TrustedAction extends ToggleableAction {
 				}
 				else {
 					self.trusted = !self.trusted;
-					context.updateModelTrustDetails(self.trusted);
+					context.model.trustedMode = self.trusted;
 				}
 				resolve(true);
 			} catch (e) {

src/sql/workbench/parts/notebook/notebookEditor.ts

@@ -91,7 +91,6 @@ export class NotebookEditor extends BaseEditor {
 			notebookUri: input.notebookUri,
 			input: input,
 			providerInfo: input.getProviderInfo(),
-			isTrusted: input.isTrusted,
 			profile: input.connectionProfile
 		};
 		bootstrapAngular(this.instantiationService,

src/sql/workbench/parts/notebook/notebookInput.ts

@@ -11,20 +11,21 @@ import * as resources from 'vs/base/common/resources';
 import * as azdata from 'azdata';
 
 import { IStandardKernelWithProvider, getProvidersForFileName, getStandardKernelsForProvider } from 'sql/workbench/parts/notebook/notebookUtils';
-import { INotebookService, DEFAULT_NOTEBOOK_PROVIDER, IProviderInfo } from 'sql/workbench/services/notebook/common/notebookService';
+import { INotebookService, DEFAULT_NOTEBOOK_PROVIDER, IProviderInfo, SerializationStateChangeType } from 'sql/workbench/services/notebook/common/notebookService';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ITextModelService } from 'vs/editor/common/services/resolverService';
-import { INotebookModel, IContentManager } from 'sql/workbench/parts/notebook/models/modelInterfaces';
+import { INotebookModel, IContentManager, NotebookContentChange } from 'sql/workbench/parts/notebook/models/modelInterfaces';
 import { TextFileEditorModel } from 'vs/workbench/services/textfile/common/textFileEditorModel';
 import { Range } from 'vs/editor/common/core/range';
 import { UntitledEditorModel } from 'vs/workbench/common/editor/untitledEditorModel';
 import { Schemas } from 'vs/base/common/network';
-import { ITextFileService, ISaveOptions } from 'vs/workbench/services/textfile/common/textfiles';
+import { ITextFileService, ISaveOptions, StateChange } from 'vs/workbench/services/textfile/common/textfiles';
 import { LocalContentManager } from 'sql/workbench/services/notebook/node/localContentManager';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 import { IDisposable } from 'vs/base/common/lifecycle';
+import { NotebookChangeType } from 'sql/workbench/parts/notebook/models/contracts';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
 
@@ -42,7 +43,7 @@ export class NotebookEditorModel extends EditorModel {
 				// Hook to content change events
 				notebook.modelReady.then(() => {
 					this._register(notebook.model.kernelChanged(e => this.updateModel()));
-					this._register(notebook.model.contentChanged(e => this.updateModel()));
+					this._register(notebook.model.contentChanged(e => this.updateModel(e)));
 				}, err => undefined);
 			}
 		}));
@@ -50,7 +51,12 @@ export class NotebookEditorModel extends EditorModel {
 		if (this.textEditorModel instanceof UntitledEditorModel) {
 			this._register(this.textEditorModel.onDidChangeDirty(e => this.setDirty(this.textEditorModel.isDirty())));
 		} else {
-			this._register(this.textEditorModel.onDidStateChange(e => this.setDirty(this.textEditorModel.isDirty())));
+			this._register(this.textEditorModel.onDidStateChange(change => {
+				this.setDirty(this.textEditorModel.isDirty());
+				if (change === StateChange.SAVED) {
+					this.sendNotebookSerializationStateChange();
+				}
+			}));
 		}
 		this.dirty = this.textEditorModel.isDirty();
 	}
@@ -86,18 +92,33 @@ export class NotebookEditorModel extends EditorModel {
 		}
 	}
 
-	public updateModel(): void {
-		let notebookModel = this.getNotebookModel();
-		if (notebookModel && this.textEditorModel && this.textEditorModel.textEditorModel) {
-			let content = JSON.stringify(notebookModel.toJSON(), undefined, '    ');
-			let model = this.textEditorModel.textEditorModel;
-			let endLine = model.getLineCount();
-			let endCol = model.getLineMaxColumn(endLine);
+	public updateModel(contentChange?: NotebookContentChange): void {
+		if (contentChange && contentChange.changeType === NotebookChangeType.TrustChanged) {
+			// This is a serializable change (in that we permanently cache trusted state, but
+			// ironically isn't cached in the JSON contents since trust doesn't persist across machines.
+			// Request serialization so trusted state is preserved but don't update the model
+			this.sendNotebookSerializationStateChange();
+		} else {
+			// For all other changes, update the backing model with the latest contents
+			let notebookModel = this.getNotebookModel();
+			if (notebookModel && this.textEditorModel && this.textEditorModel.textEditorModel) {
+				let content = JSON.stringify(notebookModel.toJSON(), undefined, '    ');
+				let model = this.textEditorModel.textEditorModel;
+				let endLine = model.getLineCount();
+				let endCol = model.getLineMaxColumn(endLine);
 
-			this.textEditorModel.textEditorModel.applyEdits([{
-				range: new Range(1, 1, endLine, endCol),
-				text: content
-			}]);
+				this.textEditorModel.textEditorModel.applyEdits([{
+					range: new Range(1, 1, endLine, endCol),
+					text: content
+				}]);
+			}
+		}
+	}
+
+	private sendNotebookSerializationStateChange() {
+		let notebookModel = this.getNotebookModel();
+		if (notebookModel) {
+			this.notebookService.serializeNotebookStateChange(this.notebookUri, SerializationStateChangeType.Saved);
 		}
 	}
 
@@ -125,7 +146,6 @@ export class NotebookInput extends EditorInput {
 	private _standardKernels: IStandardKernelWithProvider[];
 	private _connectionProfile: IConnectionProfile;
 	private _defaultKernel: azdata.nb.IKernelSpec;
-	private _isTrusted: boolean = false;
 	public hasBootstrapped = false;
 	// Holds the HTML content for the editor when the editor discards this input and loads another
 	private _parentContainer: HTMLElement;
@@ -190,13 +210,6 @@ export class NotebookInput extends EditorInput {
 			providers: this._providers ? this._providers : [DEFAULT_NOTEBOOK_PROVIDER]
 		};
 	}
-	public get isTrusted(): boolean {
-		return this._isTrusted;
-	}
-
-	public set isTrusted(value: boolean) {
-		this._isTrusted = value;
-	}
 
 	public set connectionProfile(value: IConnectionProfile) {
 		this._connectionProfile = value;

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

@@ -68,6 +68,25 @@ export interface INotebookService {
 	getMimeRegistry(): RenderMimeRegistry;
 
 	renameNotebookEditor(oldUri: URI, newUri: URI, currentEditor: INotebookEditor): void;
+
+	/**
+	 * Checks if a notebook has previously been marked as trusted, and that
+	 * the notebook has not changed on disk since that time. If the notebook
+	 * is currently dirty in the app, the previous trusted state will be used even
+	 * if it's altered on disk since the version in our UI is based on previously trusted
+	 * content.
+	 * @param notebookUri the URI identifying a notebook
+	 * @param isDirty is the notebook marked as dirty in by the text model trackers?
+	 */
+	isNotebookTrustCached(notebookUri: URI, isDirty: boolean): Promise<boolean>;
+	/**
+	 * 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
+	 */
+	serializeNotebookStateChange(notebookUri: URI, changeType: SerializationStateChangeType): void;
+
 }
 
 export interface INotebookProvider {
@@ -91,7 +110,6 @@ export interface INotebookParams extends IBootstrapParams {
 	notebookUri: URI;
 	input: NotebookInput;
 	providerInfo: Promise<IProviderInfo>;
-	isTrusted: boolean;
 	profile?: IConnectionProfile;
 	modelFactory?: ModelFactory;
 }
@@ -110,3 +128,8 @@ export interface INotebookEditor {
 	runAllCells(): Promise<boolean>;
 	clearAllOutputs(): Promise<boolean>;
 }
+
+export enum SerializationStateChangeType {
+	Saved,
+	Executed
+}

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

@@ -10,7 +10,7 @@ import { Registry } from 'vs/platform/registry/common/platform';
 
 import {
 	INotebookService, INotebookManager, INotebookProvider,
-	DEFAULT_NOTEBOOK_FILETYPE, INotebookEditor, SQL_NOTEBOOK_PROVIDER, OVERRIDE_EDITOR_THEMING_SETTING
+	DEFAULT_NOTEBOOK_FILETYPE, INotebookEditor, SQL_NOTEBOOK_PROVIDER, OVERRIDE_EDITOR_THEMING_SETTING, SerializationStateChangeType
 } from 'sql/workbench/services/notebook/common/notebookService';
 import { RenderMimeRegistry } from 'sql/workbench/parts/notebook/outputs/registry';
 import { standardRendererFactories } from 'sql/workbench/parts/notebook/outputs/factories';
@@ -35,6 +35,11 @@ import { ILifecycleService } from 'vs/platform/lifecycle/common/lifecycle';
 import { SqlNotebookProvider } from 'sql/workbench/services/notebook/sql/sqlNotebookProvider';
 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';
 
 export interface NotebookProviderProperties {
 	provider: string;
@@ -49,6 +54,18 @@ interface NotebookProvidersMemento {
 	notebookProviderCache: NotebookProviderCache;
 }
 
+interface TrustedNotebookMetadata {
+	mtime: number;
+}
+interface TrustedNotebookCache {
+	// URI goes to cached
+	[uri: string]: TrustedNotebookMetadata;
+}
+
+interface TrustedNotebooksMemento {
+	trustedNotebooksCache: TrustedNotebookCache;
+}
+
 const notebookRegistry = Registry.as<INotebookProviderRegistry>(Extensions.NotebookProviderContribution);
 
 class ProviderDescriptor {
@@ -75,7 +92,8 @@ class ProviderDescriptor {
 export class NotebookService extends Disposable implements INotebookService {
 	_serviceBrand: any;
 
-	private _memento: Memento;
+	private _providersMemento: Memento;
+	private _trustedNotebooksMemento: Memento;
 	private _mimeRegistry: RenderMimeRegistry;
 	private _providers: Map<string, ProviderDescriptor> = new Map();
 	private _managersMap: Map<string, INotebookManager[]> = new Map();
@@ -91,6 +109,8 @@ export class NotebookService extends Disposable implements INotebookService {
 	private notebookEditorVisible: IContextKey<boolean>;
 	private _themeParticipant: IDisposable;
 	private _overrideEditorThemeSetting: boolean;
+	private _trustedCacheQueue: URI[] = [];
+	private _updateTrustCacheScheduler: RunOnceScheduler;
 
 	constructor(
 		@ILifecycleService lifecycleService: ILifecycleService,
@@ -102,10 +122,15 @@ export class NotebookService extends Disposable implements INotebookService {
 		@IEditorService private readonly _editorService: IEditorService,
 		@IEditorGroupsService private readonly _editorGroupsService: IEditorGroupsService,
 		@IConfigurationService private readonly _configurationService: IConfigurationService,
+		@IFileService private readonly _fileService: IFileService,
+		@ILogService private readonly _logService: ILogService,
 		@IQueryManagementService private readonly _queryManagementService
 	) {
 		super();
-		this._memento = new Memento('notebookProviders', this._storageService);
+		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
@@ -448,7 +473,15 @@ export class NotebookService extends Disposable implements INotebookService {
 	}
 
 	private get providersMemento(): NotebookProvidersMemento {
-		return this._memento.getMemento(StorageScope.GLOBAL) as NotebookProvidersMemento;
+		return this._providersMemento.getMemento(StorageScope.GLOBAL) as NotebookProvidersMemento;
+	}
+
+	private get trustedNotebooksMemento(): TrustedNotebooksMemento {
+		let cache = this._trustedNotebooksMemento.getMemento(StorageScope.GLOBAL) as TrustedNotebooksMemento;
+		if (!cache.trustedNotebooksCache) {
+			cache.trustedNotebooksCache = {};
+		}
+		return cache;
 	}
 
 	private cleanupProviders(): void {
@@ -484,4 +517,85 @@ export class NotebookService extends Disposable implements INotebookService {
 			}
 		});
 	}
+
+	async isNotebookTrustCached(notebookUri: URI, isDirty: boolean): Promise<boolean> {
+		if (notebookUri.scheme === Schemas.untitled) {
+			return true;
+		}
+
+		let cacheInfo = this.trustedNotebooksMemento.trustedNotebooksCache[notebookUri.toString()];
+		if (!cacheInfo) {
+			// This notebook was never trusted
+			return false;
+		}
+		// This was trusted. If it's not dirty (e.g. if we're not working on our cached copy)
+		// then should verify it's not been modified on disk since that invalidates trust relationship
+		if (!isDirty) {
+			// Check mtime against mtime on disk
+			let actualMtime: number = await this.getModifiedTimeForFile(notebookUri);
+			if (actualMtime > cacheInfo.mtime) {
+				// Modified since last use, so can't guarantee trust.
+				return false;
+			}
+		}
+		return true;
+	}
+
+	private async getModifiedTimeForFile(notebookUri: URI): Promise<number> {
+		try {
+			let fstat: IFileStatWithMetadata = await this._fileService.resolve(notebookUri, {
+				resolveMetadata: true
+			});
+			return fstat ? fstat.mtime : 0;
+		} catch (err) {
+			return 0;
+		}
+	}
+
+	serializeNotebookStateChange(notebookUri: URI, changeType: SerializationStateChangeType): void {
+		let updateTrustState = changeType === SerializationStateChangeType.Saved;
+
+		if (notebookUri.scheme !== Schemas.untitled) {
+			// Cache state for non-untitled notebooks only.
+			let notebookUriString = notebookUri.toString();
+			if (updateTrustState && this._trustedCacheQueue.findIndex(uri => uri.toString() === notebookUriString)) {
+				this._trustedCacheQueue.push(notebookUri);
+				this._updateTrustCacheScheduler.schedule();
+			}
+
+			// TODO add history notification if a non-untitled notebook has a state change
+		}
+	}
+
+	private async updateTrustedCache(): Promise<void> {
+		try {
+			if (this._trustedCacheQueue.length > 0) {
+				// Copy out all items from the cache
+				let items = this._trustedCacheQueue;
+				this._trustedCacheQueue = [];
+
+				// Get all the file stats and then serialize this to a memento
+				let itemConfig = items.map(item => {
+					return { resource: item, options: { resolveMetadata: true } };
+				});
+				let metadata = await this._fileService.resolveAll(itemConfig);
+				let trustedCache = this.trustedNotebooksMemento.trustedNotebooksCache;
+				for (let i = 0; i < metadata.length; i++) {
+					let item = items[i];
+					let stat = metadata[i] && metadata[i].stat;
+					if (stat && stat.mtime) {
+						trustedCache[item.toString()] = {
+							mtime: stat.mtime
+						};
+					}
+				}
+
+				this._trustedNotebooksMemento.saveMemento();
+			}
+		} catch (err) {
+			if (this._logService) {
+				this._logService.trace(`Failed to save trust state to cache: ${toErrorMessage(err)}`);
+			}
+		}
+	}
 }