Add config level change events in ConfigModel (#937)

kartikgupta-db · web-flow · commit 84e21980d8d8 · 2023-11-14T11:03:05.000Z
## Changes
* Add config level change events in config model. Whenever a config
changes, the corresponding event is fired.

* Also, allow writing auth parameters as an override.

## Tests
&lt;!-- How is this tested? --&gt;
diff --git a/packages/databricks-vscode/src/bundle/BundleWatcher.ts b/packages/databricks-vscode/src/bundle/BundleWatcher.ts
@@ -1,5 +1,5 @@
 import {Disposable, EventEmitter, Uri, workspace} from "vscode";
-import {BundleFileSet} from "../bundle/BundleFileSet";
+import {BundleFileSet} from "./BundleFileSet";
 import {WithMutex} from "../locking";
 import path from "path";
 
diff --git a/packages/databricks-vscode/src/bundle/bundleAutocompleteProvider.ts b/packages/databricks-vscode/src/bundle/bundleAutocompleteProvider.ts
@@ -1,7 +1,7 @@
 import {CliWrapper} from "../cli/CliWrapper";
 import {ExtensionContext, extensions, Uri} from "vscode";
 import {BundleFileSet} from "./BundleFileSet";
-import {BundleWatcher} from "../file-managers/BundleWatcher";
+import {BundleWatcher} from "./BundleWatcher";
 
 export async function registerBundleAutocompleteProvider(
     cli: CliWrapper,
diff --git a/packages/databricks-vscode/src/bundle/index.ts b/packages/databricks-vscode/src/bundle/index.ts
@@ -0,0 +1,3 @@
+export {BundleWatcher} from "./BundleWatcher";
+export {BundleFileSet, parseBundleYaml, writeBundleYaml} from "./BundleFileSet";
+export {registerBundleAutocompleteProvider} from "./bundleAutocompleteProvider";
diff --git a/packages/databricks-vscode/src/configuration/BundleConfigReaderWriter.ts b/packages/databricks-vscode/src/configuration/BundleConfigReaderWriter.ts
@@ -1,23 +1,24 @@
 import {Uri} from "vscode";
-import {
-    BundleFileSet,
-    parseBundleYaml,
-    writeBundleYaml,
-} from "../bundle/BundleFileSet";
+import {BundleFileSet, parseBundleYaml, writeBundleYaml} from "../bundle";
 import {BundleTarget} from "../bundle/types";
 import {Mutex} from "../locking";
-import {RemoteUri} from "../sync/SyncDestination";
-import {BundleConfigs, isBundleConfig} from "./types";
-
-export class BundleConfigReaderWriter {
+import {BundleConfigs, ConfigReaderWriter, isBundleConfig} from "./types";
+
+/**
+ * Reads and writes bundle configs. This class does not notify when the configs change.
+ * We use the BundleWatcher to notify when the configs change.
+ */
+export class BundleConfigReaderWriter
+    implements ConfigReaderWriter<keyof BundleConfigs>
+{
     private readonly writeMutex = new Mutex();
 
     private readonly writerMapping: Record<
         keyof BundleConfigs,
         (t: BundleTarget, v: any) => BundleTarget
     > = {
         clusterId: this.setClusterId,
-        authType: this.setAuthType,
+        authParams: this.setAuthParams,
         mode: this.setMode,
         host: this.setHost,
         workspaceFsPath: this.setWorkspaceFsPath,
@@ -30,7 +31,7 @@ export class BundleConfigReaderWriter {
         ) => Promise<BundleConfigs[keyof BundleConfigs] | undefined>
     > = {
         clusterId: this.getClusterId,
-        authType: this.getAuthType,
+        authParams: this.getAuthParams,
         mode: this.getMode,
         host: this.getHost,
         workspaceFsPath: this.getWorkspaceFsPath,
@@ -68,9 +69,10 @@ export class BundleConfigReaderWriter {
         return target;
     }
 
-    public async getWorkspaceFsPath(target?: BundleTarget) {
-        const filePath = target?.workspace?.file_path;
-        return filePath ? new RemoteUri(filePath) : undefined;
+    public async getWorkspaceFsPath(
+        target?: BundleTarget
+    ): Promise<BundleConfigs["workspaceFsPath"]> {
+        return target?.workspace?.file_path;
     }
     public setWorkspaceFsPath(
         target: BundleTarget,
@@ -80,23 +82,23 @@ export class BundleConfigReaderWriter {
         target.workspace = {
             ...target.workspace,
             // eslint-disable-next-line @typescript-eslint/naming-convention
-            file_path: value?.path,
+            file_path: value,
         };
         return target;
     }
 
-    public async getAuthType(target?: BundleTarget) {
-        return target?.workspace?.auth_type;
+    /* eslint-disable @typescript-eslint/no-unused-vars */
+
+    public async getAuthParams(target?: BundleTarget) {
+        return undefined;
     }
-    public setAuthType(target: BundleTarget, value: BundleConfigs["authType"]) {
-        target = {...target};
-        target.workspace = {
-            ...target.workspace,
-            // eslint-disable-next-line @typescript-eslint/naming-convention
-            auth_type: value,
-        };
-        return target;
+    public setAuthParams(
+        target: BundleTarget,
+        value: BundleConfigs["authParams"]
+    ): BundleTarget {
+        throw new Error("Not implemented");
     }
+    /* eslint-enable @typescript-eslint/no-unused-vars */
 
     get targets() {
         return this.bundleFileSet.bundleDataCache.value.then(
@@ -168,19 +170,22 @@ export class BundleConfigReaderWriter {
     ) {
         const file = await this.getFileToWrite(key, target);
         if (file === undefined) {
-            return false;
+            throw new Error(
+                `Can't find a file to write property '${key}' of target '${target}'.`
+            );
         }
         const data = await parseBundleYaml(file);
         const targetData = data.targets?.[target];
         if (targetData === undefined) {
-            return false;
+            throw new Error(`No target '${target}' for writing '${key}.`);
         }
 
         const newTargetData = this.writerMapping[key](targetData, value);
+        if (JSON.stringify(newTargetData) === JSON.stringify(targetData)) {
+            return;
+        }
         data.targets = {...data.targets, [target]: newTargetData};
         await writeBundleYaml(file, data);
-
-        return true;
     }
 
     /**
diff --git a/packages/databricks-vscode/src/configuration/ConfigModel.ts b/packages/databricks-vscode/src/configuration/ConfigModel.ts
@@ -1,4 +1,4 @@
-import {Disposable, EventEmitter, Uri} from "vscode";
+import {Disposable, EventEmitter, Uri, Event} from "vscode";
 import {
     BundleConfigs,
     DatabricksConfigs,
@@ -8,7 +8,7 @@ import {
 import {ConfigOverrideReaderWriter} from "./ConfigOverrideReaderWriter";
 import {BundleConfigReaderWriter} from "./BundleConfigReaderWriter";
 import {Mutex} from "../locking";
-import {BundleWatcher} from "../file-managers/BundleWatcher";
+import {BundleWatcher} from "../bundle";
 import {CachedValue} from "../locking/CachedValue";
 import {StateStorage} from "../vscode-objs/StateStorage";
 
@@ -38,22 +38,43 @@ export class ConfigModel implements Disposable {
             if (this.target === undefined) {
                 return {};
             }
-            const overrides = this.overrideReaderWriter.readAll(this.target);
+            const overrides = await this.overrideReaderWriter.readAll(
+                this.target
+            );
             const bundleConfigs = await this.bundleConfigReaderWriter.readAll(
                 this.target
             );
-            const newValue = {...bundleConfigs, ...overrides};
-
-            if (JSON.stringify(oldValue) !== JSON.stringify(newValue)) {
-                this.onDidChangeEmitter.fire();
+            const newValue: DatabricksConfigs = {
+                ...bundleConfigs,
+                ...overrides,
+            };
+
+            for (const key of <(keyof DatabricksConfigs)[]>(
+                Object.keys(newValue)
+            )) {
+                if (
+                    oldValue === null ||
+                    JSON.stringify(oldValue[key]) !==
+                        JSON.stringify(newValue[key])
+                ) {
+                    this.changeEmitters.get(key)?.emitter.fire();
+                    this.onDidChangeAnyEmitter.fire();
+                }
             }
 
             return newValue;
         }
     );
 
-    private readonly onDidChangeEmitter = new EventEmitter<void>();
-    public readonly onDidChange = this.onDidChangeEmitter.event;
+    private readonly changeEmitters = new Map<
+        keyof DatabricksConfigs | "target",
+        {
+            emitter: EventEmitter<void>;
+            onDidEmit: Event<void>;
+        }
+    >();
+    private onDidChangeAnyEmitter = new EventEmitter<void>();
+    public onDidChangeAny = this.onDidChangeAnyEmitter.event;
 
     private _target: string | undefined;
 
@@ -77,7 +98,31 @@ export class ConfigModel implements Disposable {
             })
         );
     }
+    public async init() {
+        await this.readTarget();
+    }
+
+    public onDidChange<T extends keyof DatabricksConfigs | "target">(
+        key: T,
+        fn: () => any,
+        thisArgs?: any
+    ) {
+        if (!this.changeEmitters.has(key)) {
+            const emitter = new EventEmitter<void>();
+            this.changeEmitters.set(key, {
+                emitter: emitter,
+                onDidEmit: emitter.event,
+            });
+        }
 
+        const {onDidEmit} = this.changeEmitters.get(key)!;
+        return onDidEmit(fn, thisArgs);
+    }
+    /**
+     * Try to read target from bundle config.
+     * If not found, try to read from state storage.
+     * If not found, try to read the default target from bundle.
+     */
     public async readTarget() {
         const targets = Object.keys(
             (await this.bundleConfigReaderWriter.targets) ?? {}
@@ -102,6 +147,9 @@ export class ConfigModel implements Disposable {
         return this._target;
     }
 
+    /**
+     * Set target in the state storage and invalidate the configs cache.
+     */
     public async setTarget(target: string | undefined) {
         if (target === this._target) {
             return;
@@ -112,7 +160,8 @@ export class ConfigModel implements Disposable {
             await this.stateStorage.set("databricks.bundle.target", target);
         });
         await this.configCache.invalidate();
-        this.onDidChangeEmitter.fire();
+        this.changeEmitters.get("target")?.emitter.fire();
+        this.onDidChangeAnyEmitter.fire();
     }
 
     public async get<T extends keyof DatabricksConfigs>(
@@ -126,19 +175,22 @@ export class ConfigModel implements Disposable {
         key: T,
         value?: DatabricksConfigs[T],
         handleInteractiveWrite?: (file: Uri | undefined) => any
-    ): Promise<boolean> {
+    ): Promise<void> {
         // We work with 1 set of configs throughout the function.
         // No changes to the cache can happen when the global mutex is held.
         // The assumption is that user doesn't change the target mode in the middle of
         // writing a new config.
         const {mode} = {...(await this.configCache.value)};
 
         if (this.target === undefined) {
-            return false;
+            throw new Error(
+                `Can't set configuration '${key}' without selecting a target`
+            );
         }
         if (isOverrideableConfig(key)) {
             return this.overrideReaderWriter.write(key, this.target, value);
-        } else if (isBundleConfig(key)) {
+        }
+        if (isBundleConfig(key)) {
             const isInteractive = handleInteractiveWrite !== undefined;
 
             // write to bundle if not interactive and the config can be safely written to bundle
@@ -158,7 +210,6 @@ export class ConfigModel implements Disposable {
                 handleInteractiveWrite(file);
             }
         }
-        return true;
     }
 
     dispose() {
diff --git a/packages/databricks-vscode/src/configuration/ConfigOverrideReaderWriter.ts b/packages/databricks-vscode/src/configuration/ConfigOverrideReaderWriter.ts
@@ -1,9 +1,11 @@
 import {EventEmitter} from "vscode";
 import {Mutex} from "../locking";
 import {StateStorage} from "../vscode-objs/StateStorage";
-import {OverrideableConfigs} from "./types";
+import {OverrideableConfigs, ConfigReaderWriter} from "./types";
 
-export class ConfigOverrideReaderWriter {
+export class ConfigOverrideReaderWriter
+    implements ConfigReaderWriter<keyof OverrideableConfigs>
+{
     private writeMutex = new Mutex();
     private onDidChangeEmitter = new EventEmitter<void>();
     public readonly onDidChange = this.onDidChangeEmitter.event;
@@ -24,15 +26,16 @@ export class ConfigOverrideReaderWriter {
         value?: OverrideableConfigs[T]
     ) {
         const data = this.storage.get("databricks.bundle.overrides");
-        if (data[target]?.[key] !== undefined) {
-            if (data[target] === undefined) {
-                data[target] = {};
-            }
-            data[target][key] = value;
-            await this.storage.set("databricks.bundle.overrides", data);
-            this.onDidChangeEmitter.fire();
+        if (data[target] === undefined) {
+            data[target] = {};
         }
-        return true;
+        const oldValue = JSON.stringify(data[target][key]);
+        if (oldValue === JSON.stringify(value)) {
+            return;
+        }
+        data[target][key] = value;
+        await this.storage.set("databricks.bundle.overrides", data);
+        this.onDidChangeEmitter.fire();
     }
 
     /**
diff --git a/packages/databricks-vscode/src/configuration/types.ts b/packages/databricks-vscode/src/configuration/types.ts
@@ -1,19 +1,17 @@
-import {RemoteUri} from "../sync/SyncDestination";
-
 export type DatabricksConfigs = {
     host?: string;
 
     // reconcile with actual mode and auth type enums from bundle
     mode?: "dev" | "staging" | "prod";
-    authType?: string;
+    authParams?: Record<string, any>;
 
     clusterId?: string;
-    workspaceFsPath?: RemoteUri;
+    workspaceFsPath?: string;
 };
 
 export const OVERRIDEABLE_CONFIGS = [
     "clusterId",
-    "authType",
+    "authParams",
     "workspaceFsPath",
 ] as const;
 
@@ -24,7 +22,7 @@ export type OverrideableConfigs = Pick<
 
 export const BUNDLE_CONFIGS = [
     "clusterId",
-    "authType",
+    "authParams",
     "workspaceFsPath",
     "mode",
     "host",
@@ -45,3 +43,8 @@ export function isOverrideableConfig(
 export function isBundleConfig(key: any): key is keyof BundleConfigs {
     return BUNDLE_CONFIGS.includes(key);
 }
+
+export interface ConfigReaderWriter<T extends keyof DatabricksConfigs> {
+    read(key: T, target: string): Promise<DatabricksConfigs[T] | undefined>;
+    write(key: T, target: string, value?: DatabricksConfigs[T]): Promise<void>;
+}
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
@@ -32,7 +32,6 @@ import {
     WorkspaceFsCommands,
     WorkspaceFsDataProvider,
 } from "./workspace-fs";
-import {registerBundleAutocompleteProvider} from "./bundle/bundleAutocompleteProvider";
 import {CustomWhenContext} from "./vscode-objs/CustomWhenContext";
 import {StateStorage} from "./vscode-objs/StateStorage";
 import path from "node:path";
@@ -50,8 +49,11 @@ import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton";
 import {NotebookAccessVerifier} from "./language/notebooks/NotebookAccessVerifier";
 import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager";
 import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue";
-import {BundleWatcher} from "./file-managers/BundleWatcher";
-import {BundleFileSet} from "./bundle/BundleFileSet";
+import {
+    BundleWatcher,
+    BundleFileSet,
+    registerBundleAutocompleteProvider,
+} from "./bundle";
 import {showWhatsNewPopup} from "./whatsNewPopup";
 
 export async function activate(

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+export {BundleWatcher} from "./BundleWatcher";`
	`2`	`+export {BundleFileSet, parseBundleYaml, writeBundleYaml} from "./BundleFileSet";`
	`3`	`+export {registerBundleAutocompleteProvider} from "./bundleAutocompleteProvider";`