Add command to clear extension related storage (#15883)

* Add command to clear extension related storage

* News entry

* Clean up

* Fix linting tests

* Fix linting

* Add tests

* Capitalize

news/1 Enhancements/15883.md

@@ -0,0 +1 @@
+Added command `Python: Clear internal extension cache` to clear extension related cache.

package.json

@@ -103,6 +103,11 @@
             }
         ],
         "commands": [
+            {
+                "command": "python.clearPersistentStorage",
+                "title": "%python.command.python.clearPersistentStorage.title%",
+                "category": "Python"
+            },
             {
                 "command": "python.enableSourceMapSupport",
                 "title": "%python.command.python.enableSourceMapSupport.title%",

package.nls.json

@@ -34,6 +34,7 @@
     "python.command.python.enableLinting.title": "Enable/Disable Linting",
     "python.command.python.runLinting.title": "Run Linting",
     "python.command.python.enableSourceMapSupport.title": "Enable Source Map Support For Extension Debugging",
+    "python.command.python.clearPersistentStorage.title": "Clear Internal Extension Cache",
     "python.command.python.startPage.open.title": "Open Start Page",
     "python.command.python.analysis.clearCache.title": "Clear Module Analysis Cache",
     "python.command.python.analysis.restartLanguageServer.title": "Restart Language Server",

src/client/common/constants.ts

@@ -63,6 +63,7 @@ export namespace Commands {
     export const SwitchToInsidersWeekly = 'python.switchToWeeklyChannel';
     export const PickLocalProcess = 'python.pickLocalProcess';
     export const GetSelectedInterpreterPath = 'python.interpreterPath';
+    export const ClearStorage = 'python.clearPersistentStorage';
     export const ClearWorkspaceInterpreter = 'python.clearWorkspaceInterpreter';
     export const ResetInterpreterSecurityStorage = 'python.resetInterpreterSecurityStorage';
     export const OpenStartPage = 'python.startPage.open';

src/client/common/persistentState.ts

@@ -5,6 +5,9 @@
 
 import { inject, injectable, named } from 'inversify';
 import { Memento } from 'vscode';
+import { IExtensionSingleActivationService } from '../activation/types';
+import { ICommandManager } from './application/types';
+import { Commands } from './constants';
 import {
     GLOBAL_MEMENTO,
     IExtensionContext,
@@ -44,26 +47,72 @@ export class PersistentState<T> implements IPersistentState<T> {
     }
 }
 
+const GLOBAL_PERSISTENT_KEYS = 'PYTHON_EXTENSION_GLOBAL_STORAGE_KEYS';
+const WORKSPACE_PERSISTENT_KEYS = 'PYTHON_EXTENSION_WORKSPACE_STORAGE_KEYS';
+type keysStorage = { key: string; defaultValue: unknown };
+
 @injectable()
-export class PersistentStateFactory implements IPersistentStateFactory {
+export class PersistentStateFactory implements IPersistentStateFactory, IExtensionSingleActivationService {
+    private readonly globalKeysStorage = new PersistentState<keysStorage[]>(
+        this.globalState,
+        GLOBAL_PERSISTENT_KEYS,
+        [],
+    );
+    private readonly workspaceKeysStorage = new PersistentState<keysStorage[]>(
+        this.workspaceState,
+        WORKSPACE_PERSISTENT_KEYS,
+        [],
+    );
     constructor(
         @inject(IMemento) @named(GLOBAL_MEMENTO) private globalState: Memento,
         @inject(IMemento) @named(WORKSPACE_MEMENTO) private workspaceState: Memento,
+        @inject(ICommandManager) private cmdManager: ICommandManager,
     ) {}
+
+    public async activate(): Promise<void> {
+        this.cmdManager.registerCommand(Commands.ClearStorage, this.cleanAllPersistentStates.bind(this));
+    }
+
     public createGlobalPersistentState<T>(
         key: string,
         defaultValue?: T,
         expiryDurationMs?: number,
     ): IPersistentState<T> {
+        if (!this.globalKeysStorage.value.includes({ key, defaultValue })) {
+            this.globalKeysStorage.updateValue([{ key, defaultValue }, ...this.globalKeysStorage.value]).ignoreErrors();
+        }
         return new PersistentState<T>(this.globalState, key, defaultValue, expiryDurationMs);
     }
+
     public createWorkspacePersistentState<T>(
         key: string,
         defaultValue?: T,
         expiryDurationMs?: number,
     ): IPersistentState<T> {
+        if (!this.workspaceKeysStorage.value.includes({ key, defaultValue })) {
+            this.workspaceKeysStorage
+                .updateValue([{ key, defaultValue }, ...this.workspaceKeysStorage.value])
+                .ignoreErrors();
+        }
         return new PersistentState<T>(this.workspaceState, key, defaultValue, expiryDurationMs);
     }
+
+    private async cleanAllPersistentStates(): Promise<void> {
+        await Promise.all(
+            this.globalKeysStorage.value.map(async (keyContent) => {
+                const storage = this.createGlobalPersistentState(keyContent.key);
+                await storage.updateValue(keyContent.defaultValue);
+            }),
+        );
+        await Promise.all(
+            this.workspaceKeysStorage.value.map(async (keyContent) => {
+                const storage = this.createWorkspacePersistentState(keyContent.key);
+                await storage.updateValue(keyContent.defaultValue);
+            }),
+        );
+        await this.globalKeysStorage.updateValue([]);
+        await this.workspaceKeysStorage.updateValue([]);
+    }
 }
 
 /////////////////////////////
@@ -79,6 +128,10 @@ interface IPersistentStorage<T> {
  * Build a global storage object for the given key.
  */
 export function getGlobalStorage<T>(context: IExtensionContext, key: string): IPersistentStorage<T> {
+    const globalKeysStorage = new PersistentState<keysStorage[]>(context.globalState, GLOBAL_PERSISTENT_KEYS, []);
+    if (!globalKeysStorage.value.includes({ key, defaultValue: undefined })) {
+        globalKeysStorage.updateValue([{ key, defaultValue: undefined }, ...globalKeysStorage.value]).ignoreErrors();
+    }
     const raw = new PersistentState<T>(context.globalState, key);
     return {
         // We adapt between PersistentState and IPersistentStorage.

src/client/common/serviceRegistry.ts

@@ -131,6 +131,7 @@ export function registerTypes(serviceManager: IServiceManager) {
     serviceManager.addSingleton<IExtensions>(IExtensions, Extensions);
     serviceManager.addSingleton<IRandom>(IRandom, Random);
     serviceManager.addSingleton<IPersistentStateFactory>(IPersistentStateFactory, PersistentStateFactory);
+    serviceManager.addBinding(IPersistentStateFactory, IExtensionSingleActivationService);
     serviceManager.addSingleton<ITerminalServiceFactory>(ITerminalServiceFactory, TerminalServiceFactory);
     serviceManager.addSingleton<IPathUtils>(IPathUtils, PathUtils);
     serviceManager.addSingleton<IApplicationShell>(IApplicationShell, ApplicationShell);

src/test/common/persistentState.unit.test.ts

@@ -0,0 +1,90 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { expect } from 'chai';
+import * as TypeMoq from 'typemoq';
+import { Memento } from 'vscode';
+import { IExtensionSingleActivationService } from '../../client/activation/types';
+import { ICommandManager } from '../../client/common/application/types';
+import { Commands } from '../../client/common/constants';
+import { PersistentStateFactory } from '../../client/common/persistentState';
+import { IDisposable, IPersistentStateFactory } from '../../client/common/types';
+import { MockMemento } from '../mocks/mementos';
+
+suite('Persistent State', () => {
+    let cmdManager: TypeMoq.IMock<ICommandManager>;
+    let persistentStateFactory: IPersistentStateFactory & IExtensionSingleActivationService;
+    let workspaceMemento: Memento;
+    let globalMemento: Memento;
+    setup(() => {
+        cmdManager = TypeMoq.Mock.ofType<ICommandManager>();
+        workspaceMemento = new MockMemento();
+        globalMemento = new MockMemento();
+        persistentStateFactory = new PersistentStateFactory(globalMemento, workspaceMemento, cmdManager.object);
+    });
+
+    test('Global states created are restored on invoking clean storage command', async () => {
+        let clearStorageCommand: (() => Promise<void>) | undefined;
+        cmdManager
+            .setup((c) => c.registerCommand(Commands.ClearStorage, TypeMoq.It.isAny()))
+            .callback((_, c) => {
+                clearStorageCommand = c;
+            })
+            .returns(() => TypeMoq.Mock.ofType<IDisposable>().object);
+
+        // Register command to clean storage
+        await persistentStateFactory.activate();
+
+        expect(clearStorageCommand).to.not.equal(undefined, 'Callback not registered');
+
+        const globalKey1State = persistentStateFactory.createGlobalPersistentState('key1', 'defaultKey1Value');
+        await globalKey1State.updateValue('key1Value');
+        const globalKey2State = persistentStateFactory.createGlobalPersistentState<string | undefined>(
+            'key2',
+            undefined,
+        );
+        await globalKey2State.updateValue('key2Value');
+
+        // Verify states are updated correctly
+        expect(globalKey1State.value).to.equal('key1Value');
+        expect(globalKey2State.value).to.equal('key2Value');
+
+        await clearStorageCommand!(); // Invoke command
+
+        // Verify states are now reset to their default value.
+        expect(globalKey1State.value).to.equal('defaultKey1Value');
+        expect(globalKey2State.value).to.equal(undefined);
+    });
+
+    test('Workspace states created are restored on invoking clean storage command', async () => {
+        let clearStorageCommand: (() => Promise<void>) | undefined;
+        cmdManager
+            .setup((c) => c.registerCommand(Commands.ClearStorage, TypeMoq.It.isAny()))
+            .callback((_, c) => {
+                clearStorageCommand = c;
+            })
+            .returns(() => TypeMoq.Mock.ofType<IDisposable>().object);
+
+        // Register command to clean storage
+        await persistentStateFactory.activate();
+
+        expect(clearStorageCommand).to.not.equal(undefined, 'Callback not registered');
+
+        const workspaceKey1State = persistentStateFactory.createWorkspacePersistentState('key1');
+        await workspaceKey1State.updateValue('key1Value');
+        const workspaceKey2State = persistentStateFactory.createWorkspacePersistentState('key2', 'defaultKey2Value');
+        await workspaceKey2State.updateValue('key2Value');
+
+        // Verify states are updated correctly
+        expect(workspaceKey1State.value).to.equal('key1Value');
+        expect(workspaceKey2State.value).to.equal('key2Value');
+
+        await clearStorageCommand!(); // Invoke command
+
+        // Verify states are now reset to their default value.
+        expect(workspaceKey1State.value).to.equal(undefined);
+        expect(workspaceKey2State.value).to.equal('defaultKey2Value');
+    });
+});

src/test/linters/lint.provider.test.ts

@@ -7,7 +7,12 @@ import { Container } from 'inversify';
 import * as TypeMoq from 'typemoq';
 import * as vscode from 'vscode';
 import { LanguageServerType } from '../../client/activation/types';
-import { IApplicationShell, IDocumentManager, IWorkspaceService } from '../../client/common/application/types';
+import {
+    IApplicationShell,
+    ICommandManager,
+    IDocumentManager,
+    IWorkspaceService,
+} from '../../client/common/application/types';
 import { PersistentStateFactory } from '../../client/common/persistentState';
 import { IFileSystem } from '../../client/common/platform/types';
 import {
@@ -109,6 +114,10 @@ suite('Linting - Provider', () => {
         serviceManager.addSingleton<IPersistentStateFactory>(IPersistentStateFactory, PersistentStateFactory);
         serviceManager.addSingleton<vscode.Memento>(IMemento, MockMemento, GLOBAL_MEMENTO);
         serviceManager.addSingleton<vscode.Memento>(IMemento, MockMemento, WORKSPACE_MEMENTO);
+        serviceManager.addSingletonInstance<ICommandManager>(
+            ICommandManager,
+            TypeMoq.Mock.ofType<ICommandManager>().object,
+        );
         lm = new LinterManager(serviceContainer, workspaceService.object);
         serviceManager.addSingletonInstance<ILinterManager>(ILinterManager, lm);
         emitter = new vscode.EventEmitter<vscode.TextDocument>();