Refactor StateStore to make keys more explicit at the point of use. (#913)

kartikgupta-db · web-flow · commit 5b8fb235f854 · 2023-10-23T13:19:02.000Z
## Changes
Advantages of this approach
* Adding new variables does not require creating an explicit setter and
getter. Default setters and getters should work out of the box.
* Easier to see all the information in 1 location. 
* The keys are immutably exposed to point of use. It allows for
potential programatic access to these functions.
* Setters require calling the `set` function, which is async. This
prevents us from having to create even more functions just to make
setting a variable `await`able.

## Tests
&lt;!-- How is this tested? --&gt;
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
@@ -566,8 +566,10 @@ export async function activate(
             );
         })
         .finally(() => {
-            stateStorage.lastInstalledExtensionVersion =
-                packageMetadata.version;
+            stateStorage.set(
+                "databricks.lastInstalledExtensionVersion",
+                packageMetadata.version
+            );
         });
 
     CustomWhenContext.setActivated(true);
diff --git a/packages/databricks-vscode/src/language/ConfigureAutocomplete.ts b/packages/databricks-vscode/src/language/ConfigureAutocomplete.ts
@@ -106,12 +106,15 @@ export class ConfigureAutocomplete implements Disposable {
     }
 
     async configureCommand() {
-        this.stateStorage.skipAutocompleteConfigure = false;
+        this.stateStorage.set("databricks.autocompletion.skipConfigure", false);
         return this.configure(true);
     }
 
     private async configure(force = false) {
-        if (!force || this.stateStorage.skipAutocompleteConfigure) {
+        if (
+            !force ||
+            this.stateStorage.get("databricks.autocompletion.skipConfigure")
+        ) {
             return;
         }
 
@@ -142,7 +145,10 @@ export class ConfigureAutocomplete implements Disposable {
         }
 
         if (choice === "Never for this workspace") {
-            this.stateStorage.skipAutocompleteConfigure = true;
+            this.stateStorage.set(
+                "databricks.autocompletion.skipConfigure",
+                true
+            );
             return;
         }
 
diff --git a/packages/databricks-vscode/src/language/DbConnectInstallPrompt.ts b/packages/databricks-vscode/src/language/DbConnectInstallPrompt.ts
@@ -21,7 +21,9 @@ export class DbConnectInstallPrompt implements Disposable {
         if (
             advertisement &&
             executable &&
-            this.stateStorage.skippedEnvsForDbConnect.includes(executable)
+            this.stateStorage
+                .get("databricks.debugging.skipDbConnectInstallForEnvs")
+                .includes(executable)
         ) {
             return;
         }
@@ -82,7 +84,10 @@ export class DbConnectInstallPrompt implements Disposable {
 
             case "Never for this environment":
                 if (executable) {
-                    this.stateStorage.skipDbConnectInstallForEnv(executable);
+                    this.stateStorage.set(
+                        "databricks.debugging.skipDbConnectInstallForEnvs",
+                        [executable]
+                    );
                 }
                 break;
 
diff --git a/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts b/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts
@@ -45,10 +45,9 @@ export class MsPythonExtensionWrapper implements Disposable {
         if (this._terminal) {
             return this._terminal;
         }
-        const terminalName = `databricks-pip-${this.stateStorage.fixedUUID.slice(
-            0,
-            8
-        )}`;
+        const terminalName = `databricks-pip-${this.stateStorage
+            .get("databricks.fixedUUID")
+            .slice(0, 8)}`;
 
         this._terminal = window.createTerminal({
             name: terminalName,
diff --git a/packages/databricks-vscode/src/language/notebooks/NotebookAccessVerifier.ts b/packages/databricks-vscode/src/language/notebooks/NotebookAccessVerifier.ts
@@ -60,7 +60,6 @@ export class NotebookAccessVerifier extends MultiStepAccessVerifier {
                     "latest"
                 );
                 return true;
-                break;
 
             case "Change environment":
                 await this.pythonExtension.selectPythonInterpreter();
diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts
@@ -1,73 +1,100 @@
 import {randomUUID} from "crypto";
 import {ExtensionContext} from "vscode";
 
-export class StateStorage {
-    constructor(private context: ExtensionContext) {}
+/* eslint-disable @typescript-eslint/naming-convention */
+type KeyInfo<V> = {
+    location: "global" | "workspace";
+    defaultValue?: V;
+    getter?: (storage: StateStorage, value: V | undefined) => V | undefined;
+    setter?: (storage: StateStorage, value: V | undefined) => V | undefined;
+};
 
-    get fixedRandom() {
-        let randomNum = this.context.globalState.get<number>(
-            "databricks.fixedRandom"
-        );
-        if (!randomNum) {
-            randomNum = Math.random();
-            this.context.globalState.update(
-                "databricks.fixedRandom",
-                randomNum
-            );
-        }
-        return randomNum;
-    }
+function withType<V>() {
+    return function <D extends KeyInfo<V>>(data: D) {
+        return data as typeof data & {_type: V};
+    };
+}
 
-    get wsfsFeatureFlag() {
-        return true;
-    }
+const Keys = {
+    "databricks.clusterId": withType<string>()({
+        location: "workspace",
+    }),
 
-    get skipSwitchToWorkspace() {
-        return this.context.workspaceState.get(
-            "databricks.wsfs.skipSwitchToWorkspace",
-            false
-        );
-    }
+    "databricks.wsfs.skipSwitchToWorkspace": withType<boolean>()({
+        location: "workspace",
+        defaultValue: false,
+    }),
 
-    set skipSwitchToWorkspace(value: boolean) {
-        this.context.workspaceState.update(
-            "databricks.wsfs.skipSwitchToWorkspace",
-            value
-        );
-    }
+    "databricks.autocompletion.skipConfigure": withType<boolean>()({
+        location: "workspace",
+        defaultValue: false,
+    }),
 
-    get skipAutocompleteConfigure() {
-        return this.context.workspaceState.get(
-            "databricks.autocompletion.skipConfigure",
-            false
-        );
-    }
+    "databricks.fixedRandom": withType<number>()({
+        location: "global",
+        getter: (storage, value) => {
+            if (value === undefined) {
+                value = Math.random();
+                storage.set("databricks.fixedRandom", value);
+            }
+            return value;
+        },
+    }),
 
-    set skipAutocompleteConfigure(value: boolean) {
-        this.context.workspaceState.update(
-            "databricks.autocompletion.skipConfigure",
-            value
-        );
-    }
+    "databricks.fixedUUID": withType<string>()({
+        location: "workspace",
+        getter: (storage, value) => {
+            if (value === undefined) {
+                value = randomUUID();
+                storage.set("databricks.fixedUUID", value);
+            }
+            return value;
+        },
+    }),
 
-    get skippedEnvsForDbConnect() {
-        return this.context.globalState.get<string[]>(
-            "databricks.debugging.skipDbConnectInstallForEnvs",
-            []
-        );
-    }
+    "databricks.debugging.skipDbConnectInstallForEnvs": withType<string[]>()({
+        location: "global",
+        defaultValue: [],
+        setter: (storage, value) => {
+            if (value === undefined || value.length === 0) {
+                return undefined;
+            }
+            const currentEnvs: string[] = storage.get(
+                "databricks.debugging.skipDbConnectInstallForEnvs"
+            );
+            if (!currentEnvs.includes(value[0])) {
+                currentEnvs.push(value[0]);
+            }
+            return currentEnvs;
+        },
+    }),
 
-    skipDbConnectInstallForEnv(value: string) {
-        const currentEnvs = this.skippedEnvsForDbConnect;
-        if (!currentEnvs.includes(value)) {
-            currentEnvs.push(value);
-        }
-        this.context.globalState.update(
-            "databricks.debugging.skipDbConnectInstallForEnvs",
-            currentEnvs
-        );
-    }
+    "databricks.lastInstalledExtensionVersion": withType<string>()({
+        location: "workspace",
+        defaultValue: "0.0.0",
+    }),
+};
+
+type ValueType<K extends keyof typeof Keys> = (typeof Keys)[K]["_type"];
+type GetterReturnType<D extends KeyInfo<any>> = D extends {getter: infer G}
+    ? G extends (...args: any[]) => any
+        ? ReturnType<G>
+        : undefined
+    : undefined;
+
+type DefaultValue<K extends keyof typeof Keys> =
+    "defaultValue" extends keyof (typeof Keys)[K]
+        ? ValueType<K>
+        : GetterReturnType<(typeof Keys)[K]>;
+
+type KeyInfoWithType<V> = KeyInfo<V> & {
+    _type: V;
+    _getterType: V | never | undefined;
+};
+/* eslint-enable @typescript-eslint/naming-convention */
 
+export class StateStorage {
+    constructor(private context: ExtensionContext) {}
     get skippedEnvsForDatabricksSdk() {
         return this.context.globalState.get<string[]>(
             "databricks.debugging.skipDatabricksSdkInstallForEnvs",
@@ -85,29 +112,35 @@ export class StateStorage {
             currentEnvs
         );
     }
-
-    get lastInstalledExtensionVersion() {
-        return this.context.workspaceState.get<string>(
-            "databricks.lastInstalledExtensionVersion",
-            "0.0.0"
-        );
+    private getStateObject(location: "global" | "workspace") {
+        switch (location) {
+            case "workspace":
+                return this.context.workspaceState;
+            case "global":
+                return this.context.globalState;
+        }
     }
 
-    set lastInstalledExtensionVersion(value: string) {
-        this.context.workspaceState.update(
-            "databricks.lastInstalledExtensionVersion",
-            value
-        );
+    get<K extends keyof typeof Keys>(key: K): ValueType<K> | DefaultValue<K> {
+        const details = Keys[key] as KeyInfoWithType<ValueType<K>>;
+
+        const value =
+            this.getStateObject(details.location).get<ValueType<K>>(key) ??
+            details.defaultValue;
+
+        return (
+            details.getter !== undefined ? details.getter(this, value) : value
+        ) as ValueType<K> | DefaultValue<K>;
     }
 
-    get fixedUUID() {
-        let uuid = this.context.workspaceState.get<string>(
-            "databricks.fixedUUID"
-        );
-        if (!uuid) {
-            uuid = randomUUID();
-            this.context.workspaceState.update("databricks.fixedUUID", uuid);
-        }
-        return uuid;
+    async set<K extends keyof typeof Keys>(
+        key: K,
+        value: ValueType<K> | undefined
+    ) {
+        const details = Keys[key] as KeyInfoWithType<ValueType<K>>;
+        value =
+            details.setter !== undefined ? details.setter(this, value) : value;
+        await this.getStateObject(details.location).update(key, value);
+        return;
     }
 }
diff --git a/packages/databricks-vscode/src/whatsNewPopup.ts b/packages/databricks-vscode/src/whatsNewPopup.ts
@@ -44,7 +44,7 @@ export async function showWhatsNewPopup(
     }
 
     const previousVersion =
-        semver.parse(storage.lastInstalledExtensionVersion) ??
+        semver.parse(storage.get("databricks.lastInstalledExtensionVersion")) ??
         new semver.SemVer("0.0.0");
 
     // if the extension is downgraded, we do not want to show the popup
diff --git a/packages/databricks-vscode/src/workspace-fs/WorkspaceFsAccessVerifier.ts b/packages/databricks-vscode/src/workspace-fs/WorkspaceFsAccessVerifier.ts
@@ -43,7 +43,7 @@ export async function switchToWorkspacePrompt(
     });
 
     if (selection === "Don't show again") {
-        stateStorage.skipSwitchToWorkspace = true;
+        stateStorage.set("databricks.wsfs.skipSwitchToWorkspace", true);
         return;
     }
 
@@ -142,7 +142,7 @@ export class WorkspaceFsAccessVerifier implements Disposable {
             if (
                 workspaceConfigs.enableFilesInWorkspace ||
                 !(await this.isEnabledForWorkspace()) ||
-                this.stateStorage.skipSwitchToWorkspace
+                this.stateStorage.get("databricks.wsfs.skipSwitchToWorkspace")
             ) {
                 return;
             }
diff --git a/packages/databricks-vscode/src/workspace-fs/WorkspaceFsCommands.ts b/packages/databricks-vscode/src/workspace-fs/WorkspaceFsCommands.ts
@@ -44,7 +44,9 @@ export class WorkspaceFsCommands implements Disposable {
                 const element = await root?.mkdir(
                     `${path.basename(
                         this.workspaceFolder.fsPath
-                    )}-${this.stateStorage.fixedUUID.slice(0, 8)}`
+                    )}-${this.stateStorage
+                        .get("databricks.fixedUUID")
+                        .slice(0, 8)}`
                 );
                 if (element) {
                     await this.connectionManager.attachSyncDestination(

Original file line number	Diff line number	Diff line change
`@@ -106,12 +106,15 @@ export class ConfigureAutocomplete implements Disposable {`
`106`	`106`	`}`
`107`	`107`
`108`	`108`	`async configureCommand() {`
`109`		`- this.stateStorage.skipAutocompleteConfigure = false;`
	`109`	`+ this.stateStorage.set("databricks.autocompletion.skipConfigure", false);`
`110`	`110`	`return this.configure(true);`
`111`	`111`	`}`
`112`	`112`
`113`	`113`	`private async configure(force = false) {`
`114`		`- if (!force \|\| this.stateStorage.skipAutocompleteConfigure) {`
	`114`	`+ if (`
	`115`	`+ !force \|\|`
	`116`	`+ this.stateStorage.get("databricks.autocompletion.skipConfigure")`
	`117`	`+ ) {`
`115`	`118`	`return;`
`116`	`119`	`}`
`117`	`120`
`@@ -142,7 +145,10 @@ export class ConfigureAutocomplete implements Disposable {`
`142`	`145`	`}`
`143`	`146`
`144`	`147`	`if (choice === "Never for this workspace") {`
`145`		`- this.stateStorage.skipAutocompleteConfigure = true;`
	`148`	`+ this.stateStorage.set(`
	`149`	`+ "databricks.autocompletion.skipConfigure",`
	`150`	`+ true`
	`151`	`+ );`
`146`	`152`	`return;`
`147`	`153`	`}`
`148`	`154`
Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ export async function showWhatsNewPopup(`
`44`	`44`	`}`
`45`	`45`
`46`	`46`	`const previousVersion =`
`47`		`- semver.parse(storage.lastInstalledExtensionVersion) ??`
	`47`	`+ semver.parse(storage.get("databricks.lastInstalledExtensionVersion")) ??`
`48`	`48`	`new semver.SemVer("0.0.0");`
`49`	`49`
`50`	`50`	`// if the extension is downgraded, we do not want to show the popup`