From 5dc0308f65e4e3e1e64672ed35994eab1b197c59 Mon Sep 17 00:00:00 2001 From: Kartik Gupta <88345179+kartikgupta-db@users.noreply.github.com> Date: Tue, 13 Aug 2024 13:02:23 +0200 Subject: [PATCH] Allow CLI path to have spaces (#1295) ## Changes ## Tests --- .../src/bundle/BundleInitWizard.ts | 5 +- .../src/cli/CliWrapper.test.ts | 1 - .../databricks-vscode/src/cli/CliWrapper.ts | 55 +++++- .../src/configuration/LoginWizard.ts | 38 ++-- .../configuration/auth/DatabricksCliCheck.ts | 4 +- packages/databricks-vscode/src/extension.ts | 3 + .../src/language/MsPythonExtensionWrapper.ts | 3 +- .../notebooks/NotebookInitScriptManager.ts | 5 +- .../src/test/e2e/bundle_init.e2e.ts | 3 + .../src/test/e2e/wdio.conf.ts | 186 +++++++++--------- .../databricks-vscode/src/utils/shellUtils.ts | 10 +- 11 files changed, 189 insertions(+), 124 deletions(-) diff --git a/packages/databricks-vscode/src/bundle/BundleInitWizard.ts b/packages/databricks-vscode/src/bundle/BundleInitWizard.ts index 0abe577c7..56d2471de 100644 --- a/packages/databricks-vscode/src/bundle/BundleInitWizard.ts +++ b/packages/databricks-vscode/src/bundle/BundleInitWizard.ts @@ -18,6 +18,7 @@ import {Events, Telemetry} from "../telemetry"; import {OverrideableConfigModel} from "../configuration/models/OverrideableConfigModel"; import {writeFile, mkdir} from "fs/promises"; import path from "path"; +import {escapePathArgument} from "../utils/shellUtils"; export async function promptToOpenSubProjects( projects: {absolute: Uri; relative: string}[], @@ -188,12 +189,12 @@ export class BundleInitWizard { "bundle", "init", "--output-dir", - this.cli.escapePathArgument(parentFolder.fsPath), + escapePathArgument(parentFolder.fsPath), ].join(" "); const initialPrompt = `clear; echo "Executing: databricks ${args}\nFollow the steps below to create your new Databricks project.\n"`; const finalPrompt = `echo "\nPress any key to close the terminal and continue ..."; ${ShellUtils.readCmd()}; exit`; terminal.sendText( - `${initialPrompt}; ${this.cli.cliPath} ${args}; ${finalPrompt}` + `${initialPrompt}; ${this.cli.escapedCliPath} ${args}; ${finalPrompt}` ); return new Promise((resolve) => { const closeEvent = window.onDidCloseTerminal(async (t) => { diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index 3cecf27b1..8f1f5f822 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -250,7 +250,6 @@ nothing = true PATH: process.env.PATH, /* eslint-enable @typescript-eslint/naming-convention */ }), - shell: true, }, }; try { diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 9bf59a42e..d9cd4fb90 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -19,9 +19,42 @@ import {quote} from "shell-quote"; import {BundleVariableModel} from "../bundle/models/BundleVariableModel"; import {MsPythonExtensionWrapper} from "../language/MsPythonExtensionWrapper"; import path from "path"; +import {isPowershell} from "../utils/shellUtils"; const withLogContext = logging.withLogContext; -const execFile = promisify(execFileCb); +function getEscapedCommandAndAgrs( + cmd: string, + args: string[], + options: SpawnOptionsWithoutStdio +) { + if (process.platform === "win32") { + args = [ + "/d", //Disables execution of AutoRun commands, which are like .bashrc commands. + "/c", //Carries out the command specified by and then exits the command processor. + `""${cmd}" ${args.map((a) => `"${a}"`).join(" ")}"`, + ]; + cmd = "cmd.exe"; + options = {...options, windowsVerbatimArguments: true}; + } + return {cmd, args, options}; +} + +export const execFile = async ( + file: string, + args: string[], + options: any = {} +): Promise<{ + stdout: string; + stderr: string; +}> => { + const { + cmd, + args: escapedArgs, + options: escapedOptions, + } = getEscapedCommandAndAgrs(file, args, options); + const res = await promisify(execFileCb)(cmd, escapedArgs, escapedOptions); + return {stdout: res.stdout.toString(), stderr: res.stderr.toString()}; +}; export interface Command { command: string; @@ -154,13 +187,14 @@ async function runBundleCommand( }); logger?.debug(quote([cmd, ...args]), {bundleOpName}); + let options: SpawnOptionsWithoutStdio = { + cwd: workspaceFolder.fsPath, + env: removeUndefinedKeys(env), + }; + ({cmd, args, options} = getEscapedCommandAndAgrs(cmd, args, options)); try { - const p = spawn(cmd, args, { - cwd: workspaceFolder.fsPath, - env: removeUndefinedKeys(env), - shell: true, - }); + const p = spawn(cmd, args, options); const {stdout, stderr} = await waitForProcess(p, onStdOut, onStdError); logger?.info(displayLogs.end, { @@ -236,8 +270,10 @@ export class CliWrapper { }; } - escapePathArgument(arg: string): string { - return `"${arg.replaceAll('"', '\\"')}"`; + get escapedCliPath(): string { + return isPowershell() + ? `& "${this.cliPath.replace('"', '\\"')}"` + : `'${this.cliPath.replaceAll("'", "\\'")}'`; } /** @@ -300,6 +336,8 @@ export class CliWrapper { "Failed to parse Databricks Config File, please make sure it's in the correct ini format"; } else if (e.message.includes("spawn UNKNOWN")) { msg = `Failed to parse Databricks Config File using databricks CLI, please make sure you have permissions to execute this binary: "${this.cliPath}"`; + } else { + msg += e.message; } } ctx?.logger?.error(msg, e); @@ -593,7 +631,6 @@ export class CliWrapper { options: { cwd: workspaceFolder.fsPath, env, - shell: true, }, }; } diff --git a/packages/databricks-vscode/src/configuration/LoginWizard.ts b/packages/databricks-vscode/src/configuration/LoginWizard.ts index c9adab841..dea755c4d 100644 --- a/packages/databricks-vscode/src/configuration/LoginWizard.ts +++ b/packages/databricks-vscode/src/configuration/LoginWizard.ts @@ -12,7 +12,6 @@ import { ValidationMessageType, } from "../ui/MultiStepInputWizard"; import {CliWrapper, ConfigEntry} from "../cli/CliWrapper"; -import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import { AuthProvider, AuthType, @@ -31,6 +30,8 @@ import ini from "ini"; import {appendFile, copyFile} from "fs/promises"; import path from "path"; import os from "os"; +import {createFile} from "fs-extra"; +import {getDatabricksConfigFilePath} from "../utils/fileUtils"; interface AuthTypeQuickPickItem extends QuickPickItem { authType?: SdkAuthType; @@ -384,9 +385,17 @@ export async function saveNewProfile( throw new Error("Can't save empty auth provider to a profile"); } - const {path: configFilePath} = await loadConfigFile( - workspaceConfigs.databrickscfgLocation - ); + const configFilePath: string = getDatabricksConfigFilePath().fsPath; + let shouldBackup = true; + try { + await loadConfigFile(configFilePath); + } catch (e) { + shouldBackup = false; + await createFile(configFilePath); + window.showInformationMessage( + `Created a new .databrickscfg file at ${configFilePath}` + ); + } const profile: any = {}; profile[profileName] = Object.fromEntries( @@ -394,15 +403,18 @@ export async function saveNewProfile( ); const iniStr = ini.stringify(profile); const finalStr = `${os.EOL};This profile is autogenerated by the Databricks Extension for VS Code${os.EOL}${iniStr}`; - // Create a backup for .databrickscfg - const backup = path.join( - path.dirname(configFilePath), - `.databrickscfg.${Date.now()}.bak` - ); - await copyFile(configFilePath, backup); - window.showInformationMessage( - `Created a backup for .databrickscfg at ${backup}` - ); + + if (shouldBackup) { + // Create a backup for .databrickscfg + const backup = path.join( + path.dirname(configFilePath), + `.databrickscfg.${Date.now()}.bak` + ); + await copyFile(configFilePath, backup); + window.showInformationMessage( + `Created a backup for .databrickscfg at ${backup}` + ); + } // Write the new profile to .databrickscfg await appendFile(configFilePath, finalStr); diff --git a/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts b/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts index a9a362712..4dc7e60b1 100644 --- a/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts +++ b/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts @@ -1,5 +1,4 @@ import { - ExecUtils, ProductVersion, WorkspaceClient, logging, @@ -8,6 +7,7 @@ import {Disposable, window} from "vscode"; import {DatabricksCliAuthProvider} from "./AuthProvider"; import {orchestrate, OrchestrationLoopError, Step} from "./orchestrate"; import {Loggers} from "../../logger"; +import {execFile} from "../../cli/CliWrapper"; // eslint-disable-next-line @typescript-eslint/no-var-requires const extensionVersion = require("../../../package.json") @@ -92,7 +92,7 @@ export class DatabricksCliCheck implements Disposable { private async login(): Promise { try { - await ExecUtils.execFile(this.authProvider.cliPath, [ + await execFile(this.authProvider.cliPath, [ "auth", "login", "--host", diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 9a925c704..655d137eb 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -322,6 +322,9 @@ export async function activate( bundleRemoteStateModel, configModel, connectionManager, + commands.registerCommand("databricks.internal.showOutput", () => { + loggerManager.showOutputChannel("Databricks Logs"); + }), connectionManager.onDidChangeState(async () => { telemetry.setMetadata( Metadata.USER, diff --git a/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts b/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts index 9b51ae7f9..a06d04083 100644 --- a/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts +++ b/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts @@ -12,9 +12,8 @@ import {StateStorage} from "../vscode-objs/StateStorage"; import {IExtensionApi as MsPythonExtensionApi} from "./MsPythonExtensionApi"; import {Mutex} from "../locking"; import * as childProcess from "node:child_process"; -import {promisify} from "node:util"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; -export const execFile = promisify(childProcess.execFile); +import {execFile} from "../cli/CliWrapper"; export class MsPythonExtensionWrapper implements Disposable { public readonly api: MsPythonExtensionApi; diff --git a/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts b/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts index 412434059..62c494c55 100644 --- a/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts +++ b/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts @@ -17,15 +17,12 @@ import {Loggers} from "../../logger"; import {Context, context} from "@databricks/databricks-sdk/dist/context"; import {Mutex} from "../../locking"; import {MsPythonExtensionWrapper} from "../MsPythonExtensionWrapper"; -import {execFile as ef} from "child_process"; -import {promisify} from "util"; import {FileUtils} from "../../utils"; import {workspaceConfigs} from "../../vscode-objs/WorkspaceConfigs"; import {LocalUri} from "../../sync/SyncDestination"; import {DatabricksEnvFileManager} from "../../file-managers/DatabricksEnvFileManager"; import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager"; - -const execFile = promisify(ef); +import {execFile} from "../../cli/CliWrapper"; async function isDbnbTextEditor(editor?: TextEditor) { try { diff --git a/packages/databricks-vscode/src/test/e2e/bundle_init.e2e.ts b/packages/databricks-vscode/src/test/e2e/bundle_init.e2e.ts index 93d056d54..3d9fd9d09 100644 --- a/packages/databricks-vscode/src/test/e2e/bundle_init.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/bundle_init.e2e.ts @@ -67,6 +67,8 @@ describe("Bundle Init", async function () { expect(await pick.getLabel()).toBe(parentDir); await pick.select(); + await dismissNotifications(); + // Wait for the databricks cli terminal window to pop up and select all // default options for the default template const editorView = workbench.getEditorView(); @@ -74,6 +76,7 @@ describe("Bundle Init", async function () { const initTab = await getTabByTitle(title); assert(initTab, "Can't find a tab for project-init terminal wizard"); await initTab.select(); + await sleep(1000); //select temaplate type await browser.keys("default-python".split("")); diff --git a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts index 13caf977e..ec62ae7b0 100644 --- a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts +++ b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts @@ -10,13 +10,14 @@ import assert from "assert"; import fs from "fs/promises"; import {ApiError, Config, WorkspaceClient} from "@databricks/databricks-sdk"; import * as ElementCustomCommands from "./customCommands/elementCustomCommands.ts"; -import {execFile, ExecFileOptions} from "node:child_process"; +import {execFile as execFileCb} from "node:child_process"; import {cpSync, mkdirSync, rmSync} from "node:fs"; import {tmpdir} from "node:os"; import packageJson from "../../../package.json" assert {type: "json"}; import {sleep} from "wdio-vscode-service"; import {glob} from "glob"; import {getUniqueResourceName} from "./utils/commonUtils.ts"; +import {promisify} from "node:util"; const WORKSPACE_PATH = path.resolve(tmpdir(), "test-root"); @@ -24,7 +25,7 @@ const __filename = fileURLToPath(import.meta.url); const __dirname = path.dirname(__filename); const {version, name, engines} = packageJson; -const EXTENSION_DIR = path.resolve(tmpdir(), "extension-test", "extension"); +const EXTENSION_DIR = path.resolve(tmpdir(), "extension test", "extension"); const VSIX_PATH = path.resolve( __dirname, "..", @@ -34,6 +35,63 @@ const VSIX_PATH = path.resolve( ); const VSCODE_STORAGE_DIR = path.resolve(tmpdir(), "user-data-dir"); +const metaCharsRegExp = /([()\][%!^"`<>&|;, *?])/g; + +export function escapeCommand(arg: string): string { + // Escape meta chars + arg = arg.replace(metaCharsRegExp, "^$1"); + + return arg; +} + +export function escapeArgument(arg: string): string { + // Convert to string + arg = `${arg}`; + + // Algorithm below is based on https://qntm.org/cmd + + // Sequence of backslashes followed by a double quote: + // double up all the backslashes and escape the double quote + arg = arg.replace(/(\\*)"/g, '$1$1\\"'); + + // Sequence of backslashes followed by the end of the string + // (which will become a double quote later): + // double up all the backslashes + arg = arg.replace(/(\\*)$/, "$1$1"); + + // All other backslashes occur literally + + // Quote the whole thing: + arg = `"${arg}"`; + + // Escape meta chars + arg = arg.replace(metaCharsRegExp, "^$1"); + + return arg; +} + +const execFile = async ( + file: string, + args: string[], + options: any = {} +): Promise<{ + stdout: string; + stderr: string; +}> => { + if (process.platform === "win32") { + const realArgs = [escapeCommand(file)] + .concat(args.map(escapeArgument).join(" ")) + .join(" "); + + file = "cmd.exe"; + args = ["/d", "/s", "/c", `"${realArgs}"`]; + console.log("execFile", file, args); + options = {...options, windowsVerbatimArguments: true}; + } + const res = await promisify(execFileCb)(file, args, options); + return {stdout: res.stdout.toString(), stderr: res.stderr.toString()}; +}; + export const config: Options.Testrunner = { // // ==================== @@ -208,7 +266,7 @@ export const config: Options.Testrunner = { framework: "mocha", // // The number of times to retry the entire specfile when it fails as a whole - specFileRetries: 1, + specFileRetries: 0, // // Delay in seconds between the spec file retry attempts specFileRetriesDelay: 0, @@ -327,94 +385,38 @@ export const config: Options.Testrunner = { * @param {Array.} specs List of spec file paths that are to be run * @param {String} cid worker id (e.g. 0-0) */ - beforeSession: async function (config, capabilities, specs, cid) { - if (cid === "0-0") { - const binary: string = capabilities["wdio:vscodeOptions"] - .binary as string; - let cli: string; - let spawnArgs: ExecFileOptions; - switch (process.platform) { - case "win32": - cli = path.resolve(binary, "..", "bin", "code"); - spawnArgs = { - shell: true, - }; - break; - case "darwin": - cli = path.resolve( - binary, - "..", - "..", - "Resources/app/bin/code" - ); - spawnArgs = { - shell: false, - }; - break; - } - await new Promise((resolve, reject) => { - const extensionDependencies = - packageJson.extensionDependencies.flatMap((item) => [ - "--install-extension", - item, - ]); - execFile( - cli, - [ - "--extensions-dir", - EXTENSION_DIR, - ...extensionDependencies, - "--install-extension", - VSIX_PATH, - "--force", - ], - spawnArgs, - (error, stdout, stderr) => { - if (stdout) { - console.log(stdout); - } - if (error) { - console.error(stderr); - console.error(error); - reject(error); - } - resolve(undefined); - } - ); - }); - await fs.writeFile(path.join(WORKSPACE_PATH, "lock.unlock"), ""); - console.log( - `lock file ${path.join(WORKSPACE_PATH, `lock.unlock`)}` - ); - } else { - const startTime = Date.now(); - await new Promise((resolve, reject) => { - console.log( - `${cid} waiting for extension to install; file ${path.join( - WORKSPACE_PATH, - "lock.unlock" - )}` + beforeSession: async function (config, capabilities) { + const binary: string = capabilities["wdio:vscodeOptions"] + .binary as string; + let cli: string = ""; + switch (process.platform) { + case "win32": + cli = path.resolve(binary, "..", "bin", "code"); + break; + case "darwin": + cli = path.resolve( + binary, + "..", + "..", + "Resources/app/bin/code" ); - const interval = setInterval(async () => { - try { - await fs.stat(path.join(WORKSPACE_PATH, "lock.unlock")); - clearInterval(interval); - resolve(undefined); - } catch (e) { - console.log( - `${cid} waiting for extension to install; file ${path.join( - WORKSPACE_PATH, - "lock.unlock" - )}` - ); - if (Date.now() - startTime > 60_000) { - clearInterval(interval); - reject("Timeout waiting for extension to install"); - } - } - }, 5_000); - }); + break; } + const extensionDependencies = packageJson.extensionDependencies.flatMap( + (item) => ["--install-extension", item] + ); + + console.log("running vscode cli"); + const res = await execFile(cli, [ + "--extensions-dir", + EXTENSION_DIR, + ...extensionDependencies, + "--install-extension", + VSIX_PATH, + "--force", + ]); + + console.log(res.stdout, res.stderr); }, /** @@ -595,7 +597,11 @@ async function startCluster( await workspaceClient.clusters.start({ cluster_id: clusterId, }) - ).wait(); + ).wait({ + onProgress: async (state) => { + console.log(`Cluster state: ${state.state}`); + }, + }); } catch (e: unknown) { if (!(e instanceof ApiError && e.message.includes("INVALID_STATE"))) { throw e; diff --git a/packages/databricks-vscode/src/utils/shellUtils.ts b/packages/databricks-vscode/src/utils/shellUtils.ts index 8d3b5e9f4..ba6344ffb 100644 --- a/packages/databricks-vscode/src/utils/shellUtils.ts +++ b/packages/databricks-vscode/src/utils/shellUtils.ts @@ -1,8 +1,16 @@ import {env} from "vscode"; +export function isPowershell() { + return env.shell.toLowerCase().includes("powershell"); +} + export function readCmd() { - if (env.shell.toLowerCase().includes("powershell")) { + if (isPowershell()) { return "Read-Host"; } return "read"; } + +export function escapePathArgument(arg: string): string { + return `"${arg.replaceAll('"', '\\"')}"`; +}