Skip to content

Make cluster filtering opt-out #185

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kartikgupta-db merged 9 commits into from
Nov 9, 2022
Merged

Make cluster filtering opt-out #185

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
356390c
Make cluster filtering opt-out
kartikgupta-db Nov 8, 2022
f6fed0d
Additional canExecute check on attaching a new cluster
kartikgupta-db Nov 9, 2022
55698d3
Additional canExecute check on attaching a new cluster
kartikgupta-db Nov 9, 2022
27409fd
code cleanup
kartikgupta-db Nov 9, 2022
4528933
Merge remote-tracking branch 'databricks/main' into optin-cluster-fil…
kartikgupta-db Nov 9, 2022
5e331fb
setting name change
kartikgupta-db Nov 9, 2022
d3f000f
setting name change
kartikgupta-db Nov 9, 2022
d1bf79f
typo
kartikgupta-db Nov 9, 2022
6f3a250
Merge remote-tracking branch 'databricks/main' into optin-cluster-fil…
kartikgupta-db Nov 9, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 77 additions & 13 deletions packages/databricks-sdk-js/src/services/Cluster.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,19 +11,23 @@ import {
import {CancellationToken} from "../types";
import {ExecutionContext} from "./ExecutionContext";
import {WorkflowRun} from "./WorkflowRun";
import {commands} from "..";
import {commands, PermissionsService} from "..";
import {
ClusterInfo,
ClustersService,
ClusterInfoState,
ClusterInfoClusterSource,
} from "../apis/clusters";
import {Context} from "../context";
import {Context, context} from "../context";
import {User} from "../apis/scim";
import {ExposedLoggers, withLogContext} from "../logging";

export class ClusterRetriableError extends RetriableError {}
export class ClusterError extends Error {}
export class Cluster {
private clusterApi: ClustersService;
private _canExecute?: boolean;
private _hasExecutePerms?: boolean;

constructor(
private client: ApiClient,
Expand Down Expand Up @@ -120,6 +124,60 @@ export class Cluster {
this.clusterDetails = details;
}

isSingleUser() {
const modeProperty =
//TODO: deprecate data_security_mode once access_mode is available everywhere
this.details.access_mode ?? this.details.data_security_mode;
return (
modeProperty !== undefined &&
[
"SINGLE_USER",
"LEGACY_SINGLE_USER_PASSTHROUGH",
"LEGACY_SINGLE_USER_STANDARD",
//enums unique to data_security_mode
"LEGACY_SINGLE_USER",
].includes(modeProperty)
);
}

isValidSingleUser(userName?: string) {
return (
this.isSingleUser() && this.details.single_user_name === userName
);
}

get hasExecutePermsCached() {
return this._hasExecutePerms;
}

async hasExecutePerms(userDetails?: User) {
if (userDetails === undefined) {
return (this._hasExecutePerms = false);
}

if (this.isSingleUser()) {
return (this._hasExecutePerms = this.isValidSingleUser(
userDetails.userName
));
}

const permissionApi = new PermissionsService(this.client);
const perms = await permissionApi.getObjectPermissions({
object_id: this.id,
object_type: "clusters",
});

return (this._hasExecutePerms =
(perms.access_control_list ?? []).find((ac) => {
return (
ac.user_name === userDetails.userName ||
userDetails.groups
?.map((v) => v.display)
.includes(ac.group_name ?? "")
);
}) !== undefined);
}

async refresh() {
this.details = await this.clusterApi.get({
cluster_id: this.clusterDetails.cluster_id!,
Expand All @@ -145,6 +203,7 @@ export class Cluster {
});
}

this._canExecute = undefined;
await retry({
fn: async () => {
if (token?.isCancellationRequested) {
Expand Down Expand Up @@ -204,21 +263,26 @@ export class Cluster {
return await ExecutionContext.create(this.client, this, language);
}

async canExecute(): Promise<boolean> {
let context: ExecutionContext | undefined;
get canExecuteCached() {
return this._canExecute;
}

@withLogContext(ExposedLoggers.SDK)
async canExecute(@context ctx?: Context): Promise<boolean> {
let executionContext: ExecutionContext | undefined;
try {
context = await this.createExecutionContext();
let result = await context.execute("print('hello')");
if (result.result?.results?.resultType === "error") {
return false;
}
return true;
executionContext = await this.createExecutionContext();
let result = await executionContext.execute("1==1");
this._canExecute =
result.result?.results?.resultType === "error" ? false : true;
} catch (e) {
return false;
ctx?.logger?.error(`Can't execute code on cluster ${this.id}`, e);
this._canExecute = false;
} finally {
if (context) {
await context.destroy();
if (executionContext) {
await executionContext.destroy();
}
return this._canExecute ?? false;
}
}

Expand Down
11 changes: 6 additions & 5 deletions packages/databricks-vscode/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -494,28 +494,29 @@
"title": "Databricks",
"properties": {
"databricks.logs.maxFieldLength": {
"title": "Max Field Length",
"type": "number",
"default": 40,
"description": "The maximum length of each field displayed in logs outputs panel."
},
"databricks.logs.truncationDepth": {
"title": "Truncation Depth",
"type": "number",
"default": 2,
"description": "The max depth of logs to show without truncation."
},
"databricks.logs.maxArrayLength": {
"title": "Max Array Length",
"type": "number",
"default": 2,
"description": "The maximum number of items to show for array fields."
},
"databricks.logs.enabled": {
"title": "Enabled",
"type": "boolean",
"default": true,
"description": "Enable/disable logging. Reload window for changes to take effect."
},
"databricks.clusters.onlyShowAccessibleClusters": {
"type": "boolean",
"default": true,
"description": "Enable/disable filtering for only accessible clusters (clusters on which the current user can run code)"
}
}
}
Expand Down Expand Up @@ -599,4 +600,4 @@
],
"report-dir": "coverage"
}
}
}
7 changes: 7 additions & 0 deletions ...cks-vscode/src/logger/WorkspaceConfigs.ts → ...databricks-vscode/src/WorkspaceConfigs.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,4 +29,11 @@ export const workspaceConfigs = {
?.get<boolean>("logs.enabled") ?? true
);
},
get onlyShowAccessibleClusters() {
return (
workspace
.getConfiguration("databricks")
?.get<boolean>("clusters.onlyShowAccessibleClusters") ?? true
);
},
};
169 changes: 66 additions & 103 deletions packages/databricks-vscode/src/cluster/ClusterLoader.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import {
import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
import {Disposable, Event, EventEmitter} from "vscode";
import {ConnectionManager} from "../configuration/ConnectionManager";
import {Loggers} from "../logger";
import {workspaceConfigs} from "../WorkspaceConfigs";
import {sortClusters} from "./ClusterModel";

export class ClusterLoader implements Disposable {
Expand Down Expand Up @@ -60,51 +62,10 @@ export class ClusterLoader implements Disposable {
this.disposables.push(this.onDidStop(() => (this.stopped = true)));
}

private isSingleUser(c: Cluster) {
const modeProperty =
//TODO: deprecate data_security_mode once access_mode is available everywhere
c.details.access_mode ?? c.details.data_security_mode;
return (
modeProperty !== undefined &&
[
"SINGLE_USER",
"LEGACY_SINGLE_USER_PASSTHROUGH",
"LEGACY_SINGLE_USER_STANDARD",
//enums unique to data_security_mode
"LEGACY_SINGLE_USER",
].includes(modeProperty)
);
}
private isValidSingleUser(c: Cluster) {
return (
this.isSingleUser(c) &&
c.details.single_user_name ===
this.connectionManager.databricksWorkspace?.userName
);
}

private async hasPerm(c: Cluster, permissionApi: PermissionsService) {
const perms = await permissionApi.getObjectPermissions({
object_id: c.id,
object_type: "clusters",
});
return (
(perms.access_control_list ?? []).find((ac) => {
return (
ac.user_name ===
this.connectionManager.databricksWorkspace?.userName ||
this.connectionManager.databricksWorkspace?.user.groups
?.map((v) => v.display)
.includes(ac.group_name ?? "")
);
}) !== undefined
);
}

private cleanupClustersMap(clusters: Cluster[]) {
const clusterIds = clusters.map((c) => c.id);
const toDelete = [];
for (let key in this._clusters) {
for (let key of this._clusters.keys()) {
if (!clusterIds.includes(key)) {
toDelete.push(key);
}
Expand All @@ -125,57 +86,76 @@ export class ClusterLoader implements Disposable {
(await Cluster.list(apiClient))
.filter((c) => ["UI", "API"].includes(c.source))
.filter(
(c) => !this.isSingleUser(c) || this.isValidSingleUser(c)
(c) =>
!workspaceConfigs.onlyShowAccessibleClusters ||
!c.isSingleUser() ||
c.isValidSingleUser(
this.connectionManager.databricksWorkspace?.userName
)
)
.filter((c) =>
this.connectionManager.databricksWorkspace?.supportFilesInReposForCluster(
c
)
.filter(
(c) =>
!workspaceConfigs.onlyShowAccessibleClusters ||
this.connectionManager.databricksWorkspace?.supportFilesInReposForCluster(
c
)
)
);

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

// TODO: Find exact rate limit and update this.
// Rate limit is 100 on dogfood.
const maxConcurrent = 50;
const wip: Promise<void>[] = [];
for (let c of allClusters) {
if (!this.running) {
break;
}
while (wip.length === maxConcurrent) {
await Promise.race(wip);
}

for (let c of allClusters) {
if (!this.running) {
break;
}
while (wip.length === maxConcurrent) {
await Promise.race(wip);
const task = new Promise<void>((resolve) => {
c.hasExecutePerms(
this.connectionManager.databricksWorkspace?.user
)
.then((keepCluster) => {
if (!this.running) {
return resolve();
}

if (this._clusters.has(c.id) && !keepCluster) {
this._clusters.delete(c.id);
this._onDidChange.fire();
}
if (keepCluster) {
this._clusters.set(c.id, c);
this._onDidChange.fire();
}
resolve();
})
.catch((e) => {
NamedLogger.getOrCreate(Loggers.Extension).error(
`Error fetching permission for cluster ${c.name}`,
e
);
resolve();
});
});

wip.push(task);
task.then(() => {
wip.splice(wip.indexOf(task), 1);
});
}

const task = new Promise<void>((resolve) => {
this.hasPerm(c, permissionApi)
.then((keepCluster) => {
if (!this.running) {
return resolve();
}

if (this._clusters.has(c.id) && !keepCluster) {
this._clusters.delete(c.id);
this._onDidChange.fire();
}
if (keepCluster) {
this._clusters.set(c.id, c);
this._onDidChange.fire();
}
resolve();
})
.catch(resolve);
});

wip.push(task);
task.then(() => {
wip.splice(wip.indexOf(task), 1);
});
await Promise.allSettled(wip);
} else {
this._clusters = new Map(allClusters.map((c) => [c.id, c]));
this._onDidChange.fire();
}

await Promise.allSettled(wip);
this.cleanupClustersMap(allClusters);
}

Expand All @@ -189,26 +169,9 @@ export class ClusterLoader implements Disposable {
try {
await this._load();
} catch (e) {
let err = e;

/*
Standard Error class has message and stack fields set as non enumerable.
To correctly account for all such fields, we iterate over all own-properties of
the error object and accumulate them as enumerable fields in the final err object.
*/
if (Object(err) === err) {
err = {
...Object.getOwnPropertyNames(err).reduce((acc, i) => {
acc[i] = (err as any)[i];
return acc;
}, {} as any),
...(err as any),
};
}
NamedLogger.getOrCreate("Extension").log(
"error",
"Error loading clusters:",
err
NamedLogger.getOrCreate(Loggers.Extension).error(
"Error loading clusters",
e
);
}
if (!this.running) {
Expand Down
Loading