Skip to content

Commit 4a6aeae

Browse files
kartikgupta-dbkartikgupta-db
authored
Make cluster filtering opt-out (#185)
https://user-images.githubusercontent.com/88345179/200638478-eb27df60-9f39-4cf7-a884-cedbedbd549c.mov Also, add an additional permission and canExecute check when users attach a cluster. https://user-images.githubusercontent.com/88345179/200802145-cdfc7c4f-b0d0-40f3-811b-cd039c5c9460.mov
1 parent 9aef719 commit 4a6aeae

10 files changed

Lines changed: 307 additions & 124 deletions

File tree

‎packages/databricks-sdk-js/src/services/Cluster.ts‎

Lines changed: 77 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,19 +11,23 @@ import {
1111
import {CancellationToken} from "../types";
1212
import {ExecutionContext} from "./ExecutionContext";
1313
import {WorkflowRun} from "./WorkflowRun";
14-
import {commands} from "..";
14+
import {commands, PermissionsService} from "..";
1515
import {
1616
ClusterInfo,
1717
ClustersService,
1818
ClusterInfoState,
1919
ClusterInfoClusterSource,
2020
} from "../apis/clusters";
21-
import {Context} from "../context";
21+
import {Context, context} from "../context";
22+
import {User} from "../apis/scim";
23+
import {ExposedLoggers, withLogContext} from "../logging";
2224

2325
export class ClusterRetriableError extends RetriableError {}
2426
export class ClusterError extends Error {}
2527
export class Cluster {
2628
private clusterApi: ClustersService;
29+
private _canExecute?: boolean;
30+
private _hasExecutePerms?: boolean;
2731

2832
constructor(
2933
private client: ApiClient,
@@ -120,6 +124,60 @@ export class Cluster {
120124
this.clusterDetails = details;
121125
}
122126

127+
isSingleUser() {
128+
const modeProperty =
129+
//TODO: deprecate data_security_mode once access_mode is available everywhere
130+
this.details.access_mode ?? this.details.data_security_mode;
131+
return (
132+
modeProperty !== undefined &&
133+
[
134+
"SINGLE_USER",
135+
"LEGACY_SINGLE_USER_PASSTHROUGH",
136+
"LEGACY_SINGLE_USER_STANDARD",
137+
//enums unique to data_security_mode
138+
"LEGACY_SINGLE_USER",
139+
].includes(modeProperty)
140+
);
141+
}
142+
143+
isValidSingleUser(userName?: string) {
144+
return (
145+
this.isSingleUser() && this.details.single_user_name === userName
146+
);
147+
}
148+
149+
get hasExecutePermsCached() {
150+
return this._hasExecutePerms;
151+
}
152+
153+
async hasExecutePerms(userDetails?: User) {
154+
if (userDetails === undefined) {
155+
return (this._hasExecutePerms = false);
156+
}
157+
158+
if (this.isSingleUser()) {
159+
return (this._hasExecutePerms = this.isValidSingleUser(
160+
userDetails.userName
161+
));
162+
}
163+
164+
const permissionApi = new PermissionsService(this.client);
165+
const perms = await permissionApi.getObjectPermissions({
166+
object_id: this.id,
167+
object_type: "clusters",
168+
});
169+
170+
return (this._hasExecutePerms =
171+
(perms.access_control_list ?? []).find((ac) => {
172+
return (
173+
ac.user_name === userDetails.userName ||
174+
userDetails.groups
175+
?.map((v) => v.display)
176+
.includes(ac.group_name ?? "")
177+
);
178+
}) !== undefined);
179+
}
180+
123181
async refresh() {
124182
this.details = await this.clusterApi.get({
125183
cluster_id: this.clusterDetails.cluster_id!,
@@ -145,6 +203,7 @@ export class Cluster {
145203
});
146204
}
147205

206+
this._canExecute = undefined;
148207
await retry({
149208
fn: async () => {
150209
if (token?.isCancellationRequested) {
@@ -204,21 +263,26 @@ export class Cluster {
204263
return await ExecutionContext.create(this.client, this, language);
205264
}
206265

207-
async canExecute(): Promise<boolean> {
208-
let context: ExecutionContext | undefined;
266+
get canExecuteCached() {
267+
return this._canExecute;
268+
}
269+
270+
@withLogContext(ExposedLoggers.SDK)
271+
async canExecute(@context ctx?: Context): Promise<boolean> {
272+
let executionContext: ExecutionContext | undefined;
209273
try {
210-
context = await this.createExecutionContext();
211-
let result = await context.execute("print('hello')");
212-
if (result.result?.results?.resultType === "error") {
213-
return false;
214-
}
215-
return true;
274+
executionContext = await this.createExecutionContext();
275+
let result = await executionContext.execute("1==1");
276+
this._canExecute =
277+
result.result?.results?.resultType === "error" ? false : true;
216278
} catch (e) {
217-
return false;
279+
ctx?.logger?.error(`Can't execute code on cluster ${this.id}`, e);
280+
this._canExecute = false;
218281
} finally {
219-
if (context) {
220-
await context.destroy();
282+
if (executionContext) {
283+
await executionContext.destroy();
221284
}
285+
return this._canExecute ?? false;
222286
}
223287
}
224288

‎packages/databricks-vscode/package.json‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -494,28 +494,29 @@
494494
"title": "Databricks",
495495
"properties": {
496496
"databricks.logs.maxFieldLength": {
497-
"title": "Max Field Length",
498497
"type": "number",
499498
"default": 40,
500499
"description": "The maximum length of each field displayed in logs outputs panel."
501500
},
502501
"databricks.logs.truncationDepth": {
503-
"title": "Truncation Depth",
504502
"type": "number",
505503
"default": 2,
506504
"description": "The max depth of logs to show without truncation."
507505
},
508506
"databricks.logs.maxArrayLength": {
509-
"title": "Max Array Length",
510507
"type": "number",
511508
"default": 2,
512509
"description": "The maximum number of items to show for array fields."
513510
},
514511
"databricks.logs.enabled": {
515-
"title": "Enabled",
516512
"type": "boolean",
517513
"default": true,
518514
"description": "Enable/disable logging. Reload window for changes to take effect."
515+
},
516+
"databricks.clusters.onlyShowAccessibleClusters": {
517+
"type": "boolean",
518+
"default": true,
519+
"description": "Enable/disable filtering for only accessible clusters (clusters on which the current user can run code)"
519520
}
520521
}
521522
}
@@ -599,4 +600,4 @@
599600
],
600601
"report-dir": "coverage"
601602
}
602-
}
603+
}

packages/databricks-vscode/src/logger/WorkspaceConfigs.ts renamed to packages/databricks-vscode/src/WorkspaceConfigs.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,4 +29,11 @@ export const workspaceConfigs = {
2929
?.get<boolean>("logs.enabled") ?? true
3030
);
3131
},
32+
get onlyShowAccessibleClusters() {
33+
return (
34+
workspace
35+
.getConfiguration("databricks")
36+
?.get<boolean>("clusters.onlyShowAccessibleClusters") ?? true
37+
);
38+
},
3239
};

‎packages/databricks-vscode/src/cluster/ClusterLoader.ts‎

Lines changed: 66 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import {
88
import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
99
import {Disposable, Event, EventEmitter} from "vscode";
1010
import {ConnectionManager} from "../configuration/ConnectionManager";
11+
import {Loggers} from "../logger";
12+
import {workspaceConfigs} from "../WorkspaceConfigs";
1113
import {sortClusters} from "./ClusterModel";
1214

1315
export class ClusterLoader implements Disposable {
@@ -60,51 +62,10 @@ export class ClusterLoader implements Disposable {
6062
this.disposables.push(this.onDidStop(() => (this.stopped = true)));
6163
}
6264

63-
private isSingleUser(c: Cluster) {
64-
const modeProperty =
65-
//TODO: deprecate data_security_mode once access_mode is available everywhere
66-
c.details.access_mode ?? c.details.data_security_mode;
67-
return (
68-
modeProperty !== undefined &&
69-
[
70-
"SINGLE_USER",
71-
"LEGACY_SINGLE_USER_PASSTHROUGH",
72-
"LEGACY_SINGLE_USER_STANDARD",
73-
//enums unique to data_security_mode
74-
"LEGACY_SINGLE_USER",
75-
].includes(modeProperty)
76-
);
77-
}
78-
private isValidSingleUser(c: Cluster) {
79-
return (
80-
this.isSingleUser(c) &&
81-
c.details.single_user_name ===
82-
this.connectionManager.databricksWorkspace?.userName
83-
);
84-
}
85-
86-
private async hasPerm(c: Cluster, permissionApi: PermissionsService) {
87-
const perms = await permissionApi.getObjectPermissions({
88-
object_id: c.id,
89-
object_type: "clusters",
90-
});
91-
return (
92-
(perms.access_control_list ?? []).find((ac) => {
93-
return (
94-
ac.user_name ===
95-
this.connectionManager.databricksWorkspace?.userName ||
96-
this.connectionManager.databricksWorkspace?.user.groups
97-
?.map((v) => v.display)
98-
.includes(ac.group_name ?? "")
99-
);
100-
}) !== undefined
101-
);
102-
}
103-
10465
private cleanupClustersMap(clusters: Cluster[]) {
10566
const clusterIds = clusters.map((c) => c.id);
10667
const toDelete = [];
107-
for (let key in this._clusters) {
68+
for (let key of this._clusters.keys()) {
10869
if (!clusterIds.includes(key)) {
10970
toDelete.push(key);
11071
}
@@ -125,57 +86,76 @@ export class ClusterLoader implements Disposable {
12586
(await Cluster.list(apiClient))
12687
.filter((c) => ["UI", "API"].includes(c.source))
12788
.filter(
128-
(c) => !this.isSingleUser(c) || this.isValidSingleUser(c)
89+
(c) =>
90+
!workspaceConfigs.onlyShowAccessibleClusters ||
91+
!c.isSingleUser() ||
92+
c.isValidSingleUser(
93+
this.connectionManager.databricksWorkspace?.userName
94+
)
12995
)
130-
.filter((c) =>
131-
this.connectionManager.databricksWorkspace?.supportFilesInReposForCluster(
132-
c
133-
)
96+
.filter(
97+
(c) =>
98+
!workspaceConfigs.onlyShowAccessibleClusters ||
99+
this.connectionManager.databricksWorkspace?.supportFilesInReposForCluster(
100+
c
101+
)
134102
)
135103
);
136104

137-
const permissionApi = new PermissionsService(apiClient);
105+
if (workspaceConfigs.onlyShowAccessibleClusters) {
106+
// TODO: Find exact rate limit and update this.
107+
// Rate limit is 100 on dogfood.
108+
const maxConcurrent = 50;
109+
const wip: Promise<void>[] = [];
138110

139-
// TODO: Find exact rate limit and update this.
140-
// Rate limit is 100 on dogfood.
141-
const maxConcurrent = 50;
142-
const wip: Promise<void>[] = [];
111+
for (let c of allClusters) {
112+
if (!this.running) {
113+
break;
114+
}
115+
while (wip.length === maxConcurrent) {
116+
await Promise.race(wip);
117+
}
143118

144-
for (let c of allClusters) {
145-
if (!this.running) {
146-
break;
147-
}
148-
while (wip.length === maxConcurrent) {
149-
await Promise.race(wip);
119+
const task = new Promise<void>((resolve) => {
120+
c.hasExecutePerms(
121+
this.connectionManager.databricksWorkspace?.user
122+
)
123+
.then((keepCluster) => {
124+
if (!this.running) {
125+
return resolve();
126+
}
127+
128+
if (this._clusters.has(c.id) && !keepCluster) {
129+
this._clusters.delete(c.id);
130+
this._onDidChange.fire();
131+
}
132+
if (keepCluster) {
133+
this._clusters.set(c.id, c);
134+
this._onDidChange.fire();
135+
}
136+
resolve();
137+
})
138+
.catch((e) => {
139+
NamedLogger.getOrCreate(Loggers.Extension).error(
140+
`Error fetching permission for cluster ${c.name}`,
141+
e
142+
);
143+
resolve();
144+
});
145+
});
146+
147+
wip.push(task);
148+
task.then(() => {
149+
wip.splice(wip.indexOf(task), 1);
150+
});
150151
}
151152

152-
const task = new Promise<void>((resolve) => {
153-
this.hasPerm(c, permissionApi)
154-
.then((keepCluster) => {
155-
if (!this.running) {
156-
return resolve();
157-
}
158-
159-
if (this._clusters.has(c.id) && !keepCluster) {
160-
this._clusters.delete(c.id);
161-
this._onDidChange.fire();
162-
}
163-
if (keepCluster) {
164-
this._clusters.set(c.id, c);
165-
this._onDidChange.fire();
166-
}
167-
resolve();
168-
})
169-
.catch(resolve);
170-
});
171-
172-
wip.push(task);
173-
task.then(() => {
174-
wip.splice(wip.indexOf(task), 1);
175-
});
153+
await Promise.allSettled(wip);
154+
} else {
155+
this._clusters = new Map(allClusters.map((c) => [c.id, c]));
156+
this._onDidChange.fire();
176157
}
177158

178-
await Promise.allSettled(wip);
179159
this.cleanupClustersMap(allClusters);
180160
}
181161

@@ -189,26 +169,9 @@ export class ClusterLoader implements Disposable {
189169
try {
190170
await this._load();
191171
} catch (e) {
192-
let err = e;
193-
194-
/*
195-
Standard Error class has message and stack fields set as non enumerable.
196-
To correctly account for all such fields, we iterate over all own-properties of
197-
the error object and accumulate them as enumerable fields in the final err object.
198-
*/
199-
if (Object(err) === err) {
200-
err = {
201-
...Object.getOwnPropertyNames(err).reduce((acc, i) => {
202-
acc[i] = (err as any)[i];
203-
return acc;
204-
}, {} as any),
205-
...(err as any),
206-
};
207-
}
208-
NamedLogger.getOrCreate("Extension").log(
209-
"error",
210-
"Error loading clusters:",
211-
err
172+
NamedLogger.getOrCreate(Loggers.Extension).error(
173+
"Error loading clusters",
174+
e
212175
);
213176
}
214177
if (!this.running) {

0 commit comments

Comments
 (0)