From 356390c5aae42dd7b5ade77a979fdab4216065de Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Tue, 8 Nov 2022 18:48:19 +0100 Subject: [PATCH 1/7] Make cluster filtering opt-out --- .../databricks-sdk-js/src/services/Cluster.ts | 52 +++++- packages/databricks-vscode/package.json | 6 + .../src/{logger => }/WorkspaceConfigs.ts | 7 + .../src/cluster/ClusterLoader.ts | 168 +++++++----------- .../ConfigurationDataProvider.ts | 31 +++- .../databricks-vscode/src/logger/loggers.ts | 2 +- .../src/logger/outputConsoleTransport.ts | 2 +- 7 files changed, 161 insertions(+), 107 deletions(-) rename packages/databricks-vscode/src/{logger => }/WorkspaceConfigs.ts (79%) diff --git a/packages/databricks-sdk-js/src/services/Cluster.ts b/packages/databricks-sdk-js/src/services/Cluster.ts index 433b2d999..97de547ab 100644 --- a/packages/databricks-sdk-js/src/services/Cluster.ts +++ b/packages/databricks-sdk-js/src/services/Cluster.ts @@ -11,7 +11,7 @@ import { import {CancellationToken} from "../types"; import {ExecutionContext} from "./ExecutionContext"; import {WorkflowRun} from "./WorkflowRun"; -import {commands} from ".."; +import {commands, PermissionsService} from ".."; import { ClusterInfo, ClustersService, @@ -19,6 +19,7 @@ import { ClusterInfoClusterSource, } from "../apis/clusters"; import {Context} from "../context"; +import {User} from "../apis/scim"; export class ClusterRetriableError extends RetriableError {} export class ClusterError extends Error {} @@ -120,6 +121,55 @@ 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 + ); + } + + async hasExecutePerms(userDetails?: User) { + if (userDetails === undefined) { + return false; + } + + if (this.isSingleUser()) { + return this.isValidSingleUser(userDetails.userName); + } + + const permissionApi = new PermissionsService(this.client); + const perms = await permissionApi.getObjectPermissions({ + object_id: this.id, + object_type: "clusters", + }); + + return ( + (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!, diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index c2558b6f9..06dd845b9 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -516,6 +516,12 @@ "type": "boolean", "default": true, "description": "Enable/disable logging. Reload window for changes to take effect." + }, + "databricks.clusters.filteringEnabled": { + "title": "Filtering Enabled", + "type": "boolean", + "default": true, + "description": "Enabled/disable filtering for only accessible clusters (clusters on which the current user can run code)" } } } diff --git a/packages/databricks-vscode/src/logger/WorkspaceConfigs.ts b/packages/databricks-vscode/src/WorkspaceConfigs.ts similarity index 79% rename from packages/databricks-vscode/src/logger/WorkspaceConfigs.ts rename to packages/databricks-vscode/src/WorkspaceConfigs.ts index e51cc3aaa..114f57e4f 100644 --- a/packages/databricks-vscode/src/logger/WorkspaceConfigs.ts +++ b/packages/databricks-vscode/src/WorkspaceConfigs.ts @@ -29,4 +29,11 @@ export const workspaceConfigs = { ?.get("logs.enabled") ?? true ); }, + get clusterFilteringEnabled() { + return ( + workspace + .getConfiguration("databricks") + ?.get("clusters.filteringEnabled") ?? true + ); + }, }; diff --git a/packages/databricks-vscode/src/cluster/ClusterLoader.ts b/packages/databricks-vscode/src/cluster/ClusterLoader.ts index af5bed4b1..6608a042b 100644 --- a/packages/databricks-vscode/src/cluster/ClusterLoader.ts +++ b/packages/databricks-vscode/src/cluster/ClusterLoader.ts @@ -8,6 +8,7 @@ import { import {NamedLogger} from "@databricks/databricks-sdk/dist/logging"; import {Disposable, Event, EventEmitter} from "vscode"; import {ConnectionManager} from "../configuration/ConnectionManager"; +import {workspaceConfigs} from "../WorkspaceConfigs"; import {sortClusters} from "./ClusterModel"; export class ClusterLoader implements Disposable { @@ -60,51 +61,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); } @@ -125,57 +85,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.clusterFilteringEnabled || + !c.isSingleUser() || + c.isValidSingleUser( + this.connectionManager.databricksWorkspace?.userName + ) ) - .filter((c) => - this.connectionManager.databricksWorkspace?.supportFilesInReposForCluster( - c - ) + .filter( + (c) => + !workspaceConfigs.clusterFilteringEnabled || + this.connectionManager.databricksWorkspace?.supportFilesInReposForCluster( + c + ) ) ); - const permissionApi = new PermissionsService(apiClient); + if (workspaceConfigs.clusterFilteringEnabled) { + // TODO: Find exact rate limit and update this. + // Rate limit is 100 on dogfood. + const maxConcurrent = 50; + const wip: Promise[] = []; - // TODO: Find exact rate limit and update this. - // Rate limit is 100 on dogfood. - const maxConcurrent = 50; - const wip: Promise[] = []; + 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((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("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((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); } @@ -189,26 +168,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("Extension").error( + "Error loading clusters", + e ); } if (!this.running) { diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts index b68b50f09..77960a970 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts @@ -1,8 +1,10 @@ +import {NamedLogger} from "@databricks/databricks-sdk/dist/logging"; import { Disposable, Event, EventEmitter, ProviderResult, + ThemeColor, ThemeIcon, TreeDataProvider, TreeItem, @@ -43,6 +45,8 @@ export class ConfigurationDataProvider this._onDidChangeTreeData.fire(); }) ); + + this.connectionManager; } dispose() { @@ -146,9 +150,33 @@ export class ConfigurationDataProvider } if (element.id?.startsWith("CLUSTER") && cluster) { - let clusterItem = + const clusterItem = ClusterListDataProvider.clusterNodeToTreeItem(cluster); + const children = []; + + try { + if ( + !(await cluster.hasExecutePerms( + this.connectionManager.databricksWorkspace?.user + )) + ) { + children.push({ + description: + "You might not have permission to run code on this cluster", + iconPath: new ThemeIcon( + "warning", + new ThemeColor("problemsWarningIcon.foreground") + ), + }); + } + } catch (e) { + NamedLogger.getOrCreate("Extension").error( + `Error in fetching permissions for ${cluster.name}`, + e + ); + } + return [ { label: "Name:", @@ -156,6 +184,7 @@ export class ConfigurationDataProvider iconPath: clusterItem.iconPath, collapsibleState: TreeItemCollapsibleState.None, }, + ...children, ...(await ClusterListDataProvider.clusterNodeToTreeItems( cluster )), diff --git a/packages/databricks-vscode/src/logger/loggers.ts b/packages/databricks-vscode/src/logger/loggers.ts index a6a2f27e5..392f4b85d 100644 --- a/packages/databricks-vscode/src/logger/loggers.ts +++ b/packages/databricks-vscode/src/logger/loggers.ts @@ -6,7 +6,7 @@ import {window} from "vscode"; import {loggers, format, transports} from "winston"; import {getOutputConsoleTransport} from "./outputConsoleTransport"; import {unlink, access} from "fs/promises"; -import {workspaceConfigs} from "./WorkspaceConfigs"; +import {workspaceConfigs} from "../WorkspaceConfigs"; function getFileTransport(filename: string) { return new transports.File({ diff --git a/packages/databricks-vscode/src/logger/outputConsoleTransport.ts b/packages/databricks-vscode/src/logger/outputConsoleTransport.ts index b31f6c523..57db467a6 100644 --- a/packages/databricks-vscode/src/logger/outputConsoleTransport.ts +++ b/packages/databricks-vscode/src/logger/outputConsoleTransport.ts @@ -2,8 +2,8 @@ import {OutputChannel} from "vscode"; import {transports, format} from "winston"; import {OutputConsoleStream} from "./OutputConsoleStream"; import {LEVEL, MESSAGE, SPLAT} from "triple-beam"; -import {workspaceConfigs} from "./WorkspaceConfigs"; import {inspect} from "util"; +import {workspaceConfigs} from "../WorkspaceConfigs"; function processPrimitiveOrString(obj: any) { let valueStr: string; From f6fed0d477ea244cd11ffd200018bfc3f546f440 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 9 Nov 2022 10:31:12 +0100 Subject: [PATCH 2/7] Additional canExecute check on attaching a new cluster --- .../databricks-sdk-js/src/services/Cluster.ts | 35 ++++++++----- .../src/cluster/ClusterLoader.ts | 5 +- .../ConfigurationDataProvider.ts | 47 ++++++++++------- .../src/configuration/ConnectionManager.ts | 5 ++ .../databricks-vscode/src/logger/index.ts | 1 + .../databricks-vscode/src/logger/loggers.ts | 5 ++ .../databricks-vscode/src/logger/utils.ts | 52 +++++++++++++++++++ 7 files changed, 116 insertions(+), 34 deletions(-) create mode 100644 packages/databricks-vscode/src/logger/utils.ts diff --git a/packages/databricks-sdk-js/src/services/Cluster.ts b/packages/databricks-sdk-js/src/services/Cluster.ts index 97de547ab..99ea3f8ca 100644 --- a/packages/databricks-sdk-js/src/services/Cluster.ts +++ b/packages/databricks-sdk-js/src/services/Cluster.ts @@ -18,13 +18,15 @@ import { 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; constructor( private client: ApiClient, @@ -195,6 +197,7 @@ export class Cluster { }); } + this._canExecute = undefined; await retry({ fn: async () => { if (token?.isCancellationRequested) { @@ -254,21 +257,29 @@ export class Cluster { return await ExecutionContext.create(this.client, this, language); } - async canExecute(): Promise { - let context: ExecutionContext | undefined; + @withLogContext(ExposedLoggers.SDK) + async canExecute( + useCache = false, + @context ctx?: Context + ): Promise { + if (useCache) { + return this._canExecute; + } + + 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; } } diff --git a/packages/databricks-vscode/src/cluster/ClusterLoader.ts b/packages/databricks-vscode/src/cluster/ClusterLoader.ts index 6608a042b..e6b536d38 100644 --- a/packages/databricks-vscode/src/cluster/ClusterLoader.ts +++ b/packages/databricks-vscode/src/cluster/ClusterLoader.ts @@ -8,6 +8,7 @@ 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"; @@ -135,7 +136,7 @@ export class ClusterLoader implements Disposable { resolve(); }) .catch((e) => { - NamedLogger.getOrCreate("Extension").error( + NamedLogger.getOrCreate(Loggers.Extension).error( `Error fetching permission for cluster ${c.name}`, e ); @@ -168,7 +169,7 @@ export class ClusterLoader implements Disposable { try { await this._load(); } catch (e) { - NamedLogger.getOrCreate("Extension").error( + NamedLogger.getOrCreate(Loggers.Extension).error( "Error loading clusters", e ); diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts index 77960a970..9fddc8f67 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts @@ -11,6 +11,7 @@ import { TreeItemCollapsibleState, } from "vscode"; import {ClusterListDataProvider} from "../cluster/ClusterListDataProvider"; +import {Loggers, loggingUtils} from "../logger"; import {CodeSynchronizer} from "../sync/CodeSynchronizer"; import {ConnectionManager} from "./ConnectionManager"; @@ -155,26 +156,32 @@ export class ConfigurationDataProvider const children = []; - try { - if ( - !(await cluster.hasExecutePerms( - this.connectionManager.databricksWorkspace?.user - )) - ) { - children.push({ - description: - "You might not have permission to run code on this cluster", - iconPath: new ThemeIcon( - "warning", - new ThemeColor("problemsWarningIcon.foreground") - ), - }); - } - } catch (e) { - NamedLogger.getOrCreate("Extension").error( - `Error in fetching permissions for ${cluster.name}`, - e - ); + const runPerms = + cluster.state === "RUNNING" && + (await cluster?.canExecute(true)) !== undefined + ? await cluster?.canExecute(true) + : await loggingUtils.tryAndLogErrorAsync( + async () => { + return await cluster?.hasExecutePerms( + this.connectionManager.databricksWorkspace + ?.user + ); + }, + { + message: `Error in fetching permissions for ${cluster.name}`, + shouldThrow: false, + } + ); + + if (runPerms === false) { + children.push({ + description: + "You do not have permission to run code on this cluster", + iconPath: new ThemeIcon( + "warning", + new ThemeColor("problemsWarningIcon.foreground") + ), + }); } return [ diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index d26766f5f..8119aa7be 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -290,6 +290,11 @@ export class ConnectionManager { await this._projectConfigFile!.write(); } + if (cluster.state === "RUNNING") { + cluster.canExecute(false).then(() => { + this.onDidChangeClusterEmitter.fire(this.cluster); + }); + } this.updateCluster(cluster); } diff --git a/packages/databricks-vscode/src/logger/index.ts b/packages/databricks-vscode/src/logger/index.ts index ca2dd5847..32930d069 100644 --- a/packages/databricks-vscode/src/logger/index.ts +++ b/packages/databricks-vscode/src/logger/index.ts @@ -1 +1,2 @@ export * from "./loggers"; +export * as loggingUtils from "./utils"; diff --git a/packages/databricks-vscode/src/logger/loggers.ts b/packages/databricks-vscode/src/logger/loggers.ts index 392f4b85d..1600f20e4 100644 --- a/packages/databricks-vscode/src/logger/loggers.ts +++ b/packages/databricks-vscode/src/logger/loggers.ts @@ -66,3 +66,8 @@ export async function initLoggers(rootPath: string) { true ); } + +export enum Loggers { + // eslint-disable-next-line @typescript-eslint/naming-convention + Extension = "Extension", +} diff --git a/packages/databricks-vscode/src/logger/utils.ts b/packages/databricks-vscode/src/logger/utils.ts new file mode 100644 index 000000000..1b611bf28 --- /dev/null +++ b/packages/databricks-vscode/src/logger/utils.ts @@ -0,0 +1,52 @@ +import {NamedLogger} from "@databricks/databricks-sdk/dist/logging"; +import {Loggers} from "./loggers"; + +export interface TryAndLogErrorOpts { + shouldThrow: boolean; + message: string; + logger: Loggers; +} + +const defaultTryAndLogErrorOpts: TryAndLogErrorOpts = { + shouldThrow: true, + message: "", + logger: Loggers.Extension, +}; + +export async function tryAndLogErrorAsync( + fn: () => Promise, + opts: Partial = {} +): Promise { + const mergedOpts: TryAndLogErrorOpts = { + ...defaultTryAndLogErrorOpts, + ...opts, + }; + + try { + return await fn(); + } catch (e) { + NamedLogger.getOrCreate(mergedOpts.logger).error(mergedOpts.message, e); + if (mergedOpts.shouldThrow) { + throw e; + } + } +} + +export function tryAndLogError( + fn: () => T, + opts: Partial = {} +): T | undefined { + const mergedOpts: TryAndLogErrorOpts = { + ...defaultTryAndLogErrorOpts, + ...opts, + }; + + try { + return fn(); + } catch (e) { + NamedLogger.getOrCreate(mergedOpts.logger).error(mergedOpts.message, e); + if (mergedOpts.shouldThrow) { + throw e; + } + } +} From 55698d35ebf2bda9b7ed2f0d00592403ea10327b Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 9 Nov 2022 11:06:16 +0100 Subject: [PATCH 3/7] Additional canExecute check on attaching a new cluster --- .../ConfigurationDataProvider.ts | 80 +++++++++++++------ 1 file changed, 56 insertions(+), 24 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts index 9fddc8f67..0d5422979 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts @@ -156,32 +156,64 @@ export class ConfigurationDataProvider const children = []; - const runPerms = + let runPerms: + | "CAN_RUN" + | "MIGHT_RUN" + | "UNABLE_TO_RUN" + | "MIGHT_NOT_RUN" = "MIGHT_RUN"; + if ( cluster.state === "RUNNING" && (await cluster?.canExecute(true)) !== undefined - ? await cluster?.canExecute(true) - : await loggingUtils.tryAndLogErrorAsync( - async () => { - return await cluster?.hasExecutePerms( - this.connectionManager.databricksWorkspace - ?.user - ); - }, - { - message: `Error in fetching permissions for ${cluster.name}`, - shouldThrow: false, - } - ); - - if (runPerms === false) { - children.push({ - description: - "You do not have permission to run code on this cluster", - iconPath: new ThemeIcon( - "warning", - new ThemeColor("problemsWarningIcon.foreground") - ), - }); + ) { + runPerms = (await cluster?.canExecute(true)) + ? "CAN_RUN" + : "UNABLE_TO_RUN"; + } else { + runPerms = (await loggingUtils.tryAndLogErrorAsync( + async () => { + return await cluster?.hasExecutePerms( + this.connectionManager.databricksWorkspace?.user + ); + }, + { + message: `Error in fetching permissions for ${cluster.name}`, + shouldThrow: false, + } + )) + ? "MIGHT_RUN" + : "MIGHT_NOT_RUN"; + } + + switch (runPerms) { + case "CAN_RUN": + children.push({ + label: "You can run code on this cluster", + iconPath: new ThemeIcon( + "testing-passed-icon", + new ThemeColor("testing.iconPassed") + ), + }); + break; + + case "MIGHT_NOT_RUN": + children.push({ + label: "You might not have permissions to run code on this cluster", + iconPath: new ThemeIcon( + "warning", + new ThemeColor("problemsWarningIcon.foreground") + ), + }); + break; + + case "UNABLE_TO_RUN": + children.push({ + label: "You do not have permissions to run code on this cluster", + iconPath: new ThemeIcon( + "alert", + new ThemeColor("testing.iconFailed") + ), + }); + break; } return [ From 27409fd79c7906dbe56091546448699f58fa04c6 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 9 Nov 2022 11:33:08 +0100 Subject: [PATCH 4/7] code cleanup --- .../databricks-sdk-js/src/services/Cluster.ts | 29 +++++++++-------- .../ConfigurationDataProvider.ts | 21 ++++--------- .../src/configuration/ConnectionManager.ts | 31 +++++++++++++++++-- 3 files changed, 51 insertions(+), 30 deletions(-) diff --git a/packages/databricks-sdk-js/src/services/Cluster.ts b/packages/databricks-sdk-js/src/services/Cluster.ts index 99ea3f8ca..16b5ec8fe 100644 --- a/packages/databricks-sdk-js/src/services/Cluster.ts +++ b/packages/databricks-sdk-js/src/services/Cluster.ts @@ -27,6 +27,7 @@ export class ClusterError extends Error {} export class Cluster { private clusterApi: ClustersService; private _canExecute?: boolean; + private _hasExecutePerms?: boolean; constructor( private client: ApiClient, @@ -145,13 +146,19 @@ export class Cluster { ); } + get hasExecutePermsCached() { + return this._hasExecutePerms; + } + async hasExecutePerms(userDetails?: User) { if (userDetails === undefined) { - return false; + return (this._hasExecutePerms = false); } if (this.isSingleUser()) { - return this.isValidSingleUser(userDetails.userName); + return (this._hasExecutePerms = this.isValidSingleUser( + userDetails.userName + )); } const permissionApi = new PermissionsService(this.client); @@ -160,7 +167,7 @@ export class Cluster { object_type: "clusters", }); - return ( + return (this._hasExecutePerms = (perms.access_control_list ?? []).find((ac) => { return ( ac.user_name === userDetails.userName || @@ -168,8 +175,7 @@ export class Cluster { ?.map((v) => v.display) .includes(ac.group_name ?? "") ); - }) !== undefined - ); + }) !== undefined); } async refresh() { @@ -257,15 +263,12 @@ export class Cluster { return await ExecutionContext.create(this.client, this, language); } - @withLogContext(ExposedLoggers.SDK) - async canExecute( - useCache = false, - @context ctx?: Context - ): Promise { - if (useCache) { - return this._canExecute; - } + get canExecuteCached() { + return this._canExecute; + } + @withLogContext(ExposedLoggers.SDK) + async canExecute(@context ctx?: Context): Promise { let executionContext: ExecutionContext | undefined; try { executionContext = await this.createExecutionContext(); diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts index 0d5422979..c51510bff 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts @@ -163,25 +163,16 @@ export class ConfigurationDataProvider | "MIGHT_NOT_RUN" = "MIGHT_RUN"; if ( cluster.state === "RUNNING" && - (await cluster?.canExecute(true)) !== undefined + cluster?.canExecuteCached !== undefined ) { - runPerms = (await cluster?.canExecute(true)) + runPerms = cluster.canExecuteCached ? "CAN_RUN" : "UNABLE_TO_RUN"; } else { - runPerms = (await loggingUtils.tryAndLogErrorAsync( - async () => { - return await cluster?.hasExecutePerms( - this.connectionManager.databricksWorkspace?.user - ); - }, - { - message: `Error in fetching permissions for ${cluster.name}`, - shouldThrow: false, - } - )) - ? "MIGHT_RUN" - : "MIGHT_NOT_RUN"; + runPerms = + cluster.hasExecutePermsCached ?? true + ? "MIGHT_RUN" + : "MIGHT_NOT_RUN"; } switch (runPerms) { diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 8119aa7be..8ec7ef894 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -21,6 +21,8 @@ import {selectProfile} from "./selectProfileWizard"; import {ClusterManager} from "../cluster/ClusterManager"; import {workspace} from "@databricks/databricks-sdk"; import {DatabricksWorkspace} from "./DatabricksWorkspace"; +import {NamedLogger} from "@databricks/databricks-sdk/dist/logging"; +import {Loggers} from "../logger"; const extensionVersion = require("../../package.json").version; @@ -291,10 +293,35 @@ export class ConnectionManager { } if (cluster.state === "RUNNING") { - cluster.canExecute(false).then(() => { + cluster + .canExecute() + .then(() => { + this.onDidChangeClusterEmitter.fire(this.cluster); + }) + .catch((e) => { + NamedLogger.getOrCreate(Loggers.Extension).error( + `Error while running code on cluster ${ + (cluster as Cluster).id + }`, + e + ); + }); + } + + cluster + .hasExecutePerms(this.databricksWorkspace?.user) + .then(() => { this.onDidChangeClusterEmitter.fire(this.cluster); + }) + .catch((e) => { + NamedLogger.getOrCreate(Loggers.Extension).error( + `Error while fetching permission for cluster ${ + (cluster as Cluster).id + }`, + e + ); }); - } + this.updateCluster(cluster); } From 5e331fb595786461bb2ddceb7196966d429ede70 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 9 Nov 2022 16:22:16 +0100 Subject: [PATCH 5/7] setting name change --- packages/databricks-vscode/package.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 06dd845b9..da1eee116 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -518,7 +518,7 @@ "description": "Enable/disable logging. Reload window for changes to take effect." }, "databricks.clusters.filteringEnabled": { - "title": "Filtering Enabled", + "title": "Show Only Accessible Clusters", "type": "boolean", "default": true, "description": "Enabled/disable filtering for only accessible clusters (clusters on which the current user can run code)" @@ -605,4 +605,4 @@ ], "report-dir": "coverage" } -} \ No newline at end of file +} From d3f000fd47c20dadea15e593c05e5b792affbb33 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 9 Nov 2022 16:35:55 +0100 Subject: [PATCH 6/7] setting name change --- packages/databricks-vscode/package.json | 7 +------ packages/databricks-vscode/src/WorkspaceConfigs.ts | 4 ++-- packages/databricks-vscode/src/cluster/ClusterLoader.ts | 6 +++--- 3 files changed, 6 insertions(+), 11 deletions(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index da1eee116..990552790 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -494,31 +494,26 @@ "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.filteringEnabled": { - "title": "Show Only Accessible Clusters", + "databricks.clusters.onlyShowAccessibleClusters": { "type": "boolean", "default": true, "description": "Enabled/disable filtering for only accessible clusters (clusters on which the current user can run code)" diff --git a/packages/databricks-vscode/src/WorkspaceConfigs.ts b/packages/databricks-vscode/src/WorkspaceConfigs.ts index 114f57e4f..2a451c63e 100644 --- a/packages/databricks-vscode/src/WorkspaceConfigs.ts +++ b/packages/databricks-vscode/src/WorkspaceConfigs.ts @@ -29,11 +29,11 @@ export const workspaceConfigs = { ?.get("logs.enabled") ?? true ); }, - get clusterFilteringEnabled() { + get onlyShowAccessibleClusters() { return ( workspace .getConfiguration("databricks") - ?.get("clusters.filteringEnabled") ?? true + ?.get("clusters.onlyShowAccessibleClusters") ?? true ); }, }; diff --git a/packages/databricks-vscode/src/cluster/ClusterLoader.ts b/packages/databricks-vscode/src/cluster/ClusterLoader.ts index e6b536d38..0f6ca5091 100644 --- a/packages/databricks-vscode/src/cluster/ClusterLoader.ts +++ b/packages/databricks-vscode/src/cluster/ClusterLoader.ts @@ -87,7 +87,7 @@ export class ClusterLoader implements Disposable { .filter((c) => ["UI", "API"].includes(c.source)) .filter( (c) => - !workspaceConfigs.clusterFilteringEnabled || + !workspaceConfigs.onlyShowAccessibleClusters || !c.isSingleUser() || c.isValidSingleUser( this.connectionManager.databricksWorkspace?.userName @@ -95,14 +95,14 @@ export class ClusterLoader implements Disposable { ) .filter( (c) => - !workspaceConfigs.clusterFilteringEnabled || + !workspaceConfigs.onlyShowAccessibleClusters || this.connectionManager.databricksWorkspace?.supportFilesInReposForCluster( c ) ) ); - if (workspaceConfigs.clusterFilteringEnabled) { + if (workspaceConfigs.onlyShowAccessibleClusters) { // TODO: Find exact rate limit and update this. // Rate limit is 100 on dogfood. const maxConcurrent = 50; From d1bf79f31bb47c67e4d34f15c117d96bc9484ad6 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Wed, 9 Nov 2022 16:36:29 +0100 Subject: [PATCH 7/7] typo --- packages/databricks-vscode/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 990552790..6ac6581a4 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -516,7 +516,7 @@ "databricks.clusters.onlyShowAccessibleClusters": { "type": "boolean", "default": true, - "description": "Enabled/disable filtering for only accessible clusters (clusters on which the current user can run code)" + "description": "Enable/disable filtering for only accessible clusters (clusters on which the current user can run code)" } } }