Fix run on cluster for multi file setups (#463)

Author: kartikgupta-db

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

@@ -130,15 +130,9 @@ export class DatabricksRuntime implements Disposable {
                 await executionContext.destroy();
             });
 
-            const currentFsRoot =
-                this.connection.databricksWorkspace?.currentFsRoot;
             // We wait for sync to complete so that the local files are consistant
             // with the remote repo files
-            log(
-                `Synchronizing code to ${syncDestination.remoteUri.relativePath(
-                    new RemoteUri(currentFsRoot ?? "")
-                )} ...`
-            );
+            log(`Synchronizing code to ${syncDestination.remoteUri.path} ...`);
 
             this.disposables.push(
                 this.codeSynchronizer.onDidChangeState((state) => {
@@ -252,7 +246,8 @@ export class DatabricksRuntime implements Disposable {
         );
 
         const argv = [
-            syncDestination.localToRemote(new LocalUri(Uri.file(program))).path,
+            syncDestination.localToRemote(new LocalUri(Uri.file(program)))
+                .workspacePrefixPath,
             ...args,
         ];
 
@@ -276,7 +271,7 @@ export class DatabricksRuntime implements Disposable {
         );
         bootstrap = bootstrap.replace(
             '"REPO_PATH"',
-            `"${syncDestination.remoteUri.path}"`
+            `"${syncDestination.remoteUri.workspacePrefixPath}"`
         );
         bootstrap = bootstrap.replace(
             "args = []",

packages/databricks-vscode/src/test/e2e/run_on_cluster.e2e.ts

@@ -35,9 +35,18 @@ describe("Run python on cluster", async function () {
                 workspacePath: process.env["TEST_REPO_PATH"],
             })
         );
+        await fs.mkdir(path.join(projectDir, "nested"));
         await fs.writeFile(
-            path.join(projectDir, "hello.py"),
-            `spark.sql('SELECT "hello world"').show()`
+            path.join(projectDir, "nested", "hello.py"),
+            [`from lib import func`, "func(spark)"].join("\n")
+        );
+
+        await fs.writeFile(
+            path.join(projectDir, "lib.py"),
+            [
+                "def func(spark):",
+                `\tspark.sql('SELECT "hello world"').show()`,
+            ].join("\n")
         );
         await dismissNotifications();
     });