Skip to content

Commit 6b2caca

Browse files
kartikgupta-dbkartikgupta-db
authored
Use bundle validate to load interpolated view of configs after login (#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? -->
1 parent 8ce4c9a commit 6b2caca

File tree

11 files changed

+421
-44
lines changed

11 files changed

+421
-44
lines changed

‎packages/databricks-vscode/src/bundle/models/BundlePreValidateModel.ts‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
import {Disposable, Uri} from "vscode";
2+
import {BundleFileSet, BundleWatcher} from "..";
3+
import {BundleTarget} from "../types";
4+
import {CachedValue} from "../../locking/CachedValue";
5+
import {
6+
BundlePreValidateConfig,
7+
isBundlePreValidateConfigKey,
8+
} from "../../configuration/types";
9+
/**
10+
* Reads and writes bundle configs. This class does not notify when the configs change.
11+
* We use the BundleWatcher to notify when the configs change.
12+
*/
13+
export class BundlePreValidateModel implements Disposable {
14+
private disposables: Disposable[] = [];
15+
16+
private readonly stateCache = new CachedValue<
17+
BundlePreValidateConfig | undefined
18+
>(async () => {
19+
if (this.target === undefined) {
20+
return undefined;
21+
}
22+
return this.readState(this.target);
23+
});
24+
25+
public readonly onDidChange = this.stateCache.onDidChange;
26+
27+
private target: string | undefined;
28+
29+
private readonly readerMapping: Record<
30+
keyof BundlePreValidateConfig,
31+
(
32+
t?: BundleTarget
33+
) => Promise<
34+
BundlePreValidateConfig[keyof BundlePreValidateConfig] | undefined
35+
>
36+
> = {
37+
authParams: this.getAuthParams,
38+
mode: this.getMode,
39+
host: this.getHost,
40+
};
41+
42+
constructor(
43+
private readonly bundleFileSet: BundleFileSet,
44+
private readonly bunldeFileWatcher: BundleWatcher
45+
) {
46+
this.disposables.push(
47+
this.bunldeFileWatcher.onDidChange(async () => {
48+
await this.stateCache.refresh();
49+
})
50+
);
51+
}
52+
53+
private async getHost(target?: BundleTarget) {
54+
return target?.workspace?.host;
55+
}
56+
57+
private async getMode(target?: BundleTarget) {
58+
return target?.mode;
59+
}
60+
61+
/* eslint-disable @typescript-eslint/no-unused-vars */
62+
private async getAuthParams(target?: BundleTarget) {
63+
return undefined;
64+
}
65+
/* eslint-enable @typescript-eslint/no-unused-vars */
66+
67+
get targets() {
68+
return this.bundleFileSet.bundleDataCache.value.then(
69+
(data) => data?.targets
70+
);
71+
}
72+
73+
get defaultTarget() {
74+
return this.targets.then((targets) => {
75+
if (targets === undefined) {
76+
return undefined;
77+
}
78+
const defaultTarget = Object.keys(targets).find(
79+
(target) => targets[target].default
80+
);
81+
return defaultTarget;
82+
});
83+
}
84+
85+
public async setTarget(target: string | undefined) {
86+
this.target = target;
87+
await this.stateCache.refresh();
88+
}
89+
90+
private async readState(target: string) {
91+
const configs = {} as any;
92+
const targetObject = (await this.bundleFileSet.bundleDataCache.value)
93+
.targets?.[target];
94+
95+
for (const key of Object.keys(this.readerMapping)) {
96+
if (!isBundlePreValidateConfigKey(key)) {
97+
continue;
98+
}
99+
configs[key] = await this.readerMapping[key](targetObject);
100+
}
101+
return configs as BundlePreValidateConfig;
102+
}
103+
104+
public async getFileToWrite<T extends keyof BundlePreValidateConfig>(
105+
key: T
106+
) {
107+
const filesWithTarget: Uri[] = [];
108+
const filesWithConfig = (
109+
await this.bundleFileSet.findFile(async (data, file) => {
110+
const bundleTarget = data.targets?.[this.target ?? ""];
111+
if (bundleTarget) {
112+
filesWithTarget.push(file);
113+
}
114+
if (
115+
(await this.readerMapping[key](bundleTarget)) === undefined
116+
) {
117+
return false;
118+
}
119+
return true;
120+
})
121+
).map((file) => file.file);
122+
123+
if (filesWithConfig.length > 1) {
124+
throw new Error(
125+
`Multiple files found to write the config ${key} for target ${this.target}`
126+
);
127+
}
128+
129+
if (filesWithConfig.length === 0 && filesWithTarget.length === 0) {
130+
throw new Error(
131+
`No files found to write the config ${key} for target ${this.target}`
132+
);
133+
}
134+
135+
return [...filesWithConfig, ...filesWithTarget][0];
136+
}
137+
138+
public async load() {
139+
return await this.stateCache.value;
140+
}
141+
142+
public dispose() {
143+
this.disposables.forEach((d) => d.dispose());
144+
}
145+
}

‎packages/databricks-vscode/src/bundle/models/BundleValidateModel.ts‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
import {Disposable, Uri, EventEmitter} from "vscode";
2+
import {BundleWatcher} from "../BundleWatcher";
3+
import {AuthProvider} from "../../configuration/auth/AuthProvider";
4+
import {Mutex} from "../../locking";
5+
import {CliWrapper} from "../../cli/CliWrapper";
6+
import {BundleTarget} from "../types";
7+
import {CachedValue} from "../../locking/CachedValue";
8+
import {onError} from "../../utils/onErrorDecorator";
9+
import lodash from "lodash";
10+
import {workspaceConfigs} from "../../vscode-objs/WorkspaceConfigs";
11+
import {BundleValidateConfig} from "../../configuration/types";
12+
13+
type BundleValidateState = BundleValidateConfig & BundleTarget;
14+
15+
export class BundleValidateModel implements Disposable {
16+
private disposables: Disposable[] = [];
17+
private mutex = new Mutex();
18+
19+
private target: string | undefined;
20+
private authProvider: AuthProvider | undefined;
21+
22+
private readonly stateCache = new CachedValue<
23+
BundleValidateState | undefined
24+
>(this.readState.bind(this));
25+
26+
private readonly onDidChangeKeyEmitters = new Map<
27+
keyof BundleValidateState,
28+
EventEmitter<void>
29+
>();
30+
31+
onDidChangeKey(key: keyof BundleValidateState) {
32+
if (!this.onDidChangeKeyEmitters.has(key)) {
33+
this.onDidChangeKeyEmitters.set(key, new EventEmitter());
34+
}
35+
return this.onDidChangeKeyEmitters.get(key)!.event;
36+
}
37+
38+
public onDidChange = this.stateCache.onDidChange;
39+
40+
constructor(
41+
private readonly bundleWatcher: BundleWatcher,
42+
private readonly cli: CliWrapper,
43+
private readonly workspaceFolder: Uri
44+
) {
45+
this.disposables.push(
46+
this.bundleWatcher.onDidChange(async () => {
47+
await this.stateCache.refresh();
48+
}),
49+
// Emit an event for each key that changes
50+
this.stateCache.onDidChange(async ({oldValue, newValue}) => {
51+
for (const key of Object.keys({
52+
...oldValue,
53+
...newValue,
54+
}) as (keyof BundleValidateState)[]) {
55+
if (
56+
oldValue === null ||
57+
!lodash.isEqual(oldValue?.[key], newValue?.[key])
58+
) {
59+
this.onDidChangeKeyEmitters.get(key)?.fire();
60+
}
61+
}
62+
})
63+
);
64+
}
65+
66+
private readerMapping: {
67+
[K in keyof BundleValidateState]: (
68+
t?: BundleTarget
69+
) => BundleValidateState[K];
70+
} = {
71+
clusterId: (target) => target?.bundle?.compute_id,
72+
workspaceFsPath: (target) => target?.workspace?.file_path,
73+
resources: (target) => target?.resources,
74+
};
75+
76+
@Mutex.synchronise("mutex")
77+
public async setTarget(target: string | undefined) {
78+
if (this.target === target) {
79+
return;
80+
}
81+
this.target = target;
82+
this.authProvider = undefined;
83+
await this.stateCache.refresh();
84+
}
85+
86+
@Mutex.synchronise("mutex")
87+
public async setAuthProvider(authProvider: AuthProvider | undefined) {
88+
if (
89+
!lodash.isEqual(this.authProvider?.toJSON(), authProvider?.toJSON())
90+
) {
91+
this.authProvider = authProvider;
92+
await this.stateCache.refresh();
93+
}
94+
}
95+
96+
@onError({popup: {prefix: "Failed to read bundle config."}})
97+
@Mutex.synchronise("mutex")
98+
private async readState() {
99+
if (this.target === undefined || this.authProvider === undefined) {
100+
return;
101+
}
102+
103+
const targetObject = JSON.parse(
104+
await this.cli.bundleValidate(
105+
this.target,
106+
this.authProvider,
107+
this.workspaceFolder,
108+
workspaceConfigs.databrickscfgLocation
109+
)
110+
) as BundleTarget;
111+
112+
const configs: any = {};
113+
114+
for (const key of Object.keys(
115+
this.readerMapping
116+
) as (keyof BundleValidateState)[]) {
117+
configs[key] = this.readerMapping[key]?.(targetObject);
118+
}
119+
120+
return {...configs, ...targetObject} as BundleValidateState;
121+
}
122+
123+
@Mutex.synchronise("mutex")
124+
public async load<T extends keyof BundleValidateState>(
125+
keys: T[] = []
126+
): Promise<Partial<Pick<BundleValidateState, T>> | undefined> {
127+
if (keys.length === 0) {
128+
return await this.stateCache.value;
129+
}
130+
131+
const target = await this.stateCache.value;
132+
const configs: Partial<{
133+
[K in T]: BundleValidateState[K];
134+
}> = {};
135+
136+
for (const key of keys) {
137+
configs[key] = this.readerMapping[key]?.(target);
138+
}
139+
140+
return configs;
141+
}
142+
143+
dispose() {
144+
this.disposables.forEach((i) => i.dispose());
145+
}
146+
}

‎packages/databricks-vscode/src/cli/CliWrapper.ts‎

Lines changed: 31 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
import {execFile as execFileCb, spawn} from "child_process";
2-
import {ExtensionContext, window, commands} from "vscode";
2+
import {ExtensionContext, window, commands, Uri} from "vscode";
33
import {SyncDestinationMapper} from "../sync/SyncDestination";
44
import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
55
import {promisify} from "node:util";
66
import {logging} from "@databricks/databricks-sdk";
77
import {Loggers} from "../logger";
88
import {Context, context} from "@databricks/databricks-sdk/dist/context";
99
import {Cloud} from "../utils/constants";
10+
import {AuthProvider} from "../configuration/auth/AuthProvider";
11+
import {EnvVarGenerators} from "../utils";
1012

1113
const withLogContext = logging.withLogContext;
1214
const execFile = promisify(execFileCb);
@@ -86,12 +88,8 @@ export class CliWrapper {
8688
const cmd = await this.getListProfilesCommand();
8789
const res = await execFile(cmd.command, cmd.args, {
8890
env: {
89-
/* eslint-disable @typescript-eslint/naming-convention */
90-
HOME: process.env.HOME,
91-
DATABRICKS_CONFIG_FILE:
92-
configfilePath || process.env.DATABRICKS_CONFIG_FILE,
93-
DATABRICKS_OUTPUT_FORMAT: "json",
94-
/* eslint-enable @typescript-eslint/naming-convention */
91+
...EnvVarGenerators.getEnvVarsForCli(configfilePath),
92+
...EnvVarGenerators.getProxyEnvVars(),
9593
},
9694
});
9795
const profiles = JSON.parse(res.stdout).profiles || [];
@@ -134,7 +132,6 @@ export class CliWrapper {
134132
}
135133

136134
public async getBundleSchema(): Promise<string> {
137-
const execFile = promisify(execFileCb);
138135
const {stdout} = await execFile(this.cliPath, ["bundle", "schema"]);
139136
return stdout;
140137
}
@@ -172,4 +169,30 @@ export class CliWrapper {
172169
child.on("exit", resolve);
173170
});
174171
}
172+
173+
async bundleValidate(
174+
target: string,
175+
authProvider: AuthProvider,
176+
workspaceFolder: Uri,
177+
configfilePath?: string
178+
) {
179+
const {stdout, stderr} = await execFile(
180+
this.cliPath,
181+
["bundle", "validate", "--target", target],
182+
{
183+
cwd: workspaceFolder.fsPath,
184+
env: {
185+
...EnvVarGenerators.getEnvVarsForCli(configfilePath),
186+
...EnvVarGenerators.getProxyEnvVars(),
187+
...authProvider.toEnv(),
188+
},
189+
shell: true,
190+
}
191+
);
192+
193+
if (stderr !== "") {
194+
throw new Error(stderr);
195+
}
196+
return stdout;
197+
}
175198
}

‎packages/databricks-vscode/src/configuration/ConnectionCommands.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ export class ConnectionCommands implements Disposable {
200200
}
201201

202202
async selectTarget() {
203-
const targets = await this.configModel.bundleFileConfigModel.targets;
203+
const targets = await this.configModel.bundlePreValidateModel.targets;
204204
const currentTarget = this.configModel.target;
205205
if (targets === undefined) {
206206
return;

‎packages/databricks-vscode/src/configuration/ConnectionManager.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ export class ConnectionManager implements Disposable {
225225
await this.updateSyncDestinationMapper();
226226
await this.updateClusterManager();
227227
await this.configModel.set("authParams", authProvider.toJSON());
228-
228+
await this.configModel.setAuthProvider(authProvider);
229229
this.updateState("CONNECTED");
230230
});
231231
} catch (e) {

0 commit comments

Comments
 (0)