Fix backup storage by looking at the options correctly

news/2 Fixes/13981.md

@@ -0,0 +1 @@
+Backup on custom editors is being ignored.

src/client/datascience/export/exportUtil.ts

@@ -57,7 +57,7 @@ export class ExportUtil {
             await this.jupyterExporter.exportToFile(cells, tempFile.filePath, false);
             const newPath = path.join(tempDir.path, '.ipynb');
             await this.fs.copyLocal(tempFile.filePath, newPath);
-            model = await this.notebookStorage.getOrCreateModel(Uri.file(newPath));
+            model = await this.notebookStorage.getOrCreateModel({ file: Uri.file(newPath) });
         } finally {
             tempFile.dispose();
             tempDir.dispose();
@@ -67,7 +67,7 @@ export class ExportUtil {
     }
 
     public async removeSvgs(source: Uri) {
-        const model = await this.notebookStorage.getOrCreateModel(source);
+        const model = await this.notebookStorage.getOrCreateModel({ file: source });
 
         const newCells: ICell[] = [];
         for (const cell of model.cells) {

src/client/datascience/interactive-ipynb/nativeEditorProviderOld.ts

@@ -299,7 +299,7 @@ export class NativeEditorProviderOld extends NativeEditorProvider {
         this.activeEditors.set(e.file.fsPath, e);
 
         // Remove backup storage
-        this.loadModel(Uri.file(oldPath))
+        this.loadModel({ file: Uri.file(oldPath) })
             .then((m) => this.storage.deleteBackup(m))
             .ignoreErrors();
     }
@@ -347,7 +347,7 @@ export class NativeEditorProviderOld extends NativeEditorProvider {
      * I.e. document is already opened in a VSC Notebook.
      */
     private async isDocumentOpenedInVSCodeNotebook(document: TextDocument): Promise<boolean> {
-        const model = await this.loadModel(document.uri);
+        const model = await this.loadModel({ file: document.uri });
         // This is temporary code.
         return model instanceof VSCodeNotebookModel;
     }

src/client/datascience/interactive-ipynb/trustCommandHandler.ts

@@ -42,7 +42,7 @@ export class TrustCommandHandler implements IExtensionSingleActivationService {
             return;
         }
 
-        const model = await this.storageProvider.getOrCreateModel(uri);
+        const model = await this.storageProvider.getOrCreateModel({ file: uri });
         if (model.isTrusted) {
             return;
         }

src/client/datascience/interactive-window/interactiveWindow.ts

@@ -256,7 +256,7 @@ export class InteractiveWindow extends InteractiveBase implements IInteractiveWi
                 break;
 
             case InteractiveWindowMessages.ExportNotebookAs:
-                this.handleMessage(message, payload, this.exportAs);
+                this.handleMessage(message, payload, this.exportAs.bind);
                 break;
 
             case InteractiveWindowMessages.HasCellResponse:

src/client/datascience/notebook/contentProvider.ts

@@ -70,9 +70,12 @@ export class NotebookContentProvider implements INotebookContentProvider {
             };
         }
         // If there's no backup id, then skip loading dirty contents.
-        const model = await (openContext.backupId
-            ? this.notebookStorage.getOrCreateModel(uri, undefined, openContext.backupId, true)
-            : this.notebookStorage.getOrCreateModel(uri, undefined, true, true));
+        const model = await this.notebookStorage.getOrCreateModel({
+            file: uri,
+            backupId: openContext.backupId,
+            isNative: true,
+            skipLoadingDirtyContents: openContext.backupId === undefined
+        });
         if (!(model instanceof VSCodeNotebookModel)) {
             throw new Error('Incorrect NotebookModel, expected VSCodeNotebookModel');
         }
@@ -82,7 +85,7 @@ export class NotebookContentProvider implements INotebookContentProvider {
     }
     @captureTelemetry(Telemetry.Save, undefined, true)
     public async saveNotebook(document: NotebookDocument, cancellation: CancellationToken) {
-        const model = await this.notebookStorage.getOrCreateModel(document.uri, undefined, undefined, true);
+        const model = await this.notebookStorage.getOrCreateModel({ file: document.uri, isNative: true });
         if (cancellation.isCancellationRequested) {
             return;
         }
@@ -94,7 +97,7 @@ export class NotebookContentProvider implements INotebookContentProvider {
         document: NotebookDocument,
         cancellation: CancellationToken
     ): Promise<void> {
-        const model = await this.notebookStorage.getOrCreateModel(document.uri, undefined, undefined, true);
+        const model = await this.notebookStorage.getOrCreateModel({ file: document.uri, isNative: true });
         if (!cancellation.isCancellationRequested) {
             await this.notebookStorage.saveAs(model, targetResource);
         }
@@ -104,7 +107,7 @@ export class NotebookContentProvider implements INotebookContentProvider {
         _context: NotebookDocumentBackupContext,
         cancellation: CancellationToken
     ): Promise<NotebookDocumentBackup> {
-        const model = await this.notebookStorage.getOrCreateModel(document.uri, undefined, undefined, true);
+        const model = await this.notebookStorage.getOrCreateModel({ file: document.uri, isNative: true });
         const id = this.notebookStorage.generateBackupId(model);
         await this.notebookStorage.backup(model, cancellation, id);
         return {

src/client/datascience/notebook/notebookEditorProvider.ts

@@ -145,7 +145,7 @@ export class NotebookEditorProvider implements INotebookEditorProvider {
             return;
         }
         const uri = doc.uri;
-        const model = await this.storage.getOrCreateModel(uri, undefined, undefined, true);
+        const model = await this.storage.getOrCreateModel({ file: uri, isNative: true });
         if (model instanceof VSCodeNotebookModel) {
             model.associateNotebookDocument(doc);
         }

src/client/datascience/notebookStorage/nativeEditorProvider.ts

@@ -51,6 +51,7 @@ import {
     IJupyterDebugger,
     IJupyterVariableDataProviderFactory,
     IJupyterVariables,
+    IModelLoadOptions,
     INotebookEditor,
     INotebookEditorProvider,
     INotebookExporter,
@@ -128,30 +129,34 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit
         context: CustomDocumentOpenContext, // This has info about backups. right now we use our own data.
         _cancellation: CancellationToken
     ): Promise<CustomDocument> {
-        const model = await this.loadModel(uri, undefined, context.backupId);
+        const model = await this.loadModel({
+            file: uri,
+            backupId: context.backupId,
+            skipLoadingDirtyContents: context.backupId === undefined
+        });
         return {
             uri,
             dispose: () => model.dispose()
         };
     }
     public async saveCustomDocument(document: CustomDocument, cancellation: CancellationToken): Promise<void> {
-        const model = await this.loadModel(document.uri);
+        const model = await this.loadModel({ file: document.uri });
         return this.storage.save(model, cancellation);
     }
     public async saveCustomDocumentAs(document: CustomDocument, targetResource: Uri): Promise<void> {
-        const model = await this.loadModel(document.uri);
+        const model = await this.loadModel({ file: document.uri });
         return this.storage.saveAs(model, targetResource);
     }
     public async revertCustomDocument(document: CustomDocument, cancellation: CancellationToken): Promise<void> {
-        const model = await this.loadModel(document.uri);
+        const model = await this.loadModel({ file: document.uri });
         return this.storage.revert(model, cancellation);
     }
     public async backupCustomDocument(
         document: CustomDocument,
         _context: CustomDocumentBackupContext,
         cancellation: CancellationToken
     ): Promise<CustomDocumentBackup> {
-        const model = await this.loadModel(document.uri);
+        const model = await this.loadModel({ file: document.uri });
         const id = this.storage.generateBackupId(model);
         await this.storage.backup(model, cancellation, id);
         return {
@@ -167,7 +172,7 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit
 
     public async resolveCustomDocument(document: CustomDocument): Promise<void> {
         this.customDocuments.set(document.uri.fsPath, document);
-        await this.loadModel(document.uri);
+        await this.loadModel({ file: document.uri });
     }
 
     public async open(file: Uri): Promise<INotebookEditor> {
@@ -199,30 +204,26 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit
     }
 
     @captureTelemetry(Telemetry.CreateNewNotebook, undefined, false)
-    public async createNew(contents?: string, title?: string): Promise<INotebookEditor> {
+    public async createNew(possibleContents?: string, title?: string): Promise<INotebookEditor> {
         // Create a new URI for the dummy file using our root workspace path
         const uri = this.getNextNewNotebookUri(title);
 
         // Set these contents into the storage before the file opens. Make sure not
         // load from the memento storage though as this is an entirely brand new file.
-        await this.loadModel(uri, contents, true);
+        await this.loadModel({ file: uri, possibleContents, skipLoadingDirtyContents: true });
 
         return this.open(uri);
     }
 
-    public async loadModel(file: Uri, contents?: string, skipDirtyContents?: boolean): Promise<INotebookModel>;
-    // tslint:disable-next-line: unified-signatures
-    public async loadModel(file: Uri, contents?: string, backupId?: string): Promise<INotebookModel>;
-    // tslint:disable-next-line: no-any
-    public async loadModel(file: Uri, contents?: string, options?: any): Promise<INotebookModel> {
+    public async loadModel(options: IModelLoadOptions): Promise<INotebookModel> {
         // Get the model that may match this file
-        let model = [...this.models.values()].find((m) => this.fs.arePathsSame(m.file, file));
+        let model = [...this.models.values()].find((m) => this.fs.arePathsSame(m.file, options.file));
         if (!model) {
             // Every time we load a new untitled file, up the counter past the max value for this counter
-            this.untitledCounter = getNextUntitledCounter(file, this.untitledCounter);
+            this.untitledCounter = getNextUntitledCounter(options.file, this.untitledCounter);
 
             // Load our model from our storage object.
-            model = await this.storage.getOrCreateModel(file, contents, options);
+            model = await this.storage.getOrCreateModel(options);
 
             // Make sure to listen to events on the model
             this.trackModel(model);
@@ -273,7 +274,7 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit
     protected async loadNotebookEditor(resource: Uri, panel?: WebviewPanel) {
         try {
             // Get the model
-            const model = await this.loadModel(resource);
+            const model = await this.loadModel({ file: resource });
 
             // Load it (should already be visible)
             return this.createNotebookEditor(model, panel);

src/client/datascience/notebookStorage/nativeEditorStorage.ts

@@ -17,6 +17,7 @@ import {
     CellState,
     IDataScienceFileSystem,
     IJupyterExecution,
+    IModelLoadOptions,
     INotebookModel,
     INotebookStorage,
     ITrustService
@@ -75,27 +76,8 @@ export class NativeEditorStorage implements INotebookStorage {
     public get(_file: Uri): INotebookModel | undefined {
         return undefined;
     }
-    public getOrCreateModel(
-        file: Uri,
-        possibleContents?: string,
-        backupId?: string,
-        forVSCodeNotebook?: boolean
-    ): Promise<INotebookModel>;
-    public getOrCreateModel(
-        file: Uri,
-        possibleContents?: string,
-        // tslint:disable-next-line: unified-signatures
-        skipDirtyContents?: boolean,
-        forVSCodeNotebook?: boolean
-    ): Promise<INotebookModel>;
-    public getOrCreateModel(
-        file: Uri,
-        possibleContents?: string,
-        // tslint:disable-next-line: no-any
-        options?: any,
-        forVSCodeNotebook?: boolean
-    ): Promise<INotebookModel> {
-        return this.loadFromFile(file, possibleContents, options, forVSCodeNotebook);
+    public getOrCreateModel(options: IModelLoadOptions): Promise<INotebookModel> {
+        return this.loadFromFile(options);
     }
     public async save(model: INotebookModel, _cancellation: CancellationToken): Promise<void> {
         const contents = model.getContent();
@@ -154,8 +136,8 @@ export class NativeEditorStorage implements INotebookStorage {
     }
 
     public async revert(model: INotebookModel, _cancellation: CancellationToken): Promise<void> {
-        // Revert to what is in the hot exit file
-        await this.loadFromFile(model.file);
+        // Revert to what is in the real file. This is only used for the custom editor
+        await this.loadFromFile({ file: model.file, skipLoadingDirtyContents: true });
     }
 
     public async deleteBackup(model: INotebookModel, backupId: string): Promise<void> {
@@ -253,54 +235,44 @@ export class NativeEditorStorage implements INotebookStorage {
             noop();
         }
     }
-    private loadFromFile(
-        file: Uri,
-        possibleContents?: string,
-        backupId?: string,
-        forVSCodeNotebook?: boolean
-    ): Promise<INotebookModel>;
-    private loadFromFile(
-        file: Uri,
-        possibleContents?: string,
-        // tslint:disable-next-line: unified-signatures
-        skipDirtyContents?: boolean,
-        forVSCodeNotebook?: boolean
-    ): Promise<INotebookModel>;
-    private async loadFromFile(
-        file: Uri,
-        possibleContents?: string,
-        options?: boolean | string,
-        forVSCodeNotebook?: boolean
-    ): Promise<INotebookModel> {
+    private async loadFromFile(options: IModelLoadOptions): Promise<INotebookModel> {
         try {
             // Attempt to read the contents if a viable file
-            const contents = NativeEditorStorage.isUntitledFile(file) ? possibleContents : await this.fs.readFile(file);
+            const contents = NativeEditorStorage.isUntitledFile(options.file)
+                ? options.possibleContents
+                : await this.fs.readFile(options.file);
 
-            const skipDirtyContents = typeof options === 'boolean' ? options : !!options;
-            // Use backupId provided, else use static storage key.
-            const backupId =
-                typeof options === 'string' ? options : skipDirtyContents ? undefined : this.getStaticStorageKey(file);
+            // Get backup id from the options if available.
+            const backupId = options.backupId ? options.backupId : this.getStaticStorageKey(options.file);
 
             // If skipping dirty contents, delete the dirty hot exit file now
-            if (skipDirtyContents) {
-                await this.clearHotExit(file, backupId);
+            if (options.skipLoadingDirtyContents) {
+                await this.clearHotExit(options.file, backupId);
             }
 
             // See if this file was stored in storage prior to shutdown
-            const dirtyContents = skipDirtyContents ? undefined : await this.getStoredContents(file, backupId);
+            const dirtyContents = options.skipLoadingDirtyContents
+                ? undefined
+                : await this.getStoredContents(options.file, backupId);
             if (dirtyContents) {
                 // This means we're dirty. Indicate dirty and load from this content
-                return this.loadContents(file, dirtyContents, true, forVSCodeNotebook);
+                return this.loadContents(options.file, dirtyContents, true, options.isNative);
             } else {
                 // Load without setting dirty
-                return this.loadContents(file, contents, undefined, forVSCodeNotebook);
+                return this.loadContents(options.file, contents, undefined, options.isNative);
             }
         } catch (ex) {
             // May not exist at this time. Should always have a single cell though
-            traceError(`Failed to load notebook file ${file.toString()}`, ex);
+            traceError(`Failed to load notebook file ${options.file.toString()}`, ex);
             return this.factory.createModel(
-                { trusted: true, file, cells: [], crypto: this.crypto, globalMemento: this.globalStorage },
-                forVSCodeNotebook
+                {
+                    trusted: true,
+                    file: options.file,
+                    cells: [],
+                    crypto: this.crypto,
+                    globalMemento: this.globalStorage
+                },
+                options.isNative
             );
         }
     }