Add dedicated loader for command execution API (#257)

This fixes the wrong line numbers in python stack traces

Author: fjakobs

packages/databricks-vscode/resources/python/bootstrap.py

@@ -0,0 +1,44 @@
+import runpy
+import sys
+import os
+
+# values will be injected by the runner
+python_file = "PYTHON_FILE"
+repo_path = "REPO_PATH"
+args = []
+env = {}
+
+# change working directory
+os.chdir(os.path.dirname(python_file))
+
+# update python path
+sys.path.append(repo_path)
+sys.path.append(os.path.dirname(python_file))
+
+# inject command line arguments
+sys.argv = args
+
+# inject environment variables
+for key in env:
+    os.environ[key] = env[key]
+
+# provide spark globals
+user_ns = {
+    "display": display,
+    "displayHTML": displayHTML,
+    "dbutils": dbutils,
+    "table": table,
+    "sql": sql,
+    "udf": udf,
+    "getArgument": getArgument,
+    "sc": sc,
+    "spark": spark,
+    "sqlContext": sqlContext,
+}
+
+# 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)
+
+runpy.run_path(python_file, run_name="__main__", init_globals=user_ns)
+None

packages/databricks-vscode/webview-ui/job.html renamed to packages/databricks-vscode/resources/webview-ui/job.html

@@ -129,7 +129,8 @@ export async function activate(
     const runCommands = new RunCommands(connectionManager, synchronizer);
     const debugFactory = new DatabricksDebugAdapterFactory(
         connectionManager,
-        synchronizer
+        synchronizer,
+        context
     );
     const debugWorkflowFactory = new DatabricksWorkflowDebugAdapterFactory(
         connectionManager,

packages/databricks-vscode/src/extension.ts

@@ -8,22 +8,20 @@ import {
 } from "@vscode/debugadapter";
 import {DebugProtocol} from "@vscode/debugprotocol";
 import {basename} from "node:path";
-import {TextDecoder} from "node:util";
 
 import {
     DebugAdapterDescriptor,
     DebugAdapterDescriptorFactory,
     DebugAdapterInlineImplementation,
     Disposable,
     ProviderResult,
-    Uri,
     commands,
     window,
-    workspace,
+    ExtensionContext,
 } from "vscode";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {CodeSynchronizer} from "../sync/CodeSynchronizer";
-import {DatabricksRuntime, FileAccessor} from "./DatabricksRuntime";
+import {DatabricksRuntime} from "./DatabricksRuntime";
 import {Subject} from "./Subject";
 
 /**
@@ -48,42 +46,39 @@ export class DatabricksDebugAdapterFactory
 {
     constructor(
         private connection: ConnectionManager,
-        private codeSynchroniser: CodeSynchronizer
+        private codeSynchroniser: CodeSynchronizer,
+        private context: ExtensionContext
     ) {}
 
     dispose() {}
 
     createDebugAdapterDescriptor(): ProviderResult<DebugAdapterDescriptor> {
         return new DebugAdapterInlineImplementation(
-            new DatabricksDebugSession(this.connection, this.codeSynchroniser)
+            new DatabricksDebugSession(
+                this.connection,
+                this.codeSynchroniser,
+                this.context
+            )
         );
     }
 }
 
-export const workspaceFileAccessor: FileAccessor = {
-    async readFile(path: string): Promise<string> {
-        const uri = Uri.file(path);
-
-        const bytes = await workspace.fs.readFile(uri);
-        return new TextDecoder().decode(bytes);
-    },
-};
-
 export class DatabricksDebugSession extends LoggingDebugSession {
     private runtime: DatabricksRuntime;
     private _configurationDone = new Subject();
     private disposables: Disposable[] = [];
 
     constructor(
         connection: ConnectionManager,
-        codeSynchronizer: CodeSynchronizer
+        codeSynchronizer: CodeSynchronizer,
+        context: ExtensionContext
     ) {
         super();
 
         this.runtime = new DatabricksRuntime(
             connection,
-            workspaceFileAccessor,
-            codeSynchronizer
+            codeSynchronizer,
+            context
         );
 
         this.disposables.push(

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

@@ -2,21 +2,18 @@
 
 import assert from "assert";
 import {mock, when, instance, anything, verify, capture} from "ts-mockito";
-import {Disposable, Uri} from "vscode";
+import {Disposable, ExtensionContext, Uri} from "vscode";
 import {
     Cluster,
     Command,
     ExecutionContext,
     Repo,
 } from "@databricks/databricks-sdk";
-import {
-    DatabricksRuntime,
-    FileAccessor,
-    OutputEvent,
-} from "./DatabricksRuntime";
+import {DatabricksRuntime, OutputEvent} from "./DatabricksRuntime";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {SyncDestination} from "../configuration/SyncDestination";
 import {CodeSynchronizer} from "../sync/CodeSynchronizer";
+import path from "node:path";
 
 describe(__filename, () => {
     let disposables: Array<Disposable>;
@@ -66,20 +63,19 @@ describe(__filename, () => {
         );
         when(connectionManagerMock.syncDestination).thenReturn(syncDestination);
 
-        const fileAccessor: FileAccessor = {
-            async readFile(): Promise<string> {
-                return "print('43')";
-            },
-        };
-
         const connectionManager = instance<ConnectionManager>(
             connectionManagerMock
         );
 
+        const extensionContextMock = mock<ExtensionContext>();
+        when(extensionContextMock.extensionUri).thenReturn(
+            Uri.file(path.join(__dirname, "..", ".."))
+        );
+
         runtime = new DatabricksRuntime(
             connectionManager,
-            fileAccessor,
-            mock(CodeSynchronizer)
+            mock(CodeSynchronizer),
+            instance(extensionContextMock)
         );
     });
 
@@ -102,7 +98,6 @@ describe(__filename, () => {
         await runtime.start("/Desktop/workspaces/hello.py", [], {});
 
         verify(connectionManagerMock.waitForConnect()).called();
-
         assert.equal(outputs.length, 6);
         assert(outputs[0].text.includes("Connecting to cluster"));
         assert(
@@ -114,34 +109,18 @@ describe(__filename, () => {
         assert(outputs[5].text.includes("Done"));
     });
 
-    it("should have the right code with env vars", async () => {
+    it("should inject environment variables", async () => {
         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/[email protected]/test");
-import sys; sys.path.append("/Workspace/Repos/[email protected]/test")
-import sys; sys.argv = ['/Workspace/Repos/[email protected]/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')`
-        );
+        assert(code.includes(`env = {"TEST":"TEST"}`));
     });
 
     it("should have the right code without env vars", async () => {
         await runtime.start("/Desktop/workspaces/hello.py", [], {});
 
         const code = capture(executionContextMock.execute).first()[0];
-        assert.equal(
-            code,
-            `import os; os.chdir("/Workspace/Repos/[email protected]/test");
-import sys; sys.path.append("/Workspace/Repos/[email protected]/test")
-import sys; sys.argv = ['/Workspace/Repos/[email protected]/test/hello.py'];
-
-import logging; logger = spark._jvm.org.apache.log4j; logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
-print('43')`
-        );
+        assert(code.includes(`env = {}`));
     });
 
     it("should handle failed executions", async () => {