diff --git a/src/sql/workbench/contrib/notebook/browser/models/cell.ts b/src/sql/workbench/contrib/notebook/browser/models/cell.ts index 1650bbaf88..f3f421be5f 100644 --- a/src/sql/workbench/contrib/notebook/browser/models/cell.ts +++ b/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)); } } diff --git a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts index d40de963ee..87a3864ae6 100644 --- a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts +++ b/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; @@ -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 { - return this.textInput.save(groupId, options); + async save(groupId: number, options?: ITextFileSaveOptions): Promise { + let input = await this.textInput.save(groupId, options); + await this.setTrustForNewEditor(input); + return input; } - saveAs(group: number, options?: ITextFileSaveOptions): Promise { - return this.textInput.saveAs(group, options); + async saveAs(group: number, options?: ITextFileSaveOptions): Promise { + let input = await this.textInput.saveAs(group, options); + await this.setTrustForNewEditor(input); + return input; + } + + private async setTrustForNewEditor(newInput: IEditorInput | undefined): Promise { + 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[]) { diff --git a/src/sql/workbench/contrib/notebook/test/stubs.ts b/src/sql/workbench/contrib/notebook/test/stubs.ts index c7fb435b3a..9deae94de8 100644 --- a/src/sql/workbench/contrib/notebook/test/stubs.ts +++ b/src/sql/workbench/contrib/notebook/test/stubs.ts @@ -249,7 +249,7 @@ export class NotebookServiceStub implements INotebookService { isNotebookTrustCached(notebookUri: URI, isDirty: boolean): Promise { throw new Error('Method not implemented.'); } - serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel): void { + serializeNotebookStateChange(notebookUri: URI, changeType: NotebookChangeType, cell?: ICellModel, isTrusted?: boolean): Promise { throw new Error('Method not implemented.'); } navigateTo(notebookUri: URI, sectionId: string): void { diff --git a/src/sql/workbench/services/notebook/browser/notebookService.ts b/src/sql/workbench/services/notebook/browser/notebookService.ts index ec5facace9..b306d09a35 100644 --- a/src/sql/workbench/services/notebook/browser/notebookService.ts +++ b/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; /** * diff --git a/src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts b/src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts index 735c087377..290b3cc8c8 100644 --- a/src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts +++ b/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 { if (notebookUri.scheme !== Schemas.untitled) { // Conditions for saving: // 1. Not untitled. They're always trusted as we open them @@ -588,15 +585,23 @@ 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 - let notebook = find(this.listNotebookEditors(), n => n.id === notebookUriString); - if (notebook && notebook.model) { - if (notebook.model.trustedMode) { - this._trustedCacheQueue.push(notebookUri); - } else { - this._unTrustedCacheQueue.push(notebookUri); + 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) { + this._trustedCacheQueue.push(notebookUri); + } else { + this._unTrustedCacheQueue.push(notebookUri); + } + await this.updateTrustedCache(); } - this._updateTrustCacheScheduler.schedule(); } } }