Fix storage not being used

news/2 Fixes/11557.md

@@ -0,0 +1 @@
+When VS quits, make sure to save contents of notebook for next reopen.

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

@@ -139,6 +139,11 @@ export class NativeEditorOldWebView extends NativeEditor {
         // Update our title to match
         this.setTitle(path.basename(model.file.fsPath));
 
+        // Update dirty if model started out that way
+        if (this.model?.isDirty) {
+            this.setDirty().ignoreErrors();
+        }
+
         // Show ourselves
         await this.show();
         this.model?.changed(() => {

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

@@ -8,7 +8,14 @@ import { concatMultilineStringInput, splitMultilineString } from '../../../datas
 import { createCodeCell } from '../../../datascience-ui/common/cellFactory';
 import { traceError } from '../../common/logger';
 import { IFileSystem } from '../../common/platform/types';
-import { GLOBAL_MEMENTO, ICryptoUtils, IExtensionContext, IMemento, WORKSPACE_MEMENTO } from '../../common/types';
+import {
+    GLOBAL_MEMENTO,
+    ICryptoUtils,
+    IDisposableRegistry,
+    IExtensionContext,
+    IMemento,
+    WORKSPACE_MEMENTO
+} from '../../common/types';
 import { noop } from '../../common/utils/misc';
 import { PythonInterpreter } from '../../interpreter/contracts';
 import { Identifiers, KnownNotebookLanguages, Telemetry } from '../constants';
@@ -19,6 +26,8 @@ import { CellState, ICell, IJupyterExecution, IJupyterKernelSpec, INotebookModel
 
 // tslint:disable-next-line:no-require-imports no-var-requires
 import detectIndent = require('detect-indent');
+import { IDisposable } from 'monaco-editor';
+import { IWorkspaceService } from '../../common/application/types';
 import { sendTelemetryEvent } from '../../telemetry';
 import { pruneCell } from '../common';
 // tslint:disable-next-line:no-require-imports no-var-requires
@@ -35,7 +44,7 @@ interface INativeEditorStorageState {
 }
 
 @injectable()
-export class NativeEditorStorage implements INotebookModel, INotebookStorage {
+export class NativeEditorStorage implements INotebookModel, INotebookStorage, IDisposable {
     public get isDirty(): boolean {
         return this._state.changeCount !== this._state.saveChangeCount;
     }
@@ -67,15 +76,24 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
     };
     private indentAmount: string = ' ';
     private debouncedWriteToStorage = debounce(this.writeToStorage.bind(this), 250);
+    private disposed = false;
 
     constructor(
         @inject(IJupyterExecution) private jupyterExecution: IJupyterExecution,
         @inject(IFileSystem) private fileSystem: IFileSystem,
         @inject(ICryptoUtils) private crypto: ICryptoUtils,
         @inject(IExtensionContext) private context: IExtensionContext,
         @inject(IMemento) @named(GLOBAL_MEMENTO) private globalStorage: Memento,
-        @inject(IMemento) @named(WORKSPACE_MEMENTO) private localStorage: Memento
-    ) {}
+        @inject(IMemento) @named(WORKSPACE_MEMENTO) private localStorage: Memento,
+        @inject(IWorkspaceService) private workspaceService: IWorkspaceService,
+        @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry
+    ) {
+        disposableRegistry.push(this);
+    }
+
+    public dispose() {
+        this.disposed = true;
+    }
 
     public async load(file: Uri, possibleContents?: string): Promise<INotebookModel> {
         // Reload our cells
@@ -141,9 +159,21 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
         // Write but debounced (wait at least 250 ms)
         return this.debouncedWriteToStorage(filePath, specialContents, cancelToken);
     }
+
+    private async saveToStorage(): Promise<void> {
+        // Skip doing this if auto save is enabled or is untitled
+        const filesConfig = this.workspaceService.getConfiguration('files', this.file);
+        const autoSave = filesConfig.get('autoSave', 'off');
+        if (autoSave === 'off' && !this.isUntitled) {
+            // Write but debounced (wait at least 250 ms)
+            return this.storeContentsInHotExitFile(undefined);
+        }
+    }
+
     private async writeToStorage(filePath: string, contents?: string, cancelToken?: CancellationToken): Promise<void> {
         try {
-            if (!cancelToken?.isCancellationRequested) {
+            // Only write to storage if not disposed.
+            if (!cancelToken?.isCancellationRequested && !this.disposed) {
                 if (contents) {
                     await this.fileSystem.createDirectory(path.dirname(filePath));
                     if (!cancelToken?.isCancellationRequested) {
@@ -198,6 +228,10 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
         // Forward onto our listeners if necessary
         if (changed || this.isDirty !== oldDirty) {
             this._changedEmitter.fire({ ...change, newDirty: this.isDirty, oldDirty });
+            if (this.isDirty) {
+                // Save to temp storage so we don't lose the file if the user exits VS code
+                this.saveToStorage().ignoreErrors();
+            }
         }
         // Slightly different for the event we send to VS code. Skip version and file changes. Only send user events.
         if (
@@ -451,6 +485,7 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
             const dirtyContents = await this.getStoredContents();
             if (dirtyContents) {
                 // This means we're dirty. Indicate dirty and load from this content
+                this._state.saveChangeCount = -1; // Indicates dirty
                 this.loadContents(dirtyContents);
             } else {
                 // Load without setting dirty
@@ -621,8 +656,8 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
             if (data && !this.isUntitled && data.contents) {
                 return data.contents;
             }
-        } catch {
-            noop();
+        } catch (exc) {
+            traceError(`Exception reading from temporary storage for ${key}`, exc);
         }
     }
 

src/test/datascience/dataScienceIocContainer.ts

@@ -3,6 +3,8 @@
 //tslint:disable:trailing-comma no-any
 import * as child_process from 'child_process';
 import { ReactWrapper } from 'enzyme';
+import * as fs from 'fs-extra';
+import * as glob from 'glob';
 import { interfaces } from 'inversify';
 import * as os from 'os';
 import * as path from 'path';
@@ -24,6 +26,7 @@ import {
 import * as vsls from 'vsls/vscode';
 import { KernelDaemonPool } from '../../client/datascience/kernel-launcher/kernelDaemonPool';
 
+import { promisify } from 'util';
 import { LanguageServerExtensionActivationService } from '../../client/activation/activationService';
 import { LanguageServerDownloader } from '../../client/activation/common/downloader';
 import { JediExtensionActivator } from '../../client/activation/jedi';
@@ -479,6 +482,18 @@ export class DataScienceIocContainer extends UnitTestIocContainer {
     }
 
     public async dispose(): Promise<void> {
+        try {
+            // Make sure to delete any temp files written by native editor storage
+            const globPr = promisify(glob);
+            const tempLocation = this.serviceManager.get<IExtensionContext>(IExtensionContext).globalStoragePath;
+            const tempFiles = await globPr(`${tempLocation}/*.ipynb`);
+            if (tempFiles && tempFiles.length) {
+                await Promise.all(tempFiles.map((t) => fs.remove(t)));
+            }
+        } catch (exc) {
+            // tslint:disable-next-line: no-console
+            console.log(`Exception on cleanup: ${exc}`);
+        }
         await this.asyncRegistry.dispose();
         await super.dispose();
         this.disposed = true;

src/test/datascience/interactive-ipynb/nativeEditorStorage.unit.test.ts

@@ -11,6 +11,7 @@ import {
     ConfigurationChangeEvent,
     ConfigurationTarget,
     EventEmitter,
+    FileType,
     TextEditor,
     Uri,
     WorkspaceConfiguration
@@ -31,6 +32,7 @@ import { ConfigurationService } from '../../../client/common/configuration/servi
 import { CryptoUtils } from '../../../client/common/crypto';
 import { IFileSystem } from '../../../client/common/platform/types';
 import { IConfigurationService, ICryptoUtils, IDisposable, IExtensionContext } from '../../../client/common/types';
+import { sleep } from '../../../client/common/utils/async';
 import { EXTENSION_ROOT_DIR } from '../../../client/constants';
 import {
     IEditorContentChange,
@@ -87,7 +89,7 @@ class MockWorkspaceConfiguration implements WorkspaceConfiguration {
 }
 
 // tslint:disable: max-func-body-length
-suite('Data Science - Native Editor Storage', () => {
+suite('DataScience - Native Editor Storage', () => {
     let workspace: IWorkspaceService;
     let configService: IConfigurationService;
     let fileSystem: typemoq.IMock<IFileSystem>;
@@ -321,6 +323,7 @@ suite('Data Science - Native Editor Storage', () => {
 
         let listener: IWebPanelMessageListener;
         const webPanel = mock(WebPanel);
+        const startTime = Date.now();
         class WebPanelCreateMatcher extends Matcher {
             public match(value: any) {
                 listener = value.listener;
@@ -351,16 +354,26 @@ suite('Data Science - Native Editor Storage', () => {
             .returns((_a1) => {
                 return Promise.resolve(lastWriteFileValue);
             });
+        fileSystem
+            .setup((f) => f.stat(typemoq.It.isAny()))
+            .returns((_a1) => {
+                return Promise.resolve({ mtime: startTime, type: FileType.File, ctime: startTime, size: 100 });
+            });
+        storage = createStorage();
+    });
 
-        storage = new NativeEditorStorage(
+    function createStorage() {
+        return new NativeEditorStorage(
             instance(executionProvider),
             fileSystem.object, // Use typemoq so can save values in returns
             instance(crypto),
             context.object,
             globalMemento,
-            localMemento
+            localMemento,
+            instance(workspace),
+            []
         );
-    });
+    }
 
     teardown(() => {
         globalMemento.clear();
@@ -480,6 +493,47 @@ suite('Data Science - Native Editor Storage', () => {
         expect(cells).to.be.lengthOf(1);
     });
 
+    test('Editing a file and closing will keep contents', async () => {
+        await filesConfig?.update('autoSave', 'off');
+
+        await storage.load(baseUri);
+        expect(storage.isDirty).to.be.equal(false, 'Editor should not be dirty');
+        editCell(
+            [
+                {
+                    range: {
+                        startLineNumber: 2,
+                        startColumn: 1,
+                        endLineNumber: 2,
+                        endColumn: 1
+                    },
+                    rangeOffset: 4,
+                    rangeLength: 0,
+                    text: 'a',
+                    position: {
+                        lineNumber: 1,
+                        column: 1
+                    }
+                }
+            ],
+            storage.cells[1],
+            'a'
+        );
+
+        // Wait for a second so it has time to update
+        await sleep(1000);
+
+        // Recreate
+        storage = createStorage();
+        await storage.load(baseUri);
+
+        const cells = storage.cells;
+        expect(cells).to.be.lengthOf(3);
+        expect(cells[1].id).to.be.match(/NotebookImport#1/);
+        expect(concatMultilineStringInput(cells[1].data.source)).to.be.equals('b=2\nab');
+        expect(storage.isDirty).to.be.equal(true, 'Editor should be dirty');
+    });
+
     test('Opening file with local storage but no global will still open with old contents', async () => {
         await filesConfig?.update('autoSave', 'off');
         // This test is really for making sure when a user upgrades to a new extension, we still have their old storage

src/test/datascience/nativeEditor.functional.test.tsx

@@ -83,6 +83,9 @@ import {
 use(chaiAsPromised);
 
 // tslint:disable:max-func-body-length trailing-comma no-any no-multiline-string
+async function updateFileConfig(ioc: DataScienceIocContainer, key: string, value: any) {
+    return ioc.get<IWorkspaceService>(IWorkspaceService).getConfiguration('file').update(key, value);
+}
 
 suite('DataScience Native Editor', () => {
     const originalPlatform = window.navigator.platform;
@@ -641,6 +644,8 @@ df.head()`;
                 runMountedTest(
                     'RunAllCells',
                     async (wrapper) => {
+                        // Make sure we don't write to storage for the notebook. It messes up other tests
+                        await updateFileConfig(ioc, 'autoSave', 'onFocusChange');
                         addMockData(ioc, 'print(1)\na=1', 1);
                         addMockData(ioc, 'a=a+1\nprint(a)', 2);
                         addMockData(ioc, 'print(a+1)', 3);
@@ -2122,17 +2127,10 @@ df.head()`;
                         await addCell(wrapper, ioc, 'a', false);
                     }
 
-                    async function updateFileConfig(key: string, value: any) {
-                        return ioc
-                            .get<IWorkspaceService>(IWorkspaceService)
-                            .getConfiguration('file')
-                            .update(key, value);
-                    }
-
                     test('Auto save notebook every 1s', async () => {
                         // Configure notebook to save automatically ever 1s.
-                        await updateFileConfig('autoSave', 'afterDelay');
-                        await updateFileConfig('autoSaveDelay', 1_000);
+                        await updateFileConfig(ioc, 'autoSave', 'afterDelay');
+                        await updateFileConfig(ioc, 'autoSaveDelay', 1_000);
                         ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath);
 
                         /**
@@ -2165,8 +2163,8 @@ df.head()`;
 
                     test('File saved with same format', async () => {
                         // Configure notebook to save automatically ever 1s.
-                        await updateFileConfig('autoSave', 'afterDelay');
-                        await updateFileConfig('autoSaveDelay', 2_000);
+                        await updateFileConfig(ioc, 'autoSave', 'afterDelay');
+                        await updateFileConfig(ioc, 'autoSaveDelay', 2_000);
 
                         ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath);
                         const notebookFileContents = await fs.readFile(notebookFile.filePath, 'utf8');
@@ -2190,8 +2188,8 @@ df.head()`;
                         const notebookFileContents = await fs.readFile(notebookFile.filePath, 'utf8');
 
                         // Configure notebook to to never save.
-                        await updateFileConfig('autoSave', 'off');
-                        await updateFileConfig('autoSaveDelay', 1_000);
+                        await updateFileConfig(ioc, 'autoSave', 'off');
+                        await updateFileConfig(ioc, 'autoSaveDelay', 1_000);
 
                         // Update the settings and wait for the component to receive it and process it.
                         const promise = waitForMessage(ioc, InteractiveWindowMessages.SettingsUpdated);
@@ -2231,7 +2229,7 @@ df.head()`;
                         await dirtyPromise;
 
                         // Configure notebook to save when active editor changes.
-                        await updateFileConfig('autoSave', 'onFocusChange');
+                        await updateFileConfig(ioc, 'autoSave', 'onFocusChange');
                         ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath);
 
                         // Now that the notebook is dirty, change the active editor.
@@ -2263,7 +2261,7 @@ df.head()`;
                         await dirtyPromise;
 
                         // Configure notebook to save when window state changes.
-                        await updateFileConfig('autoSave', 'onWindowChange');
+                        await updateFileConfig(ioc, 'autoSave', 'onWindowChange');
                         ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath);
 
                         // Now that the notebook is dirty, change the active editor.
@@ -2289,7 +2287,7 @@ df.head()`;
                         await dirtyPromise;
 
                         // Configure notebook to save when active editor changes.
-                        await updateFileConfig('autoSave', configSetting);
+                        await updateFileConfig(ioc, 'autoSave', configSetting);
                         ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath);
 
                         // Now that the notebook is dirty, send notification about changes to window state.
@@ -2321,7 +2319,7 @@ df.head()`;
                         await dirtyPromise;
 
                         // Configure notebook to save when active editor changes.
-                        await updateFileConfig('autoSave', 'onFocusChange');
+                        await updateFileConfig(ioc, 'autoSave', 'onFocusChange');
                         ioc.forceSettingsChanged(undefined, ioc.getSettings().pythonPath);
 
                         // Force a view state change