Support environment variables in 'Run on Databricks' custom debug config (#247)

![Screenshot 2022-11-25 at 11 31
05](https://user-images.githubusercontent.com/88345179/203964712-4c92c3ed-8abd-46af-a5a6-8d29ccc29371.png)

![Screenshot 2022-11-26 at 20 27
21](https://user-images.githubusercontent.com/88345179/204106049-14c3a0b2-6575-448a-9b67-fa2f58f94cd5.png)

![Screenshot 2022-11-26 at 20 29
15](https://user-images.githubusercontent.com/88345179/204106055-dec9c903-bcf1-4d7b-93ab-927ad24e5820.png)

Author: kartikgupta-db

packages/databricks-vscode/package.json

@@ -394,6 +394,13 @@
                                     "type": "string"
                                 },
                                 "description": "Command line arguments."
+                            },
+                            "env": {
+                                "type": "object",
+                                "items": {
+                                    "type": "string"
+                                },
+                                "description": "Environment variables."
                             }
                         }
                     }
@@ -404,7 +411,8 @@
                         "request": "launch",
                         "name": "Run on Databricks",
                         "program": "${file}",
-                        "args": []
+                        "args": [],
+                        "env": {}
                     }
                 ],
                 "configurationSnippets": [

packages/databricks-vscode/src/run/DatabricksDebugAdapter.ts

@@ -38,6 +38,9 @@ interface ILaunchRequestArguments extends DebugProtocol.LaunchRequestArguments {
 
     /** Command line arguments */
     args?: string[];
+
+    /** Env variables */
+    env?: Record<string, string>;
 }
 
 export class DatabricksDebugAdapterFactory
@@ -246,7 +249,7 @@ export class DatabricksDebugSession extends LoggingDebugSession {
         await commands.executeCommand("workbench.panel.repl.view.focus");
 
         // start the program in the runtime
-        await this.runtime.start(args.program, args.args || []);
+        await this.runtime.start(args.program, args.args || [], args.env || {});
         this.sendResponse(response);
     }
 

packages/databricks-vscode/src/run/DatabricksRuntime.test.ts

@@ -99,7 +99,7 @@ describe(__filename, () => {
             })
         );
 
-        await runtime.start("/Desktop/workspaces/hello.py", []);
+        await runtime.start("/Desktop/workspaces/hello.py", [], {});
 
         verify(connectionManagerMock.waitForConnect()).called();
 
@@ -115,14 +115,15 @@ describe(__filename, () => {
     });
 
     it("should have the right code", async () => {
-        await runtime.start("/Desktop/workspaces/hello.py", []);
+        await runtime.start("/Desktop/workspaces/hello.py", [], {TEST: "TEST"});
 
         const code = capture(executionContextMock.execute).first()[0];
         assert.equal(
             code,
             `import os; os.chdir("/Workspace/Repos/fabian@databricks.com/test");
 import sys; sys.path.append("/Workspace/Repos/fabian@databricks.com/test")
 import sys; sys.argv = ['/Workspace/Repos/fabian@databricks.com/test/hello.py'];
+import os; os.environ["TEST"]='TEST';
 import logging; logger = spark._jvm.org.apache.log4j; logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
 print('43')`
         );
@@ -155,7 +156,7 @@ print('43')`
             })
         );
 
-        await runtime.start("/Desktop/workspaces/hello.py", []);
+        await runtime.start("/Desktop/workspaces/hello.py", [], {});
 
         assert.equal(outputs.length, 7);
         assert.equal(outputs[4].text, "something went wrong");

packages/databricks-vscode/src/run/DatabricksRuntime.ts

@@ -30,6 +30,8 @@ export interface FileAccessor {
     readFile(path: string): Promise<string>;
 }
 
+type EnvVars = Record<string, string>;
+
 export class DatabricksRuntime implements Disposable {
     private _onDidEndEmitter: EventEmitter<void> = new EventEmitter<void>();
     readonly onDidEnd: Event<void> = this._onDidEndEmitter.event;
@@ -60,7 +62,11 @@ export class DatabricksRuntime implements Disposable {
     /**
      * Start executing the given program.
      */
-    public async start(program: string, args: Array<string>): Promise<void> {
+    public async start(
+        program: string,
+        args: Array<string>,
+        envVars: EnvVars
+    ): Promise<void> {
         const start = Date.now();
 
         const log = (message: string, line: number) => {
@@ -168,7 +174,8 @@ export class DatabricksRuntime implements Disposable {
                     program,
                     lines,
                     args,
-                    syncDestination
+                    syncDestination,
+                    envVars
                 ),
                 undefined,
                 this.token
@@ -231,13 +238,28 @@ export class DatabricksRuntime implements Disposable {
         program: string,
         programLines: Array<string>,
         args: Array<string>,
-        syncDestination: SyncDestination
+        syncDestination: SyncDestination,
+        envVars: EnvVars
     ): string {
         const argv = [
             syncDestination.localToRemote(Uri.file(program)),
             ...args,
         ];
 
+        const envVarSetCmds = [];
+        for (const key in envVars) {
+            if (!/^[a-zA-Z_]{1,}[a-zA-Z0-9_]*$/.test(key)) {
+                this._onErrorEmitter.fire(
+                    `Invalid environment variable ${key}: Only lower and upper case letters, digits and '_'(underscore) are allowed.`
+                );
+                return "";
+            }
+            const cmd = `os.environ["${key}"]='${this.escapePythonString(
+                envVars[key]
+            )}'`;
+            envVarSetCmds.push(cmd);
+        }
+
         return [
             // set working directory
             `import os; os.chdir("${syncDestination.localToRemoteDir(
@@ -252,14 +274,16 @@ export class DatabricksRuntime implements Disposable {
                 .map((arg) => this.escapePythonString(arg))
                 .join("', '")}'];`,
 
+            `import os; ${envVarSetCmds.join("; ")};`,
+
             // Set log level to "ERROR". See https://kb.databricks.com/notebooks/cmd-c-on-object-id-p0.html
             `import logging; logger = spark._jvm.org.apache.log4j; logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)`,
             ...programLines,
         ].join("\n");
     }
 
     private escapePythonString(str: string): string {
-        return str.replace(/'/g, "\\'").replace(/\\/g, "\\\\");
+        return str.replace(/\\/g, "\\\\").replace(/'/g, "\\'");
     }
 
     public async disconnect(): Promise<void> {

packages/databricks-vscode/src/sync/CodeSynchronizer.ts

@@ -56,6 +56,7 @@ export class CodeSynchronizer implements Disposable {
     }
 
     async start(syncType: "full" | "incremental") {
+        this._state = "IN_PROGRESS";
         const task = new SyncTask(
             this.connection,
             this.cli,