VSCode: Add task to start dbx sync (#53)

Author: fjakobs

packages/databricks-vscode/package.json

@@ -42,6 +42,9 @@
         "onCommand:databricks.connection.openDatabricksConfigFile",
         "onCommand:databricks.connection.attachCluster",
         "onCommand:databricks.connection.detachCluster",
+        "onCommand:databricks.cli.startSync",
+        "onCommand:databricks.cli.startSyncFull",
+        "onCommand:databricks.cli.stopSync",
         "onCommand:databricks.cluster.refresh",
         "onCommand:databricks.cluster.filterByAll",
         "onCommand:databricks.cluster.filterByMe",
@@ -81,6 +84,18 @@
                 "title": "Detach cluster",
                 "icon": "$(debug-disconnect)"
             },
+            {
+                "command": "databricks.cli.startSync",
+                "title": "Databricks: Start synchronization"
+            },
+            {
+                "command": "databricks.cli.startSyncFull",
+                "title": "Databricks: Start synchronization (full sync)"
+            },
+            {
+                "command": "databricks.cli.stopSync",
+                "title": "Databricks: Stop synchronization"
+            },
             {
                 "command": "databricks.cluster.filterByAll",
                 "title": "All"
@@ -199,6 +214,22 @@
                 "label": "Filter clusters ...",
                 "icon": "$(filter)"
             }
+        ],
+        "problemMatchers": [
+            {
+                "name": "bricks-sync",
+                "owner": "bricks",
+                "fileLocation": "relative",
+                "pattern": {
+                    "regexp": "Putting (/Repos/.*)$",
+                    "file": 1
+                },
+                "background": {
+                    "activeOnStart": true,
+                    "beginsPattern": "Syncing from",
+                    "endsPattern": "Done(\\. Watching for changes\\.\\.\\.)$"
+                }
+            }
         ]
     },
     "vsce": {

packages/databricks-vscode/src/cli/CliCommands.ts

@@ -0,0 +1,87 @@
+import {
+    ProcessExecution,
+    Task,
+    TaskExecution,
+    TaskRevealKind,
+    tasks,
+    TaskScope,
+    window,
+    workspace,
+} from "vscode";
+import {ConnectionManager} from "../configuration/ConnectionManager";
+import {CliWrapper} from "./CliWrapper";
+
+export class CliCommands {
+    currentTaskExecution?: TaskExecution;
+
+    constructor(
+        private connection: ConnectionManager,
+        private cli: CliWrapper
+    ) {}
+
+    stopSyncCommand() {
+        return () => {
+            if (this.currentTaskExecution) {
+                this.currentTaskExecution.terminate();
+            }
+        };
+    }
+
+    startSyncCommand(syncType: "full" | "incremental") {
+        return async () => {
+            const workspacePath = workspace.rootPath;
+            const me = this.connection.me;
+            const pathMapper = this.connection.pathMapper;
+            const profile = this.connection.profile;
+
+            if (!workspacePath) {
+                window.showErrorMessage(
+                    "Can't start sync: No workspace opened!"
+                );
+                return;
+            }
+
+            if (!me || !profile) {
+                window.showErrorMessage(
+                    "Can't start sync: Databricks connection not configured!"
+                );
+                return;
+            }
+            if (!pathMapper) {
+                window.showErrorMessage(
+                    "Can't start sync: Databricks synchronization destination not configured!"
+                );
+                return;
+            }
+
+            if (this.currentTaskExecution) {
+                this.currentTaskExecution.terminate();
+            }
+
+            const {command, args} = this.cli.getSyncCommand(
+                profile,
+                me,
+                pathMapper,
+                syncType
+            );
+
+            const task = new Task(
+                {
+                    type: "bricks",
+                },
+                TaskScope.Workspace,
+                "$(rocket) Databricks Sync",
+                "bricks",
+                new ProcessExecution(command, args, {
+                    cwd: workspacePath,
+                })
+            );
+            task.isBackground = true;
+            task.problemMatchers = ["$bricks-sync"];
+            task.presentationOptions.echo = true;
+            task.presentationOptions.reveal = TaskRevealKind.Always;
+
+            this.currentTaskExecution = await tasks.executeTask(task);
+        };
+    }
+}

packages/databricks-vscode/src/cli/CliWrapper.test.ts

@@ -0,0 +1,65 @@
+import * as assert from "assert";
+import {Uri} from "vscode";
+import {PathMapper} from "../configuration/PathMapper";
+
+import {CliWrapper} from "./CliWrapper";
+
+describe(__filename, () => {
+    it("should create sync command", () => {
+        const cli = new CliWrapper();
+        const mapper = new PathMapper(
+            Uri.file(
+                "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices"
+            ),
+            Uri.file("/Users/fabian.jakobs/Desktop/notebook-best-practices")
+        );
+
+        const {command, args} = cli.getSyncCommand(
+            "DEFAULT",
+            "fabian@databricks.com",
+            mapper,
+            "incremental"
+        );
+
+        assert.equal(
+            [command, ...args].join(" "),
+            "dbx sync repo --profile DEFAULT --user fabian@databricks.com --dest-repo notebook-best-practices"
+        );
+    });
+
+    it("should create full sync command", () => {
+        const cli = new CliWrapper();
+        const mapper = new PathMapper(
+            Uri.file(
+                "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices"
+            ),
+            Uri.file("/Users/fabian.jakobs/Desktop/notebook-best-practices")
+        );
+
+        const {command, args} = cli.getSyncCommand(
+            "DEFAULT",
+            "fabian@databricks.com",
+            mapper,
+            "full"
+        );
+
+        assert.equal(
+            [command, ...args].join(" "),
+            "dbx sync repo --profile DEFAULT --user fabian@databricks.com --dest-repo notebook-best-practices --full-sync"
+        );
+    });
+
+    it("should create an 'add profile' command", () => {
+        const cli = new CliWrapper();
+
+        const {command, args} = cli.getAddProfileCommand(
+            "DEFAULT",
+            new URL("https://databricks.com")
+        );
+
+        assert.equal(
+            [command, ...args].join(" "),
+            "databricks configure --profile DEFAULT --host https://databricks.com/ --token"
+        );
+    });
+});

packages/databricks-vscode/src/cli/CliWrapper.ts

@@ -1,4 +1,10 @@
 import {spawn} from "child_process";
+import {PathMapper} from "../configuration/PathMapper";
+
+interface Command {
+    command: string;
+    args: string[];
+}
 
 /**
  * Entrypoint for all wrapped CLI commands
@@ -9,26 +15,58 @@ import {spawn} from "child_process";
 export class CliWrapper {
     constructor() {}
 
+    /**
+     * Constructs the dbx sync command
+     */
+    getSyncCommand(
+        profile: string,
+        me: string,
+        pathMapper: PathMapper,
+        syncType: "full" | "incremental"
+    ): Command {
+        const command = "dbx";
+        const args = [
+            "sync",
+            "repo",
+            "--profile",
+            profile,
+            "--user",
+            me,
+            "--dest-repo",
+            pathMapper.remoteWorkspaceName,
+        ];
+
+        if (syncType === "full") {
+            args.push("--full-sync");
+        }
+
+        return {command, args};
+    }
+
+    getAddProfileCommand(profile: string, host: URL): Command {
+        return {
+            command: "databricks",
+            args: [
+                "configure",
+                "--profile",
+                profile,
+                "--host",
+                host.href,
+                "--token",
+            ],
+        };
+    }
+
     async addProfile(
         name: string,
         host: URL,
         token: string
     ): Promise<{stdout: string; stderr: string}> {
         return new Promise((resolve, reject) => {
-            let child = spawn(
-                "databricks",
-                [
-                    "configure",
-                    "--profile",
-                    name,
-                    "--host",
-                    host.href,
-                    "--token",
-                ],
-                {
-                    stdio: ["pipe", 0, 0],
-                }
-            );
+            const {command, args} = this.getAddProfileCommand(name, host);
+            let child = spawn(command, args, {
+                stdio: ["pipe", 0, 0],
+            });
 
             child.stdin!.write(token);
             child.stdin!.end();

packages/databricks-vscode/src/configuration/PathMapper.test.ts

@@ -38,4 +38,15 @@ describe(__filename, () => {
             "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices/jobs"
         );
     });
+
+    it("should get repo name", async () => {
+        let mapper = new PathMapper(
+            Uri.file(
+                "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices"
+            ),
+            Uri.file("/Users/fabian.jakobs/Desktop/notebook-best-practices")
+        );
+
+        assert.equal(mapper.remoteWorkspaceName, "notebook-best-practices");
+    });
 });

packages/databricks-vscode/src/configuration/PathMapper.ts

@@ -6,11 +6,7 @@ import {Uri} from "vscode";
  * on the Databricks driver
  */
 export class PathMapper {
-    readonly repo: string;
-
-    constructor(readonly repoPath: Uri, readonly workspacePath: Uri) {
-        this.repo = path.basename(repoPath.path);
-    }
+    constructor(readonly repoPath: Uri, readonly workspacePath: Uri) {}
 
     localToRemoteDir(localPath: Uri): string {
         return path.dirname(this.localToRemote(localPath));
@@ -20,4 +16,8 @@ export class PathMapper {
         let relativePath = localPath.path.replace(this.workspacePath.path, "");
         return Uri.joinPath(this.repoPath, relativePath).path;
     }
+
+    get remoteWorkspaceName(): string {
+        return path.basename(this.repoPath.path);
+    }
 }

packages/databricks-vscode/src/extension.test.ts

@@ -7,6 +7,7 @@ import {ClusterModel} from "./cluster/ClusterModel";
 import {ClusterCommands} from "./cluster/ClusterCommands";
 import {ConfigurationDataProvider} from "./configuration/ConfigurationDataProvider";
 import {WorkflowCommands} from "./workflow/WorkflowCommands";
+import {CliCommands} from "./cli/CliCommands";
 
 export function activate(context: ExtensionContext) {
     let cli = new CliWrapper();
@@ -77,6 +78,7 @@ export function activate(context: ExtensionContext) {
     context.subscriptions.push(
         clusterModel,
         clusterTreeDataProvider,
+        window.registerTreeDataProvider("clusterList", clusterTreeDataProvider),
 
         commands.registerCommand(
             "databricks.cluster.refresh",
@@ -99,7 +101,26 @@ export function activate(context: ExtensionContext) {
             clusterCommands
         )
     );
-    window.registerTreeDataProvider("clusterList", clusterTreeDataProvider);
+
+    // Tasks
+    const cliCommands = new CliCommands(connectionManager, cli);
+    context.subscriptions.push(
+        commands.registerCommand(
+            "databricks.cli.startSync",
+            cliCommands.startSyncCommand("incremental"),
+            cliCommands
+        ),
+        commands.registerCommand(
+            "databricks.cli.startSyncFull",
+            cliCommands.startSyncCommand("full"),
+            cliCommands
+        ),
+        commands.registerCommand(
+            "databricks.cli.stopSync",
+            cliCommands.stopSyncCommand(),
+            cliCommands
+        )
+    );
 }
 
 // this method is called when your extension is deactivated