Skip to content

Commit

Permalink
Use bundle validate to load interpolated view of configs after login (
Browse files Browse the repository at this point in the history
#979)

## Changes
* Add a `ConfigModel` for loading configs - using `bundle validate` -
**after** auth is completed. This also interpolates bundle variables and
params. Right now we use it for workspace path and clusterId. In the
future, we would also want to use it to pull a pre deploy view of the
resources.

**Note**: We still need more discussions on how to handle params. We can
have a separate PR for that.

## Tests
<!-- How is this tested? -->
  • Loading branch information
kartikgupta-db authored Jan 9, 2024
1 parent 8ce4c9a commit 6b2caca
Show file tree
Hide file tree
Showing 11 changed files with 421 additions and 44 deletions.
145 changes: 145 additions & 0 deletions packages/databricks-vscode/src/bundle/models/BundlePreValidateModel.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,145 @@
import {Disposable, Uri} from "vscode";
import {BundleFileSet, BundleWatcher} from "..";
import {BundleTarget} from "../types";
import {CachedValue} from "../../locking/CachedValue";
import {
BundlePreValidateConfig,
isBundlePreValidateConfigKey,
} from "../../configuration/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 BundlePreValidateModel implements Disposable {
private disposables: Disposable[] = [];

private readonly stateCache = new CachedValue<
BundlePreValidateConfig | undefined
>(async () => {
if (this.target === undefined) {
return undefined;
}
return this.readState(this.target);
});

public readonly onDidChange = this.stateCache.onDidChange;

private target: string | undefined;

private readonly readerMapping: Record<
keyof BundlePreValidateConfig,
(
t?: BundleTarget
) => Promise<
BundlePreValidateConfig[keyof BundlePreValidateConfig] | undefined
>
> = {
authParams: this.getAuthParams,
mode: this.getMode,
host: this.getHost,
};

constructor(
private readonly bundleFileSet: BundleFileSet,
private readonly bunldeFileWatcher: BundleWatcher
) {
this.disposables.push(
this.bunldeFileWatcher.onDidChange(async () => {
await this.stateCache.refresh();
})
);
}

private async getHost(target?: BundleTarget) {
return target?.workspace?.host;
}

private async getMode(target?: BundleTarget) {
return target?.mode;
}

/* eslint-disable @typescript-eslint/no-unused-vars */
private async getAuthParams(target?: BundleTarget) {
return undefined;
}
/* eslint-enable @typescript-eslint/no-unused-vars */

get targets() {
return this.bundleFileSet.bundleDataCache.value.then(
(data) => data?.targets
);
}

get defaultTarget() {
return this.targets.then((targets) => {
if (targets === undefined) {
return undefined;
}
const defaultTarget = Object.keys(targets).find(
(target) => targets[target].default
);
return defaultTarget;
});
}

public async setTarget(target: string | undefined) {
this.target = target;
await this.stateCache.refresh();
}

private async readState(target: string) {
const configs = {} as any;
const targetObject = (await this.bundleFileSet.bundleDataCache.value)
.targets?.[target];

for (const key of Object.keys(this.readerMapping)) {
if (!isBundlePreValidateConfigKey(key)) {
continue;
}
configs[key] = await this.readerMapping[key](targetObject);
}
return configs as BundlePreValidateConfig;
}

public async getFileToWrite<T extends keyof BundlePreValidateConfig>(
key: T
) {
const filesWithTarget: Uri[] = [];
const filesWithConfig = (
await this.bundleFileSet.findFile(async (data, file) => {
const bundleTarget = data.targets?.[this.target ?? ""];
if (bundleTarget) {
filesWithTarget.push(file);
}
if (
(await this.readerMapping[key](bundleTarget)) === undefined
) {
return false;
}
return true;
})
).map((file) => file.file);

if (filesWithConfig.length > 1) {
throw new Error(
`Multiple files found to write the config ${key} for target ${this.target}`
);
}

if (filesWithConfig.length === 0 && filesWithTarget.length === 0) {
throw new Error(
`No files found to write the config ${key} for target ${this.target}`
);
}

return [...filesWithConfig, ...filesWithTarget][0];
}

public async load() {
return await this.stateCache.value;
}

public dispose() {
this.disposables.forEach((d) => d.dispose());
}
}
146 changes: 146 additions & 0 deletions packages/databricks-vscode/src/bundle/models/BundleValidateModel.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
import {Disposable, Uri, EventEmitter} from "vscode";
import {BundleWatcher} from "../BundleWatcher";
import {AuthProvider} from "../../configuration/auth/AuthProvider";
import {Mutex} from "../../locking";
import {CliWrapper} from "../../cli/CliWrapper";
import {BundleTarget} from "../types";
import {CachedValue} from "../../locking/CachedValue";
import {onError} from "../../utils/onErrorDecorator";
import lodash from "lodash";
import {workspaceConfigs} from "../../vscode-objs/WorkspaceConfigs";
import {BundleValidateConfig} from "../../configuration/types";

type BundleValidateState = BundleValidateConfig & BundleTarget;

export class BundleValidateModel implements Disposable {
private disposables: Disposable[] = [];
private mutex = new Mutex();

private target: string | undefined;
private authProvider: AuthProvider | undefined;

private readonly stateCache = new CachedValue<
BundleValidateState | undefined
>(this.readState.bind(this));

private readonly onDidChangeKeyEmitters = new Map<
keyof BundleValidateState,
EventEmitter<void>
>();

onDidChangeKey(key: keyof BundleValidateState) {
if (!this.onDidChangeKeyEmitters.has(key)) {
this.onDidChangeKeyEmitters.set(key, new EventEmitter());
}
return this.onDidChangeKeyEmitters.get(key)!.event;
}

public onDidChange = this.stateCache.onDidChange;

constructor(
private readonly bundleWatcher: BundleWatcher,
private readonly cli: CliWrapper,
private readonly workspaceFolder: Uri
) {
this.disposables.push(
this.bundleWatcher.onDidChange(async () => {
await this.stateCache.refresh();
}),
// Emit an event for each key that changes
this.stateCache.onDidChange(async ({oldValue, newValue}) => {
for (const key of Object.keys({
...oldValue,
...newValue,
}) as (keyof BundleValidateState)[]) {
if (
oldValue === null ||
!lodash.isEqual(oldValue?.[key], newValue?.[key])
) {
this.onDidChangeKeyEmitters.get(key)?.fire();
}
}
})
);
}

private readerMapping: {
[K in keyof BundleValidateState]: (
t?: BundleTarget
) => BundleValidateState[K];
} = {
clusterId: (target) => target?.bundle?.compute_id,
workspaceFsPath: (target) => target?.workspace?.file_path,
resources: (target) => target?.resources,
};

@Mutex.synchronise("mutex")
public async setTarget(target: string | undefined) {
if (this.target === target) {
return;
}
this.target = target;
this.authProvider = undefined;
await this.stateCache.refresh();
}

@Mutex.synchronise("mutex")
public async setAuthProvider(authProvider: AuthProvider | undefined) {
if (
!lodash.isEqual(this.authProvider?.toJSON(), authProvider?.toJSON())
) {
this.authProvider = authProvider;
await this.stateCache.refresh();
}
}

@onError({popup: {prefix: "Failed to read bundle config."}})
@Mutex.synchronise("mutex")
private async readState() {
if (this.target === undefined || this.authProvider === undefined) {
return;
}

const targetObject = JSON.parse(
await this.cli.bundleValidate(
this.target,
this.authProvider,
this.workspaceFolder,
workspaceConfigs.databrickscfgLocation
)
) as BundleTarget;

const configs: any = {};

for (const key of Object.keys(
this.readerMapping
) as (keyof BundleValidateState)[]) {
configs[key] = this.readerMapping[key]?.(targetObject);
}

return {...configs, ...targetObject} as BundleValidateState;
}

@Mutex.synchronise("mutex")
public async load<T extends keyof BundleValidateState>(
keys: T[] = []
): Promise<Partial<Pick<BundleValidateState, T>> | undefined> {
if (keys.length === 0) {
return await this.stateCache.value;
}

const target = await this.stateCache.value;
const configs: Partial<{
[K in T]: BundleValidateState[K];
}> = {};

for (const key of keys) {
configs[key] = this.readerMapping[key]?.(target);
}

return configs;
}

dispose() {
this.disposables.forEach((i) => i.dispose());
}
}
39 changes: 31 additions & 8 deletions packages/databricks-vscode/src/cli/CliWrapper.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
import {execFile as execFileCb, spawn} from "child_process";
import {ExtensionContext, window, commands} from "vscode";
import {ExtensionContext, window, commands, Uri} from "vscode";
import {SyncDestinationMapper} from "../sync/SyncDestination";
import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
import {promisify} from "node:util";
import {logging} from "@databricks/databricks-sdk";
import {Loggers} from "../logger";
import {Context, context} from "@databricks/databricks-sdk/dist/context";
import {Cloud} from "../utils/constants";
import {AuthProvider} from "../configuration/auth/AuthProvider";
import {EnvVarGenerators} from "../utils";

const withLogContext = logging.withLogContext;
const execFile = promisify(execFileCb);
Expand Down Expand Up @@ -86,12 +88,8 @@ export class CliWrapper {
const cmd = await this.getListProfilesCommand();
const res = await execFile(cmd.command, cmd.args, {
env: {
/* eslint-disable @typescript-eslint/naming-convention */
HOME: process.env.HOME,
DATABRICKS_CONFIG_FILE:
configfilePath || process.env.DATABRICKS_CONFIG_FILE,
DATABRICKS_OUTPUT_FORMAT: "json",
/* eslint-enable @typescript-eslint/naming-convention */
...EnvVarGenerators.getEnvVarsForCli(configfilePath),
...EnvVarGenerators.getProxyEnvVars(),
},
});
const profiles = JSON.parse(res.stdout).profiles || [];
Expand Down Expand Up @@ -134,7 +132,6 @@ export class CliWrapper {
}

public async getBundleSchema(): Promise<string> {
const execFile = promisify(execFileCb);
const {stdout} = await execFile(this.cliPath, ["bundle", "schema"]);
return stdout;
}
Expand Down Expand Up @@ -172,4 +169,30 @@ export class CliWrapper {
child.on("exit", resolve);
});
}

async bundleValidate(
target: string,
authProvider: AuthProvider,
workspaceFolder: Uri,
configfilePath?: string
) {
const {stdout, stderr} = await execFile(
this.cliPath,
["bundle", "validate", "--target", target],
{
cwd: workspaceFolder.fsPath,
env: {
...EnvVarGenerators.getEnvVarsForCli(configfilePath),
...EnvVarGenerators.getProxyEnvVars(),
...authProvider.toEnv(),
},
shell: true,
}
);

if (stderr !== "") {
throw new Error(stderr);
}
return stdout;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ export class ConnectionCommands implements Disposable {
}

async selectTarget() {
const targets = await this.configModel.bundleFileConfigModel.targets;
const targets = await this.configModel.bundlePreValidateModel.targets;
const currentTarget = this.configModel.target;
if (targets === undefined) {
return;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ export class ConnectionManager implements Disposable {
await this.updateSyncDestinationMapper();
await this.updateClusterManager();
await this.configModel.set("authParams", authProvider.toJSON());

await this.configModel.setAuthProvider(authProvider);
this.updateState("CONNECTED");
});
} catch (e) {
Expand Down
Loading

0 comments on commit 6b2caca

Please sign in to comment.