diff --git a/packages/databricks-vscode-types/index.ts b/packages/databricks-vscode-types/index.ts index ab129cdf3..c0f1792fb 100644 --- a/packages/databricks-vscode-types/index.ts +++ b/packages/databricks-vscode-types/index.ts @@ -12,7 +12,6 @@ export interface PublicApi { connectionManager: { onDidChangeState: Event; - login(interactive?: boolean, force?: boolean): Promise; waitForConnect(): Promise; get state(): ConnectionState; diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 0dd472cdf..ffe5d20e6 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -211,9 +211,13 @@ { "command": "databricks.call", "title": "Call", - "enablement": "databricks.context.activated && databricks.context.loggedIn", + "enablement": "databricks.context.activated", "category": "Databricks" }, + { + "command": "databricks.internal.clearOverrides", + "title": "Clear workspace overrides" + }, { "command": "databricks.notebookInitScript.verify", "title": "Verify Databricks notebook init scripts", diff --git a/packages/databricks-vscode/src/bundle/types.ts b/packages/databricks-vscode/src/bundle/types.ts index 2a320a1c5..e28250bcf 100644 --- a/packages/databricks-vscode/src/bundle/types.ts +++ b/packages/databricks-vscode/src/bundle/types.ts @@ -1,3 +1,4 @@ -import {BundleSchema} from "./BundleSchema"; +export {type BundleSchema} from "./BundleSchema"; +import {BundleSchema} from "./BundleSchema"; export type BundleTarget = Required["targets"][string]; diff --git a/packages/databricks-vscode/src/cluster/ClusterManager.ts b/packages/databricks-vscode/src/cluster/ClusterManager.ts index 55030124e..36aeafabb 100644 --- a/packages/databricks-vscode/src/cluster/ClusterManager.ts +++ b/packages/databricks-vscode/src/cluster/ClusterManager.ts @@ -16,8 +16,13 @@ export class ClusterManager implements Disposable { private setInterval() { this.refreshTimer = setInterval(async () => { + const oldState = this.cluster.state; await this.cluster.refresh(); - this.onChange(this.cluster.state); + if ( + JSON.stringify(oldState) !== JSON.stringify(this.cluster.state) + ) { + this.onChange(this.cluster.state); + } }, this.refreshTimeout.toMillSeconds().value); } diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts index 79d7cf3f6..30509a973 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts @@ -14,6 +14,8 @@ import {StateStorage} from "../vscode-objs/StateStorage"; import {WorkspaceFsAccessVerifier} from "../workspace-fs"; import {FeatureManager} from "../feature-manager/FeatureManager"; import {Telemetry} from "../telemetry"; +import {ConfigModel} from "./ConfigModel"; +import {expect} from "chai"; describe(__filename, () => { let connectionManagerMock: ConnectionManager; @@ -21,7 +23,7 @@ describe(__filename, () => { let onChangeClusterListener: (e: Cluster) => void; let onChangeSyncDestinationListener: (e: SyncDestinationMapper) => void; let sync: CodeSynchronizer; - + let mockConfigModel: ConfigModel; beforeEach(() => { disposables = []; connectionManagerMock = mock(ConnectionManager); @@ -56,6 +58,13 @@ describe(__filename, () => { dispose() {}, }); sync = instance(syncMock); + + mockConfigModel = mock(ConfigModel); + mockConfigModel.onDidChangeAny = () => { + return { + dispose() {}, + }; + }; }); afterEach(() => { @@ -74,7 +83,8 @@ describe(__filename, () => { instance(mock(StateStorage)), instance(mock(WorkspaceFsAccessVerifier)), instance(mock(FeatureManager<"debugging.dbconnect">)), - instance(mock(Telemetry)) + instance(mock(Telemetry)), + mockConfigModel ); disposables.push(provider); @@ -98,7 +108,8 @@ describe(__filename, () => { instance(mock(StateStorage)), instance(mock(WorkspaceFsAccessVerifier)), instance(mock(FeatureManager<"debugging.dbconnect">)), - instance(mock(Telemetry)) + instance(mock(Telemetry)), + mockConfigModel ); disposables.push(provider); @@ -122,16 +133,17 @@ describe(__filename, () => { instance(mock(StateStorage)), instance(mock(WorkspaceFsAccessVerifier)), instance(mock(FeatureManager<"debugging.dbconnect">)), - instance(mock(Telemetry)) + instance(mock(Telemetry)), + mockConfigModel ); disposables.push(provider); const children = await resolveProviderResult(provider.getChildren()); assert(children); - assert.equal(children.length, 0); + assert.equal(children.length, 4); }); - it("should return cluster children", async () => { + it("should return children", async () => { const mockApiClient = mock(ApiClient); when(mockApiClient.host).thenResolve( new URL("https://www.example.com") @@ -155,12 +167,13 @@ describe(__filename, () => { instance(mock(StateStorage)), instance(mock(WorkspaceFsAccessVerifier)), instance(mock(FeatureManager<"debugging.dbconnect">)), - instance(mock(Telemetry)) + instance(mock(Telemetry)), + mockConfigModel ); disposables.push(provider); const children = await resolveProviderResult(provider.getChildren()); - assert.deepEqual(children, [ + expect(children).to.include.deep.members([ { collapsibleState: 2, contextValue: "workspace", diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts index 7a76c943e..aa36a7400 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts @@ -21,6 +21,7 @@ import { } from "../workspace-fs"; import {FeatureManager} from "../feature-manager/FeatureManager"; import {Telemetry} from "../telemetry"; +import {ConfigModel} from "./ConfigModel"; export type ConfigurationTreeItem = TreeItem & { url?: string; @@ -47,7 +48,8 @@ export class ConfigurationDataProvider private readonly stateStorage: StateStorage, private readonly wsfsAccessVerifier: WorkspaceFsAccessVerifier, private readonly featureManager: FeatureManager, - private readonly telemetry: Telemetry + private readonly telemetry: Telemetry, + private readonly configModel: ConfigModel ) { this.disposables.push( this.connectionManager.onDidChangeState(() => { @@ -64,6 +66,9 @@ export class ConfigurationDataProvider }), this.wsfsAccessVerifier.onDidChangeState(() => { this._onDidChangeTreeData.fire(); + }), + this.configModel.onDidChangeAny(() => { + this._onDidChangeTreeData.fire(); }) ); @@ -82,13 +87,12 @@ export class ConfigurationDataProvider element?: ConfigurationTreeItem | undefined ): Promise> { switch (this.connectionManager.state) { + case "DISCONNECTED": case "CONNECTED": break; case "CONNECTING": await this.connectionManager.waitForConnect(); break; - case "DISCONNECTED": - return []; } const cluster = this.connectionManager.cluster; @@ -96,14 +100,52 @@ export class ConfigurationDataProvider if (!element) { const children: Array = []; - children.push({ - label: `Workspace`, - iconPath: new ThemeIcon("account"), - id: "WORKSPACE", - collapsibleState: TreeItemCollapsibleState.Expanded, - contextValue: "workspace", - url: this.connectionManager.databricksWorkspace?.host?.toString(), - }); + children.push( + { + label: + this.configModel.target !== undefined + ? `Bundle Target - ${this.configModel.target}` + : `Bundle Target - "None selected"`, + id: "BUNDLE-TARGET", + collapsibleState: TreeItemCollapsibleState.Expanded, + contextValue: "bundleTarget", + command: { + title: "Call", + command: "databricks.call", + arguments: [ + async () => { + const targets = await this.configModel + .bundleConfigReaderWriter.targets; + if (targets === undefined) { + return; + } + + const selectedTarget = + await window.showQuickPick( + Object.keys(targets), + {title: "Select bundle target"} + ); + if (selectedTarget === undefined) { + return; + } + const currentTarget = this.configModel.target; + if (currentTarget !== selectedTarget) { + this._onDidChangeTreeData.fire(); + } + this.configModel.setTarget(selectedTarget); + }, + ], + }, + }, + { + label: `Workspace`, + iconPath: new ThemeIcon("account"), + id: "WORKSPACE", + collapsibleState: TreeItemCollapsibleState.Expanded, + contextValue: "workspace", + url: this.connectionManager.databricksWorkspace?.host?.toString(), + } + ); if (cluster) { let contextValue: @@ -362,6 +404,23 @@ export class ConfigurationDataProvider return children; } + if (element.id === "BUNDLE-TARGET") { + if (this.configModel.target === undefined) { + return []; + } else { + return [ + { + label: "Host", + description: await this.configModel.get("host"), + }, + { + label: "Mode", + description: await this.configModel.get("mode"), + }, + ]; + } + } + return []; } diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index c2c2ef100..43930f352 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -104,7 +104,7 @@ export class ConnectionCommands implements Disposable { */ attachClusterCommand() { return async (cluster: Cluster) => { - await this.connectionManager.attachCluster(cluster); + await this.connectionManager.attachCluster(cluster.id); }; } @@ -168,7 +168,7 @@ export class ConnectionCommands implements Disposable { const selectedItem = quickPick.selectedItems[0]; if ("cluster" in selectedItem) { const cluster = selectedItem.cluster; - await this.connectionManager.attachCluster(cluster); + await this.connectionManager.attachCluster(cluster.id); } else { await UrlUtils.openExternal( `${ diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index f7c20d363..e3f7600f4 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -1,30 +1,21 @@ import {WorkspaceClient, ApiClient, logging} from "@databricks/databricks-sdk"; import {Cluster, WorkspaceFsEntity, WorkspaceFsUtils} from "../sdk-extensions"; -import { - env, - EventEmitter, - Uri, - window, - workspace as vscodeWorkspace, -} from "vscode"; +import {EventEmitter, Uri, window, Disposable} from "vscode"; import {CliWrapper} from "../cli/CliWrapper"; import { SyncDestinationMapper, RemoteUri, LocalUri, } from "../sync/SyncDestination"; -import { - ConfigFileError, - ProjectConfig, - ProjectConfigFile, -} from "../file-managers/ProjectConfigFile"; import {configureWorkspaceWizard} from "./configureWorkspaceWizard"; import {ClusterManager} from "../cluster/ClusterManager"; import {DatabricksWorkspace} from "./DatabricksWorkspace"; -import {Loggers} from "../logger"; import {CustomWhenContext} from "../vscode-objs/CustomWhenContext"; import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; -import {StateStorage} from "../vscode-objs/StateStorage"; +import {ConfigModel} from "./ConfigModel"; +import {onError} from "../utils/onErrorDecorator"; +import {AuthProvider} from "./auth/AuthProvider"; +import {Mutex} from "../locking"; // eslint-disable-next-line @typescript-eslint/naming-convention const {NamedLogger} = logging; @@ -36,11 +27,14 @@ export type ConnectionState = "CONNECTED" | "CONNECTING" | "DISCONNECTED"; * It's responsible for reading and validating the project configuration * and for providing instances of the APIClient, Cluster and Workspace classes */ -export class ConnectionManager { +export class ConnectionManager implements Disposable { + private disposables: Disposable[] = []; + private _state: ConnectionState = "DISCONNECTED"; + private loginLogoutMutex: Mutex = new Mutex(); + private _workspaceClient?: WorkspaceClient; private _syncDestinationMapper?: SyncDestinationMapper; - private _projectConfigFile?: ProjectConfigFile; private _clusterManager?: ClusterManager; private _databricksWorkspace?: DatabricksWorkspace; @@ -60,10 +54,94 @@ export class ConnectionManager { public metadataServiceUrl?: string; + private async updateSyncDestinationMapper() { + const workspacePath = await this.configModel.get("workspaceFsPath"); + const remoteUri = workspacePath + ? new RemoteUri(workspacePath) + : undefined; + if (remoteUri?.path === this._syncDestinationMapper?.remoteUri.path) { + return; + } + if (remoteUri === undefined) { + this._syncDestinationMapper = undefined; + this.onDidChangeSyncDestinationEmitter.fire( + this._syncDestinationMapper + ); + return; + } + if (!(await remoteUri.validate(this))) { + window.showErrorMessage( + `Invalid sync destination ${workspacePath}` + ); + this.detachSyncDestination(); + return; + } + this._syncDestinationMapper = new SyncDestinationMapper( + new LocalUri(this.workspaceUri), + remoteUri + ); + } + + private async updateClusterManager() { + const clusterId = await this.configModel.get("clusterId"); + if (clusterId === this._clusterManager?.cluster?.id) { + return; + } + this._clusterManager?.dispose(); + this._clusterManager = + clusterId !== undefined && this.apiClient !== undefined + ? new ClusterManager( + await Cluster.fromClusterId(this.apiClient, clusterId), + () => { + this.onDidChangeClusterEmitter.fire(this.cluster); + } + ) + : undefined; + this.onDidChangeClusterEmitter.fire(this.cluster); + } + constructor( private cli: CliWrapper, - private stateStorage: StateStorage - ) {} + private readonly configModel: ConfigModel, + private readonly workspaceUri: Uri + ) { + this.disposables.push( + this.configModel.onDidChange( + "workspaceFsPath", + this.updateSyncDestinationMapper, + this + ), + this.configModel.onDidChange( + "clusterId", + this.updateClusterManager, + this + ), + this.configModel.onDidChange( + "target", + this.loginWithSavedAuth, + this + ), + this.configModel.onDidChange("authParams", async () => { + const config = await this.configModel.getS("authParams"); + if (config === undefined) { + return; + } + + // We only react to auth changes coming from the bundle. + // If an override is set, then all settings must have gone + // through this class, which means that we have already logged in + // using those settings, so we don't double login. + if (config.source === "bundle") { + await this.loginWithSavedAuth(); + } + }) + ); + } + + public async init() { + await this.configModel.init(); + await this.loginWithSavedAuth(); + } get state(): ConnectionState { return this._state; @@ -98,126 +176,67 @@ export class ConnectionManager { return this._workspaceClient?.apiClient; } - async login(interactive = false, force = false): Promise { - try { - await this._login(interactive, force); - } catch (e) { - NamedLogger.getOrCreate("Extension").error("Login Error:", e); - if (interactive && e instanceof Error) { - window.showErrorMessage(`Login error: ${e.message}`); - } - this.updateState("DISCONNECTED"); + @onError({ + popup: {prefix: "Can't login with saved auth. "}, + }) + private async loginWithSavedAuth() { + const authParams = await this.configModel.get("authParams"); + if (authParams === undefined) { await this.logout(); + return; } + await this.login(AuthProvider.fromJSON(authParams, this.cli.cliPath)); } - private async _login(interactive = false, force = false): Promise { + private async login( + authProvider: AuthProvider, + force = false + ): Promise { if (force) { await this.logout(); } if (this.state !== "DISCONNECTED") { return; } - this.updateState("CONNECTING"); - - let projectConfigFile: ProjectConfigFile; - let workspaceClient: WorkspaceClient; + if (!(await authProvider.check())) { + // We return without any state changes. This ensures that + // if users move from a working auth type to an invalid + // auth, the old auth will continue working and they will not + // have to re authenticate even the old one. + return; + } try { - try { - projectConfigFile = await ProjectConfigFile.load( - vscodeWorkspace.rootPath!, - this.cli.cliPath - ); - } catch (e) { - if ( - e instanceof ConfigFileError && - e.message.startsWith("Project config file does not exist") - ) { - this.updateState("DISCONNECTED"); - await this.logout(); - return; - } else { - throw e; - } - } + await this.loginLogoutMutex.synchronise(async () => { + this.updateState("CONNECTING"); - if (!(await projectConfigFile.authProvider.check())) { - throw new Error( - `Can't login with ${projectConfigFile.authProvider.describe()}.` + const databricksWorkspace = await DatabricksWorkspace.load( + authProvider.getWorkspaceClient(), + authProvider ); - } - - workspaceClient = - projectConfigFile.authProvider.getWorkspaceClient(); + this._databricksWorkspace = databricksWorkspace; + this._workspaceClient = authProvider.getWorkspaceClient(); - await workspaceClient.config.authenticate(new Headers()); - - this._databricksWorkspace = await DatabricksWorkspace.load( - workspaceClient, - projectConfigFile.authProvider - ); - } catch (e: any) { - const message = `Can't login to Databricks: ${e.message}`; - NamedLogger.getOrCreate("Extension").error(message, e); - window.showErrorMessage(message); + await this._databricksWorkspace.optionalEnableFilesInReposPopup( + this._workspaceClient + ); - this.updateState("DISCONNECTED"); - await this.logout(); - return; - } + await this.createWsFsRootDirectory(this._workspaceClient); + await this.updateSyncDestinationMapper(); + await this.updateClusterManager(); + await this.configModel.set("authParams", authProvider.toJSON()); - if ( - workspaceConfigs.syncDestinationType === "repo" && - (!this._databricksWorkspace.isReposEnabled || - !this._databricksWorkspace.isFilesInReposEnabled) - ) { - let message = ""; - if (!this._databricksWorkspace.isReposEnabled) { - message = - "Repos are not enabled for this workspace. Please enable it in the Databricks UI."; - } else if (!this._databricksWorkspace.isFilesInReposEnabled) { - message = - "Files in Repos is not enabled for this workspace. Please enable it in the Databricks UI."; - } - NamedLogger.getOrCreate("Extension").error(message); - if (interactive) { - const result = await window.showWarningMessage( - message, - "Open Databricks UI" + this.updateState("CONNECTED"); + }); + } catch (e) { + NamedLogger.getOrCreate("Extension").error(`Login failed`, e); + if (e instanceof Error) { + await window.showWarningMessage( + `Login failed with error: "${e.message}"."` ); - if (result === "Open Databricks UI") { - const host = await workspaceClient.apiClient.host; - await env.openExternal( - Uri.parse( - host.toString() + - "#setting/accounts/workspace-settings" - ) - ); - } } + await this.logout(); } - - this._workspaceClient = workspaceClient; - this._projectConfigFile = projectConfigFile; - - if (projectConfigFile.clusterId) { - await this.attachCluster(projectConfigFile.clusterId, true); - } else { - this.updateCluster(undefined); - } - - if (projectConfigFile.workspacePath) { - await this.attachSyncDestination( - new RemoteUri(projectConfigFile.workspacePath), - true - ); - } else { - this.updateSyncDestination(undefined); - } - - await this.createWsFsRootDirectory(workspaceClient); - this.updateState("CONNECTED"); } async createWsFsRootDirectory(wsClient: WorkspaceClient) { @@ -249,239 +268,77 @@ export class ConnectionManager { } } + @onError({popup: {prefix: "Can't logout. "}}) + @Mutex.synchronise("loginLogoutMutex") async logout() { - if (this._state === "DISCONNECTED") { - return; - } else { - await this.waitForConnect(); - } - - this._projectConfigFile = undefined; this._workspaceClient = undefined; this._databricksWorkspace = undefined; - this.updateCluster(undefined); - this.updateSyncDestination(undefined); + await this.updateClusterManager(); + await this.updateSyncDestinationMapper(); + await this.configModel.set("authParams", undefined); this.updateState("DISCONNECTED"); } + @onError({ + popup: {prefix: "Can't configure workspace. "}, + }) async configureWorkspace() { - let config: ProjectConfig | undefined; - while (true) { - config = await configureWorkspaceWizard( - this.cli, - this.databricksWorkspace?.host?.toString() || - config?.authProvider?.host.toString() - ); - - if (!config) { - return; - } - - if (!(await config.authProvider.check())) { - return; - } - - try { - const workspaceClient = - config.authProvider.getWorkspaceClient(); - - await workspaceClient.config.authenticate(new Headers()); - - await DatabricksWorkspace.load( - workspaceClient, - config.authProvider - ); - } catch (e: any) { - NamedLogger.getOrCreate("Extension").error( - `Connection using "${config.authProvider.describe()}" failed`, - e - ); - - const response = await window.showWarningMessage( - `Connection using "${config.authProvider.describe()}" failed with error: "${ - e.message - }"."`, - "Retry", - "Cancel" - ); - - switch (response) { - case "Retry": - continue; - - case "Cancel": - return; - } - } - - break; + const host = await this.configModel.get("host"); + if (host === undefined) { + return; } - await this.writeConfigFile(config); - window.showInformationMessage( - `connected to: ${config.authProvider.host}` - ); - - await this.login(true, true); - } - - private async writeConfigFile(config: ProjectConfig) { - const projectConfigFile = new ProjectConfigFile( - config, - vscodeWorkspace.rootPath!, - this.cli.cliPath - ); + const config = await configureWorkspaceWizard(this.cli, host); + if (!config) { + return; + } - await projectConfigFile.write(); + await this.login(config.authProvider); } - async attachCluster( - cluster: Cluster | string, - skipWrite = false - ): Promise { - try { - if (typeof cluster === "string") { - cluster = await Cluster.fromClusterId( - this._workspaceClient!.apiClient, - cluster - ); - } - - if ( - JSON.stringify(this.cluster?.details) === - JSON.stringify(cluster.details) - ) { - return; - } - - if (!skipWrite) { - this._projectConfigFile!.clusterId = cluster.id; - await this._projectConfigFile!.write(); - } - - if (cluster.state === "RUNNING") { - 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); - } catch (e) { - NamedLogger.getOrCreate("Extension").error( - "Attach Cluster error", - e - ); - window.showErrorMessage( - `Error in attaching cluster ${ - typeof cluster === "string" ? cluster : cluster.id - }` - ); - await this.detachCluster(); + @onError({ + popup: {prefix: "Can't attach cluster. "}, + }) + async attachCluster(clusterId: string): Promise { + if (this.cluster?.id === clusterId) { + return; } + await this.configModel.set("clusterId", clusterId); } + @onError({ + popup: {prefix: "Can't detach cluster. "}, + }) async detachCluster(): Promise { - if (!this.cluster && this._projectConfigFile?.clusterId === undefined) { - return; - } - - if (this._projectConfigFile) { - this._projectConfigFile.clusterId = undefined; - await this._projectConfigFile.write(); - } - - this.updateCluster(undefined); + await this.configModel.set("clusterId", undefined); } - async attachSyncDestination( - remoteWorkspace: RemoteUri, - skipWrite = false - ): Promise { - try { - if ( - !vscodeWorkspace.workspaceFolders || - !vscodeWorkspace.workspaceFolders.length - ) { - // TODO how do we handle this? - return; - } - if ( - this.workspaceClient === undefined || - this.databricksWorkspace === undefined - ) { - throw new Error( - "Can't attach a Sync Destination when profile is not connected" - ); - } - if (!(await remoteWorkspace.validate(this))) { - await this.detachSyncDestination(); - return; - } - - if (!skipWrite) { - this._projectConfigFile!.workspacePath = remoteWorkspace.uri; - await this._projectConfigFile!.write(); - } - - const wsUri = vscodeWorkspace.workspaceFolders[0].uri; - this.updateSyncDestination( - new SyncDestinationMapper(new LocalUri(wsUri), remoteWorkspace) - ); - } catch (e) { - NamedLogger.getOrCreate("Extension").error( - "Attach Sync Destination error", - e - ); + @onError({ + popup: {prefix: "Can't attach sync destination. "}, + }) + async attachSyncDestination(remoteWorkspace: RemoteUri): Promise { + if (!(await remoteWorkspace.validate(this))) { + await this.detachSyncDestination(); window.showErrorMessage( - `Error in attaching sync destination ${remoteWorkspace.path}` + `Can't attach sync destination ${remoteWorkspace.path}` ); - await this.detachSyncDestination(); + return; } + await this.configModel.set("workspaceFsPath", remoteWorkspace.path); } + @onError({ + popup: {prefix: "Can't detach sync destination. "}, + }) async detachSyncDestination(): Promise { - if ( - !this._syncDestinationMapper && - this._projectConfigFile?.workspacePath === undefined - ) { - return; - } - - if (this._projectConfigFile) { - this._projectConfigFile.workspacePath = undefined; - await this._projectConfigFile.write(); - } - - this.updateSyncDestination(undefined); + await this.configModel.set("workspaceFsPath", undefined); } private updateState(newState: ConnectionState) { - if (newState === "DISCONNECTED") { - this._databricksWorkspace = undefined; + if (!this.loginLogoutMutex.locked) { + throw new Error( + "updateState must be called after aquireing the state mutex" + ); } if (this._state !== newState) { this._state = newState; @@ -490,29 +347,6 @@ export class ConnectionManager { CustomWhenContext.setLoggedIn(this._state === "CONNECTED"); } - private updateCluster(newCluster: Cluster | undefined) { - if (this.cluster !== newCluster) { - this._clusterManager?.dispose(); - this._clusterManager = newCluster - ? new ClusterManager(newCluster, () => { - this.onDidChangeClusterEmitter.fire(this.cluster); - }) - : undefined; - this.onDidChangeClusterEmitter.fire(this.cluster); - } - } - - private updateSyncDestination( - newSyncDestination: SyncDestinationMapper | undefined - ) { - if (this._syncDestinationMapper !== newSyncDestination) { - this._syncDestinationMapper = newSyncDestination; - this.onDidChangeSyncDestinationEmitter.fire( - this._syncDestinationMapper - ); - } - } - async startCluster() { await this._clusterManager?.start(() => { this.onDidChangeClusterEmitter.fire(this.cluster); @@ -537,4 +371,8 @@ export class ConnectionManager { }); } } + + async dispose() { + this.disposables.forEach((d) => d.dispose()); + } } diff --git a/packages/databricks-vscode/src/configuration/DatabricksWorkspace.ts b/packages/databricks-vscode/src/configuration/DatabricksWorkspace.ts index 37fe410ec..7cf5f51e5 100644 --- a/packages/databricks-vscode/src/configuration/DatabricksWorkspace.ts +++ b/packages/databricks-vscode/src/configuration/DatabricksWorkspace.ts @@ -1,7 +1,7 @@ import {WorkspaceClient, iam, logging} from "@databricks/databricks-sdk"; import {Cluster, WorkspaceConf, WorkspaceConfProps} from "../sdk-extensions"; import {Context, context} from "@databricks/databricks-sdk/dist/context"; -import {Uri} from "vscode"; +import {Uri, window, env} from "vscode"; import {Loggers} from "../logger"; import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import {AuthProvider} from "./auth/AuthProvider"; @@ -117,4 +117,35 @@ export class DatabricksWorkspace { return new DatabricksWorkspace(authProvider, me, state); } + + public async optionalEnableFilesInReposPopup( + workspaceClient: WorkspaceClient + ) { + if ( + workspaceConfigs.syncDestinationType === "repo" && + (!this.isReposEnabled || !this.isFilesInReposEnabled) + ) { + let message = ""; + if (!this.isReposEnabled) { + message = + "Repos are not enabled for this workspace. Please enable it in the Databricks UI."; + } else if (!this.isFilesInReposEnabled) { + message = + "Files in Repos is not enabled for this workspace. Please enable it in the Databricks UI."; + } + logging.NamedLogger.getOrCreate("Extension").error(message); + const result = await window.showWarningMessage( + message, + "Open Databricks UI" + ); + if (result === "Open Databricks UI") { + const host = await workspaceClient.apiClient.host; + await env.openExternal( + Uri.parse( + host.toString() + "#setting/accounts/workspace-settings" + ) + ); + } + } + } } diff --git a/packages/databricks-vscode/src/configuration/configureWorkspaceWizard.ts b/packages/databricks-vscode/src/configuration/configureWorkspaceWizard.ts index 9d14f471a..01b9ae089 100644 --- a/packages/databricks-vscode/src/configuration/configureWorkspaceWizard.ts +++ b/packages/databricks-vscode/src/configuration/configureWorkspaceWizard.ts @@ -9,7 +9,6 @@ import { } from "../utils/urlUtils"; import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import {AuthProvider, AuthType} from "./auth/AuthProvider"; -import {ProjectConfig} from "../file-managers/ProjectConfigFile"; interface AuthTypeQuickPickItem extends QuickPickItem { authType: AuthType | "new-profile" | "none"; @@ -27,7 +26,7 @@ interface State { export async function configureWorkspaceWizard( cliWrapper: CliWrapper, host?: string -): Promise { +): Promise<{authProvider: AuthProvider} | undefined> { const title = "Configure Databricks Workspace"; async function collectInputs(): Promise { @@ -42,10 +41,8 @@ export async function configureWorkspaceWizard( const items: Array = []; if (state.host) { - items.push({ - label: state.host.toString(), - detail: "Currently selected host", - }); + state.host = normalizeHost(state.host.toString()); + return (input: MultiStepInput) => selectAuthMethod(input, state); } if (process.env.DATABRICKS_HOST) { diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 02458d082..397d9cd4e 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -18,7 +18,6 @@ import {DatabricksDebugAdapterFactory} from "./run/DatabricksDebugAdapter"; import {DatabricksWorkflowDebugAdapterFactory} from "./run/DatabricksWorkflowDebugAdapter"; import {SyncCommands} from "./sync/SyncCommands"; import {CodeSynchronizer} from "./sync/CodeSynchronizer"; -import {ProjectConfigFileWatcher} from "./file-managers/ProjectConfigFileWatcher"; import {QuickstartCommands} from "./quickstart/QuickstartCommands"; import {showQuickStartOnFirstUse} from "./quickstart/QuickStart"; import {PublicApi} from "@databricks/databricks-vscode-types"; @@ -55,6 +54,9 @@ import { registerBundleAutocompleteProvider, } from "./bundle"; import {showWhatsNewPopup} from "./whatsNewPopup"; +import {ConfigModel} from "./configuration/ConfigModel"; +import {ConfigOverrideReaderWriter} from "./configuration/ConfigOverrideReaderWriter"; +import {BundleConfigReaderWriter} from "./configuration/BundleConfigReaderWriter"; export async function activate( context: ExtensionContext @@ -89,6 +91,7 @@ export async function activate( return undefined; } + const workspaceUri = workspace.workspaceFolders[0].uri; const stateStorage = new StateStorage(context); // Add the databricks binary to the PATH environment variable in terminals @@ -146,9 +149,28 @@ export async function activate( workspace.onDidChangeConfiguration(updateFeatureContexts) ); + const bundleFileSet = new BundleFileSet(workspace.workspaceFolders[0].uri); + const bundleFileWatcher = new BundleWatcher(bundleFileSet); + context.subscriptions.push(bundleFileWatcher); + + const overrideReaderWriter = new ConfigOverrideReaderWriter(stateStorage); + const bundleConfigReaderWriter = new BundleConfigReaderWriter( + bundleFileSet + ); + const confgiModel = new ConfigModel( + overrideReaderWriter, + bundleConfigReaderWriter, + stateStorage, + bundleFileWatcher + ); + // Configuration group const cli = new CliWrapper(context); - const connectionManager = new ConnectionManager(cli, stateStorage); + const connectionManager = new ConnectionManager( + cli, + confgiModel, + workspaceUri + ); context.subscriptions.push( connectionManager.onDidChangeState(async (state) => { telemetry.setMetadata( @@ -344,7 +366,8 @@ export async function activate( stateStorage, wsfsAccessVerifier, featureManager, - telemetry + telemetry, + confgiModel ); context.subscriptions.push( @@ -496,14 +519,6 @@ export async function activate( ) ); - context.subscriptions.push( - new ProjectConfigFileWatcher( - connectionManager, - workspace.rootPath!, - cli.cliPath - ) - ); - // Quickstart const quickstartCommands = new QuickstartCommands(context); context.subscriptions.push( @@ -535,11 +550,6 @@ export async function activate( } }) ); - - const bundleFileSet = new BundleFileSet(workspace.workspaceFolders[0].uri); - const bundleFileWatcher = new BundleWatcher(bundleFileSet); - context.subscriptions.push(bundleFileWatcher); - // generate a json schema for bundle root and load a custom provider into // redhat.vscode-yaml extension to validate bundle config files with this schema registerBundleAutocompleteProvider( @@ -554,13 +564,6 @@ export async function activate( ); }); - connectionManager.login(false).catch((e) => { - logging.NamedLogger.getOrCreate(Loggers.Extension).error( - "Login error", - e - ); - }); - setDbnbCellLimits( workspace.workspaceFolders[0].uri, connectionManager @@ -585,6 +588,15 @@ export async function activate( ); }); + context.subscriptions.push( + commands.registerCommand("databricks.internal.clearOverrides", () => { + stateStorage.set("databricks.bundle.overrides", undefined); + }) + ); + + connectionManager.init().catch((e) => { + window.showErrorMessage(e); + }); CustomWhenContext.setActivated(true); telemetry.recordEvent(Events.EXTENSION_ACTIVATED); diff --git a/packages/databricks-vscode/src/file-managers/ProjectConfigFile.test.ts b/packages/databricks-vscode/src/file-managers/ProjectConfigFile.test.ts deleted file mode 100644 index 2bfb37d21..000000000 --- a/packages/databricks-vscode/src/file-managers/ProjectConfigFile.test.ts +++ /dev/null @@ -1,118 +0,0 @@ -/* eslint-disable @typescript-eslint/naming-convention */ - -import {mkdir, mkdtemp, readFile, writeFile} from "fs/promises"; -import {ProjectConfig, ProjectConfigFile} from "./ProjectConfigFile"; -import * as assert from "assert"; -import path from "path"; -import * as os from "os"; -import {ProfileAuthProvider} from "../configuration/auth/AuthProvider"; -import {Uri} from "vscode"; - -describe(__filename, () => { - let tempDir: string; - before(async () => { - tempDir = await mkdtemp(path.join(os.tmpdir(), "ProjectConfTests-")); - }); - - it("should write config file", async () => { - const authProvider = new ProfileAuthProvider( - new URL("https://000000000000.00.azuredatabricks.net/"), - "testProfile" - ); - const expected: ProjectConfig = { - authProvider: authProvider, - clusterId: "testClusterId", - workspacePath: Uri.from({scheme: "wsfs", path: "workspacePath"}), - }; - await new ProjectConfigFile(expected, tempDir, "databricks").write(); - - const rawData = await readFile( - ProjectConfigFile.getProjectConfigFilePath(tempDir), - {encoding: "utf-8"} - ); - const actual = JSON.parse(rawData); - assert.deepEqual(actual, { - host: "https://000000000000.00.azuredatabricks.net/", - authType: "profile", - profile: "testProfile", - workspacePath: "workspacePath", - clusterId: "testClusterId", - }); - }); - - it("should load config file", async () => { - const configFile = ProjectConfigFile.getProjectConfigFilePath(tempDir); - await mkdir(path.dirname(configFile), {recursive: true}); - - const config = { - host: "https://000000000000.00.azuredatabricks.net/", - authType: "profile", - profile: "testProfile", - workspacePath: "workspacePath", - clusterId: "testClusterId", - }; - await writeFile(configFile, JSON.stringify(config), { - encoding: "utf-8", - }); - - const actual = await ProjectConfigFile.load(tempDir, "databricks"); - assert.equal(actual.host.toString(), config.host); - assert.ok(actual.authProvider instanceof ProfileAuthProvider); - assert.equal(actual.authProvider.authType, config.authType); - assert.deepStrictEqual(actual.authProvider.toJSON(), { - host: config.host.toString(), - authType: config.authType, - profile: config.profile, - }); - assert.deepEqual( - actual.workspacePath, - Uri.from({scheme: "wsfs", path: config.workspacePath}) - ); - assert.equal(actual.clusterId, config.clusterId); - }); - - it("should load old config file format", async () => { - const configFile = ProjectConfigFile.getProjectConfigFilePath(tempDir); - await mkdir(path.dirname(configFile), {recursive: true}); - - const config = { - profile: "testProfile", - workspacePath: "workspacePath", - clusterId: "testClusterId", - }; - await writeFile(configFile, JSON.stringify(config), { - encoding: "utf-8", - }); - - await writeFile( - path.join(tempDir, ".databrickscfg"), - `[testProfile] -host = https://000000000000.00.azuredatabricks.net/ -token = testToken`, - { - encoding: "utf-8", - } - ); - - process.env.DATABRICKS_CONFIG_FILE = path.join( - tempDir, - ".databrickscfg" - ); - const actual = await ProjectConfigFile.load(tempDir, "databricks"); - assert.equal( - actual.host.toString(), - "https://000000000000.00.azuredatabricks.net/" - ); - assert.ok(actual.authProvider instanceof ProfileAuthProvider); - assert.deepStrictEqual(actual.authProvider.toJSON(), { - host: "https://000000000000.00.azuredatabricks.net/", - authType: "profile", - profile: config.profile, - }); - assert.deepEqual( - actual.workspacePath, - Uri.from({scheme: "wsfs", path: config.workspacePath}) - ); - assert.equal(actual.clusterId, config.clusterId); - }); -}); diff --git a/packages/databricks-vscode/src/file-managers/ProjectConfigFile.ts b/packages/databricks-vscode/src/file-managers/ProjectConfigFile.ts deleted file mode 100644 index 6d7ea89fa..000000000 --- a/packages/databricks-vscode/src/file-managers/ProjectConfigFile.ts +++ /dev/null @@ -1,163 +0,0 @@ -import path from "node:path"; -import fs from "node:fs/promises"; -import { - AuthProvider, - ProfileAuthProvider, -} from "../configuration/auth/AuthProvider"; -import {Uri} from "vscode"; -import {Loggers} from "../logger"; -import {Config, logging} from "@databricks/databricks-sdk"; -import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; - -export interface ProjectConfig { - authProvider: AuthProvider; - clusterId?: string; - workspacePath?: Uri; -} - -export class ConfigFileError extends Error {} - -export class ProjectConfigFile { - constructor( - private config: ProjectConfig, - readonly rootPath: string, - readonly cliPath: string - ) {} - - get host() { - return this.config.authProvider.host; - } - - get authProvider() { - return this.config.authProvider; - } - - get clusterId() { - return this.config.clusterId; - } - - set clusterId(clusterId: string | undefined) { - this.config.clusterId = clusterId; - } - - get workspacePath(): Uri | undefined { - return this.config.workspacePath; - } - - set workspacePath(workspacePath: Uri | undefined) { - this.config.workspacePath = workspacePath; - } - - toJSON(): Record { - return { - ...this.config.authProvider.toJSON(), - clusterId: this.clusterId, - workspacePath: this.workspacePath?.path, - }; - } - - async write() { - try { - const originalConfig = await ProjectConfigFile.load( - this.rootPath, - this.cliPath - ); - if ( - JSON.stringify(originalConfig.toJSON(), null, 2) === - JSON.stringify(this.toJSON(), null, 2) - ) { - return; - } - } catch (e) {} - - const fileName = ProjectConfigFile.getProjectConfigFilePath( - this.rootPath - ); - await fs.mkdir(path.dirname(fileName), {recursive: true}); - - await fs.writeFile(fileName, JSON.stringify(this, null, 2), { - encoding: "utf-8", - }); - } - - static async importOldConfig(config: any): Promise { - const sdkConfig = new Config({ - profile: config.profile, - configFile: - workspaceConfigs.databrickscfgLocation ?? - process.env.DATABRICKS_CONFIG_FILE, - env: {}, - }); - - await sdkConfig.ensureResolved(); - - return new ProfileAuthProvider( - new URL(sdkConfig.host!), - sdkConfig.profile! - ); - } - - static async load( - rootPath: string, - cliPath: string - ): Promise { - const projectConfigFilePath = this.getProjectConfigFilePath(rootPath); - - let rawConfig; - try { - rawConfig = await fs.readFile(projectConfigFilePath, { - encoding: "utf-8", - }); - } catch (e: any) { - if (e.code && e.code === "ENOENT") { - throw new ConfigFileError( - `Project config file does not exist: ${projectConfigFilePath}` - ); - } else { - throw e; - } - } - - let authProvider: AuthProvider; - let config: any; - try { - config = JSON.parse(rawConfig); - if (!config.authType && config.profile) { - authProvider = await this.importOldConfig(config); - } else { - authProvider = AuthProvider.fromJSON(config, cliPath); - } - } catch (e: any) { - logging.NamedLogger.getOrCreate(Loggers.Extension).error( - "Error parsing project config file", - e - ); - throw new ConfigFileError( - `Error parsing project config file: ${e.message}` - ); - } - return new ProjectConfigFile( - { - authProvider: authProvider!, - clusterId: config.clusterId, - workspacePath: - config.workspacePath !== undefined - ? Uri.from({ - scheme: "wsfs", - path: config.workspacePath, - }) - : undefined, - }, - rootPath, - cliPath - ); - } - - static getProjectConfigFilePath(rootPath?: string): string { - if (!rootPath) { - throw new Error("Not in a VSCode workspace"); - } - const cwd = path.normalize(rootPath); - return path.join(cwd, ".databricks", "project.json"); - } -} diff --git a/packages/databricks-vscode/src/file-managers/ProjectConfigFileWatcher.ts b/packages/databricks-vscode/src/file-managers/ProjectConfigFileWatcher.ts deleted file mode 100644 index 7f4e8287a..000000000 --- a/packages/databricks-vscode/src/file-managers/ProjectConfigFileWatcher.ts +++ /dev/null @@ -1,78 +0,0 @@ -import {Disposable, workspace} from "vscode"; -import {ConnectionManager} from "../configuration/ConnectionManager"; -import {ProjectConfigFile} from "./ProjectConfigFile"; -import {RemoteUri} from "../sync/SyncDestination"; - -export class ProjectConfigFileWatcher implements Disposable { - private disposables: Array = []; - constructor( - readonly connectionManager: ConnectionManager, - rootPath: string, - cliPath: string - ) { - const fileSystemWatcher = workspace.createFileSystemWatcher( - ProjectConfigFile.getProjectConfigFilePath(rootPath) - ); - - this.disposables.push( - fileSystemWatcher, - fileSystemWatcher.onDidCreate(async () => { - switch (this.connectionManager.state) { - case "DISCONNECTED": - await this.connectionManager.login(); - break; - case "CONNECTING": - await this.connectionManager.waitForConnect(); - break; - case "CONNECTED": - return; - } - }, this), - fileSystemWatcher.onDidChange(async () => { - const configFile = await ProjectConfigFile.load( - rootPath, - cliPath - ); - if (this.connectionManager.state === "CONNECTING") { - await this.connectionManager.waitForConnect(); - } - if ( - configFile.host.toString() !== - connectionManager.databricksWorkspace?.host.toString() || - configFile.authProvider.authType !== - connectionManager.databricksWorkspace?.authProvider - .authType - ) { - await connectionManager.login(false, true); - } - if (connectionManager.cluster?.id !== configFile.clusterId) { - if (configFile.clusterId) { - await connectionManager.attachCluster( - configFile.clusterId - ); - } else { - await connectionManager.detachCluster(); - } - } - if ( - connectionManager.syncDestinationMapper?.remoteUri.path !== - configFile.workspacePath?.path - ) { - if (configFile.workspacePath) { - await connectionManager.attachSyncDestination( - new RemoteUri(configFile.workspacePath?.path) - ); - } else { - await connectionManager.detachSyncDestination(); - } - } - }, this), - fileSystemWatcher.onDidDelete(async () => { - await connectionManager.logout(); - }, this) - ); - } - dispose() { - this.disposables.forEach((item) => item.dispose()); - } -} diff --git a/packages/databricks-vscode/src/locking/Mutex.ts b/packages/databricks-vscode/src/locking/Mutex.ts index 3c64af3e8..e886da260 100644 --- a/packages/databricks-vscode/src/locking/Mutex.ts +++ b/packages/databricks-vscode/src/locking/Mutex.ts @@ -25,10 +25,10 @@ export class Mutex { return this._locked; } - async synchronise(fn: () => Promise) { + async synchronise(fn: () => Promise) { await this.wait(); try { - await fn(); + return await fn(); } finally { this.signal(); } @@ -43,12 +43,9 @@ export class Mutex { const original = descriptor.value; descriptor.value = async function (...args: any[]) { const mutex = (this as any)[mutexKey] as Mutex; - await mutex.wait(); - try { + mutex.synchronise(async () => { return await original.apply(this, args); - } finally { - mutex.signal(); - } + }); }; }; } diff --git a/packages/databricks-vscode/src/utils/onErrorDecorator.ts b/packages/databricks-vscode/src/utils/onErrorDecorator.ts new file mode 100644 index 000000000..982614b25 --- /dev/null +++ b/packages/databricks-vscode/src/utils/onErrorDecorator.ts @@ -0,0 +1,81 @@ +import {window} from "vscode"; +import { + getContextParamIndex, + logging, + Context, +} from "@databricks/databricks-sdk"; +import {Loggers} from "../logger"; + +interface Props { + popup?: + | { + prefix?: string; + } + | true; + log?: + | { + prefix?: string; + } + | true; +} + +const defaultProps: Props = { + log: true, +}; + +export function onError(props: Props) { + props = {...defaultProps, ...props}; + return function showErrorDecorator( + target: any, + propertyKey: string, + descriptor: PropertyDescriptor + ) { + let contextParamIndex: number = -1; + // Find the @context if it exists. We want to use it for logging whenever possible. + if (props.log !== undefined) { + try { + logging.withLogContext(Loggers.Extension)( + target, + propertyKey, + descriptor + ); + contextParamIndex = getContextParamIndex(target, propertyKey); + } catch (e) {} + } + const originalMethod = descriptor.value; + + descriptor.value = async function (...args: any[]) { + try { + return await originalMethod.apply(this, args); + } catch (e) { + if (!(e instanceof Error)) { + throw e; + } + + let prefix = ""; + if (props.popup !== undefined) { + prefix = + typeof props.popup !== "boolean" + ? props.popup.prefix ?? "" + : ""; + window.showErrorMessage(prefix + e.message); + } + if (props.log !== undefined) { + // If we do not have a context, we create a new logger. + const logger = + contextParamIndex !== -1 + ? (args[contextParamIndex] as Context).logger + : logging.NamedLogger.getOrCreate( + Loggers.Extension + ); + prefix = + props.log === true + ? prefix + : props.log.prefix ?? prefix; + logger?.error(prefix + e.message, e); + } + return; + } + }; + }; +}