Configure autocomplete as a part of env panel (#1554)

## Changes
"Configure autocomplete" logic was disconnected from the enc setup, and
it was still possible to install pyspark into the global python
installation.

Or re-install it, even if we already installed it before through env
panel.

Now "Setup Databricks builtins" appears in the env panel as an optional
step if all the previous env checks are green.

Also fix a bug where we were settig up builtins file in the bundle root,
while the python extension expects a workspace root.

<img width="417" alt="Screenshot 2025-02-12 at 11 11 13"
src="https://github.com/user-attachments/assets/3377a10c-b884-4a40-b6f8-0cb6fc53f139"
/>


<img width="516" alt="Screenshot 2025-02-12 at 11 11 24"
src="https://github.com/user-attachments/assets/8ea3800c-9a1f-4fc6-91aa-00758f67bd78"
/>




<!-- Summary of your changes that are easy to understand -->

Author: ilia-db

packages/databricks-vscode/src/extension.ts

@@ -357,7 +357,19 @@ export async function activate(
         )
     );
 
-    const clusterModel = new ClusterModel(connectionManager);
+    const configureAutocomplete = new ConfigureAutocomplete(
+        context,
+        stateStorage,
+        workspaceFolderManager
+    );
+    context.subscriptions.push(
+        configureAutocomplete,
+        telemetry.registerCommand(
+            "databricks.autocomplete.configure",
+            configureAutocomplete.configureCommand,
+            configureAutocomplete
+        )
+    );
 
     const environmentDependenciesInstaller =
         new EnvironmentDependenciesInstaller(
@@ -371,7 +383,8 @@ export async function activate(
             new EnvironmentDependenciesVerifier(
                 connectionManager,
                 pythonExtensionWrapper,
-                environmentDependenciesInstaller
+                environmentDependenciesInstaller,
+                configureAutocomplete
             )
     );
     const environmentCommands = new EnvironmentCommands(
@@ -458,22 +471,6 @@ export async function activate(
     );
     featureManager.isEnabled("environment.dependencies");
 
-    const configureAutocomplete = new ConfigureAutocomplete(
-        context,
-        stateStorage,
-        workspaceFolderManager,
-        pythonExtensionWrapper,
-        environmentDependenciesInstaller
-    );
-    context.subscriptions.push(
-        configureAutocomplete,
-        telemetry.registerCommand(
-            "databricks.autocomplete.configure",
-            configureAutocomplete.configureCommand,
-            configureAutocomplete
-        )
-    );
-
     const codeSynchroniser = new CodeSynchronizer(
         connectionManager,
         configModel,
@@ -517,6 +514,8 @@ export async function activate(
         configModel
     );
 
+    const clusterModel = new ClusterModel(connectionManager);
+
     const connectionCommands = new ConnectionCommands(
         workspaceFsCommands,
         connectionManager,

packages/databricks-vscode/src/feature-manager/FeatureManager.ts

@@ -18,6 +18,7 @@ export interface FeatureState {
 export interface FeatureStepState {
     id: string;
     available: boolean;
+    optional?: boolean;
     title?: string;
     message?: string;
     warning?: string;

packages/databricks-vscode/src/feature-manager/MultiStepAccessVerfier.ts

@@ -35,7 +35,9 @@ export abstract class MultiStepAccessVerifier implements Feature {
         this.state.steps.set(stepState.id, stepState);
         const oldAvailability = this.state.available;
         this.state.available = Array.from(this.state.steps).reduce(
-            (result, current) => result && current[1].available,
+            (result, current) =>
+                result &&
+                (current[1].available || current[1].optional === true),
             true
         );
         const nonComparableFields: Array<keyof FeatureStepState> = ["action"];
@@ -64,14 +66,16 @@ export abstract class MultiStepAccessVerifier implements Feature {
         title: string,
         message?: string,
         action?: FeatureEnableAction,
-        isDisabledByFf?: boolean
+        isDisabledByFf?: boolean,
+        optional?: boolean
     ) {
         return this.updateStep({
             id: id,
             available: false,
             title,
             message,
             action,
+            optional,
             isDisabledByFf: isDisabledByFf === true,
         });
     }

packages/databricks-vscode/src/language/ConfigureAutocomplete.ts

@@ -1,17 +1,14 @@
-import {logging} from "@databricks/databricks-sdk";
-import {appendFile, mkdir, readdir, readFile} from "fs/promises";
+import {appendFile, mkdir, stat, readFile} from "fs/promises";
 import path from "path";
 import {
     ExtensionContext,
     window,
     Disposable,
     workspace,
     ConfigurationTarget,
+    EventEmitter,
 } from "vscode";
-import {Loggers} from "../logger";
 import {StateStorage} from "../vscode-objs/StateStorage";
-import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
-import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller";
 import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 
 async function getImportString(context: ExtensionContext) {
@@ -31,29 +28,21 @@ async function getImportString(context: ExtensionContext) {
     }
 }
 
-type StepResult = "Skip" | "Cancel" | "Error" | undefined;
-
-interface Step {
-    fn: (dryRun: boolean) => Promise<StepResult>;
-    required?: boolean;
-}
-
 export class ConfigureAutocomplete implements Disposable {
     private disposables: Disposable[] = [];
+    private _onDidUpdateEmitter = new EventEmitter<void>();
+    public onDidUpdate = this._onDidUpdateEmitter.event;
 
     constructor(
         private readonly context: ExtensionContext,
         private readonly stateStorage: StateStorage,
-        private readonly workspaceFolderManager: WorkspaceFolderManager,
-        private readonly pythonExtension: MsPythonExtensionWrapper,
-        private readonly environmentDependenciesInstaller: EnvironmentDependenciesInstaller
+        private readonly workspaceFolderManager: WorkspaceFolderManager
     ) {
-        //Remove any type stubs that users already have. We will now start using the python SDK (installed with databricks-connect)
+        // Remove any type stubs that users already have. We will now start using the python SDK (installed with databricks-connect)
         let extraPaths =
             workspace
                 .getConfiguration("python")
                 .get<Array<string>>("analysis.extraPaths") ?? [];
-
         extraPaths = extraPaths.filter(
             (value) =>
                 !value.endsWith(path.join("resources", "python", "stubs"))
@@ -65,120 +54,17 @@ export class ConfigureAutocomplete implements Disposable {
                 extraPaths,
                 ConfigurationTarget.Global
             );
-        this.configure();
     }
 
     dispose() {
         this.disposables.forEach((i) => i.dispose());
     }
 
-    private async tryStep(fn: () => Promise<StepResult>) {
-        try {
-            return await fn();
-        } catch (e) {
-            logging.NamedLogger.getOrCreate(Loggers.Extension).error(
-                "Error configuring autocomplete",
-                e
-            );
-
-            if (e instanceof Error) {
-                window.showErrorMessage(
-                    `Error configuring autocomplete: ${e.message}`
-                );
-            }
-            return "Error";
-        }
-    }
-
-    /*
-        Skip run if all the required steps return "Skip". 
-    */
-    private async shouldSkipRun(steps: Step[]) {
-        for (const {fn, required} of steps) {
-            const result = await this.tryStep(() => fn(true));
-            if (result === "Error") {
-                return true;
-            }
-            if (result !== "Skip" && required) {
-                return false;
-            }
-        }
-        return true;
-    }
-
-    async configureCommand() {
-        this.stateStorage.set("databricks.autocompletion.skipConfigure", false);
-        return this.configure(true);
-    }
-
-    private async configure(force = false) {
-        if (
-            !force ||
-            this.stateStorage.get("databricks.autocompletion.skipConfigure")
-        ) {
-            return;
-        }
-
-        const steps = [
-            {
-                fn: async (dryRun = false) => await this.installPyspark(dryRun),
-            },
-            {
-                fn: async (dryRun = false) => this.addBuiltinsFile(dryRun),
-                required: true,
-            },
-        ];
-
-        // Force is only set when running from command pallet and we do a fresh configure if it is set.
-        if (!force && (await this.shouldSkipRun(steps))) {
-            return;
-        }
-
-        const choice = await window.showInformationMessage(
-            "Do you want to configure autocompletion for Databricks specific globals (dbutils etc)?",
-            "Configure",
-            "Cancel",
-            "Never for this workspace"
-        );
-
-        if (choice === "Cancel" || choice === undefined) {
-            return;
-        }
-
-        if (choice === "Never for this workspace") {
-            this.stateStorage.set(
-                "databricks.autocompletion.skipConfigure",
-                true
-            );
-            return;
-        }
-
-        for (const {fn} of steps) {
-            const result = await this.tryStep(() => fn(false));
-            if (result === "Error" || result === "Cancel") {
-                return;
-            }
-        }
-    }
-
-    private async installPyspark(dryRun = false): Promise<StepResult> {
-        const executable = await this.pythonExtension.getPythonExecutable();
-
-        if (executable === undefined) {
-            return "Skip";
-        }
-
-        if (dryRun) {
-            return;
-        }
-        this.environmentDependenciesInstaller.show(false);
-    }
-
     private get projectRootPath() {
-        return this.workspaceFolderManager.activeProjectUri.fsPath;
+        return this.workspaceFolderManager.activeWorkspaceFolder.uri.fsPath;
     }
 
-    private async addBuiltinsFile(dryRun = false): Promise<StepResult> {
+    private get builtinsPath() {
         const stubPath = workspace
             .getConfiguration("python")
             .get<string>("analysis.stubPath");
@@ -187,45 +73,76 @@ export class ConfigureAutocomplete implements Disposable {
             ? path.join(this.projectRootPath, stubPath)
             : this.projectRootPath;
 
+        return path.join(builtinsDir, "__builtins__.pyi");
+    }
+
+    public async configureCommand() {
+        await this.stateStorage.set(
+            "databricks.autocompletion.skipConfigure",
+            false
+        );
+        return this.setupBuiltins();
+    }
+
+    public async shouldSetupBuiltins(): Promise<boolean> {
+        if (this.stateStorage.get("databricks.autocompletion.skipConfigure")) {
+            return false;
+        }
+
+        const builtinsPath = this.builtinsPath;
+
         let builtinsFileExists = false;
         try {
-            builtinsFileExists = (await readdir(builtinsDir)).includes(
-                "__builtins__.pyi"
-            );
-        } catch (e) {}
-
-        const builtinsPath = path.join(builtinsDir, "__builtins__.pyi");
+            const stats = await stat(builtinsPath);
+            builtinsFileExists = stats.isFile();
+        } catch (e) {
+            builtinsFileExists = false;
+        }
 
         const importString = await getImportString(this.context);
         if (importString === undefined) {
-            return "Error";
+            return false;
         }
 
         if (
             builtinsFileExists &&
             (await readFile(builtinsPath, "utf-8")).includes(importString)
         ) {
-            return "Skip";
+            return false;
         }
 
-        if (dryRun) {
+        return true;
+    }
+
+    private async setupBuiltins(): Promise<void> {
+        if (!(await this.shouldSetupBuiltins())) {
             return;
         }
 
-        const messageString = `${
-            builtinsFileExists ? "Update" : "Create"
-        } ${builtinsPath} ?`;
+        const importString = await getImportString(this.context);
+        if (!importString) {
+            return;
+        }
+
+        const builtinsPath = this.builtinsPath;
+        const messageString = `Create ${builtinsPath}?`;
         const choice = await window.showInformationMessage(
             messageString,
             "Continue",
-            "Cancel"
+            "Cancel",
+            "Never for this workspace"
         );
 
-        if (choice === "Cancel" || choice === undefined) {
-            return "Cancel";
+        if (choice === "Never for this workspace") {
+            await this.stateStorage.set(
+                "databricks.autocompletion.skipConfigure",
+                true
+            );
+            this._onDidUpdateEmitter.fire();
+        } else if (choice === "Continue") {
+            await mkdir(path.dirname(builtinsPath), {recursive: true});
+            await appendFile(builtinsPath, `\n${importString}\n`);
+            this._onDidUpdateEmitter.fire();
         }
-
-        await mkdir(path.dirname(builtinsPath), {recursive: true});
-        await appendFile(builtinsPath, `\n${importString}\n`);
     }
 }

packages/databricks-vscode/src/language/EnvironmentCommands.ts

@@ -26,12 +26,6 @@ export class EnvironmentCommands {
             "environment.dependencies",
             true
         );
-        if (state.available) {
-            window.showInformationMessage(
-                "Python environment and Databricks Connect are already set up."
-            );
-            return true;
-        }
         for (const [, step] of state.steps) {
             if (step.available || (stepId && step.id !== stepId)) {
                 continue;
@@ -43,6 +37,12 @@ export class EnvironmentCommands {
                 return false;
             }
         }
+        if (state.available) {
+            window.showInformationMessage(
+                "Python environment and Databricks Connect are set up."
+            );
+            return true;
+        }
     }
 
     async refresh() {