diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index c5d08d953..03d15e732 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -57,6 +57,7 @@ "onCommand:databricks.run.runEditorContents", "onCommand:databricks.quickstart.open", "onCommand:databricks.logs.openFolder", + "onCommand:databricks.autocomplete.configure", "onView:configurationView", "onView:clusterView", "onTaskType:databricks", @@ -171,6 +172,11 @@ "title": "Open full logs", "category": "Databricks" }, + { + "command": "databricks.autocomplete.configure", + "title": "Configure autocomplete for Databricks globals", + "category": "Databricks" + }, { "command": "databricks.cluster.start", "title": "Start Cluster", @@ -526,6 +532,9 @@ } ] }, + "extensionDependencies": [ + "ms-python.python" + ], "vsce": { "dependencies": false, "useYarn": false diff --git a/packages/databricks-vscode/resources/python/stubs/databricks/sdk/runtime.pyi b/packages/databricks-vscode/resources/python/stubs/databricks/sdk/runtime.pyi new file mode 100644 index 000000000..4f53dd311 --- /dev/null +++ b/packages/databricks-vscode/resources/python/stubs/databricks/sdk/runtime.pyi @@ -0,0 +1,317 @@ +import typing +from pyspark.sql.session import SparkSession +from pyspark.sql.functions import udf as U +from pyspark.sql.context import SQLContext + +udf = U +spark: SparkSession +sc = spark.sparkContext +sqlContext: SQLContext +sql = sqlContext.sql +table = sqlContext.table + +def displayHTML(html): + """ + Display HTML data. + + Parameters + ---------- + data : URL or HTML string + If data is a URL, display the resource at that URL, the resource is loaded dynamically by the browser. + Otherwise data should be the HTML to be displayed. + + See also: + IPython.display.HTML + IPython.display.display_html + """ + ... + +def display(input=None, *args, **kwargs): + """ + Display plots or data. + + Display plot: + - display() # no-op + - display(matplotlib.figure.Figure) + + Display dataset: + - display(spark.DataFrame) + - display(list) # if list can be converted to DataFrame, e.g., list of named tuples + - display(pandas.DataFrame) + - display(koalas.DataFrame) + - display(pyspark.pandas.DataFrame) + + Display any other value that has a _repr_html_() method + + For Spark 2.0 and 2.1: + - display(DataFrame, streamName='optional', trigger=optional pyspark.sql.streaming.Trigger, + checkpointLocation='optional') + + For Spark 2.2+: + - display(DataFrame, streamName='optional', trigger=optional interval like '1 second', + checkpointLocation='optional') + """ + ... + +class dbutils: + class credentials: + """ + Utilities for interacting with credentials within notebooks + """ + + def assumeRole(role: str) -> bool: + """ + Sets the role ARN to assume when looking for credentials to authenticate with S3 + """ + ... + def showCurrentRole() -> typing.List[str]: + """ + Shows the currently set role + """ + ... + def showRoles() -> typing.List[str]: + """ + Shows the set of possibly assumed roles + """ + ... + + class data: + """ + Utilities for understanding and interacting with datasets (EXPERIMENTAL) + """ + + def summarize(df: any, precise: bool) -> None: + """ + Summarize a Spark DataFrame and visualize the statistics to get quick insights + """ + ... + + class fs: + """ + Manipulates the Databricks filesystem (DBFS) from the console + """ + + def cp(from_: str, to: str, recurse: bool = False) -> bool: + """ + Copies a file or directory, possibly across FileSystems + """ + ... + def head(file: str, maxBytes: int = 65536) -> str: + """ + Returns up to the first 'maxBytes' bytes of the given file as a String encoded in UTF-8 + """ + ... + def ls(dir: str) -> typing.List[str]: + """ + Lists the contents of a directory + """ + ... + def mkdirs(dir: str) -> bool: + """ + Creates the given directory if it does not exist, also creating any necessary parent directories + """ + ... + def mv(from_: str, to: str, recurse: bool = False) -> bool: + """ + Moves a file or directory, possibly across FileSystems + """ + ... + def put(file: str, contents: str, overwrite: bool = False) -> bool: + """ + Writes the given String out to a file, encoded in UTF-8 + """ + ... + def rm(dir: str, recurse: bool = False) -> bool: + """ + Removes a file or directory + """ + ... + def mount( + source: str, + mountPoint: str, + encryptionType: str = "", + owner: str = "", + extraConfigs: typing.Map[str, str] = {}, + ) -> bool: + """ + Mounts the given source directory into DBFS at the given mount point + """ + ... + def unmount(mountPoint: str) -> bool: + """ + Deletes a DBFS mount point + """ + ... + def updateMount( + source: str, + mountPoint: str, + encryptionType: str = "", + owner: str = "", + extraConfigs: typing.Map[str, str] = {}, + ) -> bool: + """ + Similar to mount(), but updates an existing mount point (if present) instead of creating a new one + """ + ... + def mounts() -> typing.List[str]: + """ + Displays information about what is mounted within DBFS + """ + ... + def refreshMounts() -> bool: + """ + Forces all machines in this cluster to refresh their mount cache, ensuring they receive the most recent information + """ + ... + + class jobs: + """ + Utilities for leveraging jobs features + """ + + class taskValues: + """ + Provides utilities for leveraging job task values + """ + + def get( + taskKey: str, key: str, default: any = None, debugValue: any = None + ) -> None: + """ + Returns the latest task value that belongs to the current job run + """ + ... + def set(key: str, value: any) -> None: + """ + Sets a task value on the current task run + """ + ... + + class library: + """ + Utilities for session isolated libraries + """ + + def restartPython() -> None: + """ + Restart python process for the current notebook session + """ + ... + + class notebook: + """ + Utilities for the control flow of a notebook (EXPERIMENTAL) + """ + + def exit(value: str) -> None: + """ + This method lets you exit a notebook with a value + """ + ... + def run(path: str, timeoutSeconds: int, arguments: typing.Map[str, str]) -> str: + """ + This method runs a notebook and returns its exit value + """ + ... + + class secrets: + """ + Provides utilities for leveraging secrets within notebooks + """ + + def get(scope: str, key: str) -> str: + """ + Gets the string representation of a secret value with scope and key + """ + ... + def list(scope: str) -> typing.List[str]: + """ + Lists secret metadata for secrets within a scope + """ + ... + def listScopes() -> typing.List[str]: + """ + Lists secret scopes + """ + ... + + class widgets: + """ + provides utilities for working with notebook widgets. You can create different types of widgets and get their bound value + """ + + def get(name: str) -> str: + """Returns the current value of a widget with give name. + + :param name: Name of the argument to be accessed + :return: Current value of the widget or default value + """ + ... + def getArgument(name: str, defaultValue: typing.Optional[str] = None) -> str: + """Returns the current value of a widget with give name. + + :param name: Name of the argument to be accessed + :param defaultValue: (Deprecated) default value + :return: Current value of the widget or default value + """ + ... + def text(name: str, defaultValue: str, label: str = None): + """Creates a text input widget with given name, default value and optional label for + display + :param name: Name of argument associated with the new input widget + :param defaultValue: Default value of the input widget + :param label: Optional label string for display in notebook and dashboard + """ + ... + def dropdown( + name: str, defaultValue: str, choices: typing.List[str], label: str = None + ): + """Creates a dropdown input widget with given specification. + :param name: Name of argument associated with the new input widget + :param defaultValue: Default value of the input widget (must be one of choices) + :param choices: List of choices for the dropdown input widget + :param label: Optional label string for display in notebook and dashboard + """ + ... + def combobox( + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None, + ): + """Creates a combobox input widget with given specification. + :param name: Name of argument associated with the new input widget + :param defaultValue: Default value of the input widget + :param choices: List of choices for the dropdown input widget + :param label: Optional label string for display in notebook and dashboard + """ + ... + def multiselect( + name: str, + defaultValue: str, + choices: typing.List[str], + label: typing.Optional[str] = None, + ): + """Creates a multiselect input widget with given specification. + :param name: Name of argument associated with the new input widget + :param defaultValue: Default value of the input widget (must be one of choices) + :param choices: List of choices for the dropdown input widget + :param label: Optional label string for display in notebook and dashboard + """ + ... + def remove(name: str): + """Removes given input widget. If widget does not exist it will throw an error. + :param name: Name of argument associated with input widget to be removed + """ + ... + def removeAll(): + """Removes all input widgets in the notebook.""" + ... + @property + def meta() -> MetaUtils: + """ + Methods to hook into the compiler (EXPERIMENTAL) + """ + ... + +getArgument = dbutils.widgets.getArgument diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 09a640455..9a245e614 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -19,6 +19,7 @@ import {LoggerManager, Loggers} from "./logger"; import {NamedLogger} from "@databricks/databricks-sdk/dist/logging"; import {workspaceConfigs} from "./WorkspaceConfigs"; import {PackageJsonUtils, UtilsCommands} from "./utils"; +import {ConfigureAutocomplete} from "./language/ConfigureAutocomplete"; export async function activate( context: ExtensionContext @@ -51,6 +52,19 @@ export async function activate( metadata: await PackageJsonUtils.getMetadata(context), }); + const configureAutocomplete = new ConfigureAutocomplete( + context, + workspace.workspaceFolders[0].uri.fsPath + ); + context.subscriptions.push( + configureAutocomplete, + commands.registerCommand( + "databricks.autocomplete.configure", + configureAutocomplete.configureCommand, + configureAutocomplete + ) + ); + context.subscriptions.push( commands.registerCommand( "databricks.logs.openFolder", diff --git a/packages/databricks-vscode/src/language/ConfigureAutocomplete.ts b/packages/databricks-vscode/src/language/ConfigureAutocomplete.ts new file mode 100644 index 000000000..4ecf6cd61 --- /dev/null +++ b/packages/databricks-vscode/src/language/ConfigureAutocomplete.ts @@ -0,0 +1,291 @@ +import {NamedLogger} from "@databricks/databricks-sdk/dist/logging"; +import {appendFile, mkdir, readdir, readFile} from "fs/promises"; +import path from "path"; +import { + ExtensionContext, + extensions, + Uri, + window, + Event, + Disposable, + workspace, + ConfigurationTarget, +} from "vscode"; +import {Loggers} from "../logger"; + +type Resource = Uri | undefined; + +// Refer https://github.com/microsoft/vscode-python/blob/main/src/client/apiTypes.ts +interface IPythonExtension { + /** + * Return internal settings within the extension which are stored in VSCode storage + */ + settings: { + /** + * An event that is emitted when execution details (for a resource) change. For instance, when interpreter configuration changes. + */ + readonly onDidChangeExecutionDetails: Event; + /** + * Returns all the details the consumer needs to execute code within the selected environment, + * corresponding to the specified resource taking into account any workspace-specific settings + * for the workspace to which this resource belongs. + * @param {Resource} [resource] A resource for which the setting is asked for. + * * When no resource is provided, the setting scoped to the first workspace folder is returned. + * * If no folder is present, it returns the global setting. + * @returns {({ execCommand: string[] | undefined })} + */ + getExecutionDetails(resource?: Resource): { + /** + * E.g of execution commands returned could be, + * * `['']` + * * `['']` + * * `['conda', 'run', 'python']` which is used to run from within Conda environments. + * or something similar for some other Python environments. + * + * @type {(string[] | undefined)} When return value is `undefined`, it means no interpreter is set. + * Otherwise, join the items returned using space to construct the full execution command. + */ + execCommand: string[] | undefined; + }; + }; +} + +const importString = "from databricks.sdk.runtime import *"; + +type StepResult = "Skip" | "Cancel" | "Error" | undefined; + +interface Step { + fn: (dryRun: boolean) => Promise; + required?: boolean; +} + +export class ConfigureAutocomplete implements Disposable { + private disposables: Disposable[] = []; + private _onPythonChangeEventListenerAdded = false; + + constructor( + private readonly context: ExtensionContext, + private readonly workspaceFolder: string + ) { + this.configure(); + } + + dispose() { + this.disposables.forEach((i) => i.dispose()); + } + + private async tryStep(fn: () => Promise) { + try { + return await fn(); + } catch (e) { + 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() { + return this.configure(true); + } + + private async configure(force = false) { + const pythonExtension = extensions.getExtension("ms-python.python"); + if (pythonExtension === undefined) { + window.showWarningMessage( + "VSCode Extension for Databricks requires Microsoft Python extension for providing autocompletion. Autocompletion will be disabled now." + ); + return; + } + if (!pythonExtension.isActive) { + await pythonExtension.activate(); + } + /* + We hook into the python extension, so that whenever user changes python interpreter (or environment), + we prompt them to go through the configuration steps again. For now this only involves installing pyspark. + In future, we would want them to install databricks sdk in the new environment. + */ + if (!this._onPythonChangeEventListenerAdded) { + this.disposables.push( + pythonExtension.exports.settings.onDidChangeExecutionDetails( + () => this.configure(true), + this + ) + ); + this._onPythonChangeEventListenerAdded = true; + } + + const steps = [ + { + fn: async (dryRun = false) => + await this.installPyspark( + pythonExtension.exports as IPythonExtension, + dryRun + ), + }, + { + fn: async (dryRun = false) => this.updateExtraPaths(dryRun), + required: true, + }, + { + 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( + "To allow autocompletion for Databricks specific globals (like dbutils), we need to install pyspark and add (or modify) __builtins__.pyi file to your project", + "Continue", + "Cancel" + ); + + if (choice === "Cancel" || choice === undefined) { + return; + } + + for (const {fn} of steps) { + const result = await this.tryStep(() => fn(false)); + if (result === "Error" || result === "Cancel") { + return; + } + } + } + + private async installPyspark( + pythonExtension: IPythonExtension, + dryRun = false + ): Promise { + const execCommandParts = pythonExtension.settings.getExecutionDetails( + workspace.workspaceFolders?.[0].uri + ).execCommand; + + if (execCommandParts === undefined) { + return "Skip"; + } + + if (dryRun) { + return; + } + const choice = await window.showInformationMessage( + ["Install pyspark in local env?"].join("\n"), + "Install PySpark", + "Continue without PySpark", + "Cancel" + ); + + if (choice === "Cancel" || choice === undefined) { + return "Cancel"; + } + if (choice === "Continue without PySpark") { + return "Skip"; + } + + //TODO: Make sure that pyspark is not updated if it is already installed + const execCommand = execCommandParts + .concat(["-m", "pip", "install", "pyspark"]) + .join(" "); + + const terminal = window.createTerminal("pip"); + this.disposables.push(terminal); + terminal.sendText(execCommand); + terminal.show(); + } + + private async updateExtraPaths(dryRun = false): Promise { + const extraPaths = + workspace + .getConfiguration("python") + .get>("analysis.extraPaths") ?? []; + const stubPath = this.context.asAbsolutePath( + path.join("resources", "python", "stubs") + ); + if (extraPaths.includes(stubPath)) { + return "Skip"; + } + if (dryRun) { + return; + } + extraPaths.push(stubPath); + workspace + .getConfiguration("python") + .update( + "analysis.extraPaths", + extraPaths, + ConfigurationTarget.Global + ); + } + + private async addBuiltinsFile(dryRun = false): Promise { + const stubPath = workspace + .getConfiguration("python") + .get("analysis.stubPath"); + + const builtinsDir = stubPath + ? path.join(this.workspaceFolder, stubPath) + : this.workspaceFolder; + + let builtinsFileExists = false; + try { + builtinsFileExists = (await readdir(builtinsDir)).includes( + "__builtins__.pyi" + ); + } catch (e) {} + + const builtinsPath = path.join(builtinsDir, "__builtins__.pyi"); + + if ( + builtinsFileExists && + (await readFile(builtinsPath, "utf-8")).includes(importString) + ) { + return "Skip"; + } + + if (dryRun) { + return; + } + + const messageString = `${ + builtinsFileExists ? "Update" : "Create" + } ${builtinsPath} ?`; + const choice = await window.showInformationMessage( + messageString, + "Continue", + "Cancel" + ); + + if (choice === "Cancel" || choice === undefined) { + return "Cancel"; + } + + await mkdir(path.dirname(builtinsPath), {recursive: true}); + await appendFile(builtinsPath, `\n${importString}\n`); + } +} diff --git a/packages/databricks-vscode/src/test/e2e/configure.e2e.ts b/packages/databricks-vscode/src/test/e2e/configure.e2e.ts index a02eb197d..a9362018f 100644 --- a/packages/databricks-vscode/src/test/e2e/configure.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/configure.e2e.ts @@ -1,7 +1,11 @@ import assert from "node:assert"; import path from "node:path"; import * as fs from "fs/promises"; -import {getViewSection, waitForTreeItems} from "./utils"; +import { + getViewSection, + waitForPythonExtension, + waitForTreeItems, +} from "./utils"; import { CustomTreeSection, InputBox, @@ -10,12 +14,14 @@ import { Workbench, } from "wdio-vscode-service"; -describe("Configure Databricks Extension", () => { +describe("Configure Databricks Extension", async function () { // this will be populated by the tests let clusterId: string; let projectDir: string; let workbench: Workbench; + this.timeout(3 * 60 * 1000); + before(async function () { assert( process.env.TEST_DEFAULT_CLUSTER_ID, @@ -29,6 +35,7 @@ describe("Configure Databricks Extension", () => { projectDir = process.env.WORKSPACE_PATH; workbench = await browser.getWorkbench(); + await waitForPythonExtension(); }); it("should open VSCode", async function () { diff --git a/packages/databricks-vscode/src/test/e2e/run_job_on_cluster.e2e.ts b/packages/databricks-vscode/src/test/e2e/run_job_on_cluster.e2e.ts index 8bc180064..17680e148 100644 --- a/packages/databricks-vscode/src/test/e2e/run_job_on_cluster.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/run_job_on_cluster.e2e.ts @@ -1,7 +1,12 @@ import path from "node:path"; import * as fs from "fs/promises"; import assert from "node:assert"; -import {getViewSection, getViewSubSection, waitForTreeItems} from "./utils"; +import { + getViewSection, + getViewSubSection, + waitForPythonExtension, + waitForTreeItems, +} from "./utils"; import {sleep, TreeItem} from "wdio-vscode-service"; describe("Run python on cluster", async function () { @@ -40,6 +45,7 @@ describe("Run python on cluster", async function () { const section = await getViewSection("CONFIGURATION"); assert(section); await waitForTreeItems(section); + await waitForPythonExtension(); }); beforeEach(async () => { diff --git a/packages/databricks-vscode/src/test/e2e/run_on_cluster.e2e.ts b/packages/databricks-vscode/src/test/e2e/run_on_cluster.e2e.ts index 30aebaaa0..7a68d249d 100644 --- a/packages/databricks-vscode/src/test/e2e/run_on_cluster.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/run_on_cluster.e2e.ts @@ -1,11 +1,17 @@ import path from "node:path"; import * as fs from "fs/promises"; import assert from "node:assert"; -import {getViewSection, getViewSubSection, waitForTreeItems} from "./utils"; +import { + getViewSection, + getViewSubSection, + waitForPythonExtension, + waitForTreeItems, +} from "./utils"; import {sleep} from "wdio-vscode-service"; -describe("Run python on cluster", () => { +describe("Run python on cluster", async function () { let projectDir: string; + this.timeout(3 * 60 * 1000); before(async () => { assert(process.env.TEST_DEFAULT_CLUSTER_ID); @@ -27,6 +33,7 @@ describe("Run python on cluster", () => { path.join(projectDir, "hello.py"), `spark.sql('SELECT "hello world"').show()` ); + await waitForPythonExtension(); }); it("should connect to Databricks", async () => { diff --git a/packages/databricks-vscode/src/test/e2e/utils.ts b/packages/databricks-vscode/src/test/e2e/utils.ts index 4922bbc36..732b9c46f 100644 --- a/packages/databricks-vscode/src/test/e2e/utils.ts +++ b/packages/databricks-vscode/src/test/e2e/utils.ts @@ -1,4 +1,4 @@ -import * as assert from "node:assert"; +import assert from "node:assert"; import { CustomTreeSection, sleep, @@ -81,3 +81,43 @@ export async function waitForTreeItems( await new Promise((resolve) => setTimeout(resolve, 200)); } } + +export async function waitForPythonExtension() { + const section = await getViewSection("CONFIGURATION"); + assert(section); + const welcome = await section.findWelcomeContent(); + assert(welcome); + sleep(1000); + const workbench = await browser.getWorkbench(); + const notifs = await workbench.getNotifications(); + for (const n of notifs) { + if ( + (await n.getActions()).find( + (btn) => btn.getTitle() === "Install and Reload" + ) !== undefined + ) { + await n.takeAction("Install and Reload"); + } + } + + await browser.waitUntil( + async () => + ( + await ( + await workbench.getEditorView().getActiveTab() + )?.getTitle() + )?.includes("README.quickstart.md") === true, + { + timeout: 120000, + timeoutMsg: + "Timeout when installing python extension and reloading", + } + ); + + sleep(500); + try { + for (const n of notifs) { + await n.dismiss(); + } + } catch {} +} diff --git a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts index faab61c8d..1241d37db 100644 --- a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts +++ b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts @@ -100,6 +100,9 @@ export const config: Options.Testrunner = { extensionPath: process.env.CI ? path.resolve(__dirname, "..", "..", "..", "extension") : path.resolve(__dirname, "..", "..", ".."), + vscodeArgs: { + disableExtensions: false, + }, workspacePath: WORKSPACE_PATH, userSettings: { "editor.fontSize": 14,