diff --git a/packages/databricks-vscode/src/feature-manager/FeatureManager.ts b/packages/databricks-vscode/src/feature-manager/FeatureManager.ts index 03e5ad783..39c37fc81 100644 --- a/packages/databricks-vscode/src/feature-manager/FeatureManager.ts +++ b/packages/databricks-vscode/src/feature-manager/FeatureManager.ts @@ -20,6 +20,7 @@ export interface FeatureStepState { available: boolean; title?: string; message?: string; + warning?: string; action?: FeatureEnableAction; isDisabledByFf?: boolean; } diff --git a/packages/databricks-vscode/src/feature-manager/MultiStepAccessVerfier.ts b/packages/databricks-vscode/src/feature-manager/MultiStepAccessVerfier.ts index adfbf3b64..d41b9ac2d 100644 --- a/packages/databricks-vscode/src/feature-manager/MultiStepAccessVerfier.ts +++ b/packages/databricks-vscode/src/feature-manager/MultiStepAccessVerfier.ts @@ -76,12 +76,13 @@ export abstract class MultiStepAccessVerifier implements Feature { }); } - acceptStep(id: string, title?: string, message?: string) { + acceptStep(id: string, title?: string, message?: string, warning?: string) { return this.updateStep({ id: id, available: true, title, message, + warning, }); } diff --git a/packages/databricks-vscode/src/language/EnvironmentCommands.ts b/packages/databricks-vscode/src/language/EnvironmentCommands.ts index 855be8952..e59c44994 100644 --- a/packages/databricks-vscode/src/language/EnvironmentCommands.ts +++ b/packages/databricks-vscode/src/language/EnvironmentCommands.ts @@ -1,4 +1,4 @@ -import {window} from "vscode"; +import {window, commands} from "vscode"; import {FeatureManager} from "../feature-manager/FeatureManager"; import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper"; import {Cluster} from "../sdk-extensions"; @@ -12,6 +12,7 @@ export class EnvironmentCommands { ) {} async setup(stepId?: string) { + commands.executeCommand("configurationView.focus"); await window.withProgress( {location: {viewId: "configurationView"}}, () => this._setup(stepId) @@ -24,6 +25,9 @@ export class EnvironmentCommands { true ); if (state.available) { + window.showInformationMessage( + "Python environment and Databricks Connect are already set up." + ); return true; } for (const [, step] of state.steps) { diff --git a/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts b/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts index 4c9a30a14..85cb4b1c9 100644 --- a/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts +++ b/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts @@ -161,38 +161,6 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier { async checkPythonEnvironment(): Promise { const env = await this.pythonExtension.pythonEnvironment; - const dbrVersionParts = - this.connectionManager.cluster?.dbrVersion || []; - // DBR 13 and 14 require python 3.10 - if ( - (dbrVersionParts[0] === 13 || dbrVersionParts[0] === 14) && - !this.matchEnvironmentVersion(env, 3, 10) - ) { - return this.rejectStep( - "checkPythonEnvironment", - "Activate an environment with Python 3.10", - `The python version should match DBR ${ - dbrVersionParts[0] - } requirements. ${this.printEnvironment(env)}`, - this.selectPythonInterpreter.bind(this) - ); - } - // DBR 15 requires python 3.11 - if ( - dbrVersionParts[0] === 15 && - !this.matchEnvironmentVersion(env, 3, 11) - ) { - return this.rejectStep( - "checkPythonEnvironment", - "Activate an environment with Python 3.11", - `The version should match DBR ${ - dbrVersionParts[0] - } requirements. ${this.printEnvironment(env)}`, - this.selectPythonInterpreter.bind(this) - ); - } - // If we don't know DBR version (no cluster is connected or new version is released and the extension isn't updated yet), - // we still check that environment is active and has python >= 3.10 const envVersionTooLow = env?.version && (env.version.major !== 3 || env.version.minor < 10); const noEnvironment = !env?.environment; @@ -215,10 +183,30 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier { this.selectPythonInterpreter.bind(this) ); } + const dbrVersionParts = + this.connectionManager.cluster?.dbrVersion || []; + let warning; + if ( + (dbrVersionParts[0] === 13 || dbrVersionParts[0] === 14) && + !this.matchEnvironmentVersion(env, 3, 10) + ) { + warning = `Use python 3.10 to match DBR ${ + dbrVersionParts[0] + } requirements. ${this.printEnvironment(env)}`; + } + if ( + dbrVersionParts[0] === 15 && + !this.matchEnvironmentVersion(env, 3, 11) + ) { + warning = `Use python 3.11 to match DBR ${ + dbrVersionParts[0] + } requirements. ${this.printEnvironment(env)}`; + } return this.acceptStep( "checkPythonEnvironment", `Active Environment: ${env.environment.name}`, - env.executable.uri?.fsPath + env.executable.uri?.fsPath, + warning ); } diff --git a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts index ffd3ba230..9a715fc24 100644 --- a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts +++ b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts @@ -8,7 +8,7 @@ import path from "node:path"; import {fileURLToPath} from "url"; import assert from "assert"; import fs from "fs/promises"; -import {ApiError, Config, WorkspaceClient} from "@databricks/databricks-sdk"; +import {Config, WorkspaceClient} from "@databricks/databricks-sdk"; import * as ElementCustomCommands from "./customCommands/elementCustomCommands.ts"; import {execFile as execFileCb} from "node:child_process"; import {cpSync, mkdirSync, rmSync} from "node:fs"; @@ -588,26 +588,44 @@ function getWorkspaceClient(config: Config) { async function startCluster( workspaceClient: WorkspaceClient, - clusterId: string + clusterId: string, + attempt = 0 ) { - console.log(`Starting cluster: ${clusterId}`); - - try { - await ( - await workspaceClient.clusters.start({ - cluster_id: clusterId, - }) - ).wait({ - onProgress: async (state) => { - console.log(`Cluster state: ${state.state}`); - }, - }); - } catch (e: unknown) { - if (!(e instanceof ApiError && e.message.includes("INVALID_STATE"))) { - throw e; - } - console.log(e.message); + console.log(`Cluster ID: ${clusterId}`); + if (attempt > 100) { + throw new Error("Failed to start the cluster: too many attempts"); + } + const cluster = await workspaceClient.clusters.get({ + cluster_id: clusterId, + }); + console.log(`Cluster State: ${cluster.state}`); + switch (cluster.state) { + case "RUNNING": + console.log("Cluster is already running"); + break; + case "TERMINATED": + case "ERROR": + case "UNKNOWN": + console.log("Starting the cluster..."); + await ( + await workspaceClient.clusters.start({ + cluster_id: clusterId, + }) + ).wait({ + onProgress: async (state) => { + console.log(`Cluster state: ${state.state}`); + }, + }); + break; + case "PENDING": + case "RESIZING": + case "TERMINATING": + case "RESTARTING": + console.log("Waiting and retrying..."); + await sleep(10000); + await startCluster(workspaceClient, clusterId, attempt + 1); + break; + default: + throw new Error(`Unknown cluster state: ${cluster.state}`); } - - console.log(`Cluster started`); } diff --git a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts index 2f7e7e717..2d210b79a 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts @@ -64,7 +64,7 @@ export class EnvironmentComponent extends BaseComponent { const environmentState = await this.featureManager.isEnabled( "environment.dependencies" ); - const children = []; + const children: ConfigurationTreeItem[] = []; for (const [id, step] of environmentState.steps) { if (!step.available) { children.push({ @@ -93,6 +93,13 @@ export class EnvironmentComponent extends BaseComponent { tooltip: step.message, iconPath: new ThemeIcon("check"), }); + if (step.warning) { + children.push({ + contextValue: getItemContext(id, true), + label: step.warning, + iconPath: new ThemeIcon("warning"), + }); + } } } return children;