Added cancellation support for model runs

Author: mdesmet

src/dbt_client/commandProcessExecution.ts

@@ -1,11 +1,26 @@
 import { ChildProcess, spawn } from "child_process";
-import { EventEmitter } from "vscode";
+import { CancellationToken, Disposable, EventEmitter } from "vscode";
 
-export class CommandProcessExecution {
+export class CommandProcessExecution implements Disposable {
   private readonly commandProcess: ChildProcess;
+  private disposables: Disposable[] = [];
 
-  constructor(command: string, args?: string[], cwd?: string) {
+  constructor(
+    command: string,
+    args?: string[],
+    cwd?: string,
+    token?: CancellationToken
+  ) {
     this.commandProcess = spawn(command, args, { cwd: cwd });
+    if (token !== undefined) {
+      this.disposables.push(token.onCancellationRequested(() => {
+        this.commandProcess.kill('SIGINT');
+      }));
+    }
+  }
+  
+  dispose() {
+    this.disposables.forEach(disposable => disposable.dispose());
   }
 
   async complete(): Promise<string> {
@@ -34,7 +49,9 @@ export class CommandProcessExecution {
     });
   }
 
-  async completeWithTerminalOutput(writeEmitter: EventEmitter<string>): Promise<void> {
+  async completeWithTerminalOutput(
+    writeEmitter: EventEmitter<string>
+  ): Promise<void> {
     return new Promise((resolve, reject) => {
       this.commandProcess.stdout!.on("data", (chunk) => {
         writeEmitter.fire(`\r${this.formatText(chunk.toString())}`);
@@ -43,7 +60,7 @@ export class CommandProcessExecution {
         writeEmitter.fire(`\r${this.formatText(chunk.toString())}`);
       });
       this.commandProcess.once("close", () => {
-        writeEmitter.fire('\r\n\r\n');
+        writeEmitter.fire("\r\n\r\n");
         resolve();
       });
 

src/dbt_client/dbtClient.ts

@@ -1,4 +1,4 @@
-import { Disposable, EventEmitter, Terminal, window } from "vscode";
+import { CancellationToken, Disposable, EventEmitter, Terminal, window } from "vscode";
 import {
   OnSourceFileChanged,
   SourceFileChangedEvent,
@@ -16,7 +16,6 @@ export class DBTClient implements OnSourceFileChanged, Disposable {
   private readonly writeEmitter = new EventEmitter<string>();
   private readonly queue: DBTCommandQueue = new DBTCommandQueue();
   private dbtInstalled?: boolean;
-  private notYetShownErrorMessage = true;
   private terminal?: Terminal;
 
   constructor(pythonPath: string) {
@@ -40,6 +39,7 @@ export class DBTClient implements OnSourceFileChanged, Disposable {
     try {
       this.raiseDBTInstallationCheckEvent();
       await checkDBTInstalledProcess.complete();
+      checkDBTInstalledProcess.dispose();
     } catch (err) {
       if (err.match(DBTClient.IS_INSTALLED)) {
         this.checkIfDBTIsUpToDate(err);
@@ -51,47 +51,52 @@ export class DBTClient implements OnSourceFileChanged, Disposable {
 
   addCommandToQueue(command: DBTCommand) {
     if (!this.dbtInstalled) {
-      this.notYetShownErrorMessage &&
+      if (command.focus) {
         window.showErrorMessage(
           "Please ensure DBT is installed in your selected Python environment."
         );
-      this.notYetShownErrorMessage = false;
+      }
       return;
     }
 
     this.queue.addToQueue({
-      command: () => this.executeCommandImmediately(command),
+      command: (token) => this.executeCommandImmediately(command, token),
       statusMessage: command.statusMessage,
+      focus: command.focus
     });
   }
 
-  executeCommandImmediately(command: DBTCommand) {
-    return this.executeCommand(command).completeWithTerminalOutput(
+  async executeCommandImmediately(command: DBTCommand, token?: CancellationToken) {
+    const process = this.executeCommand(command, token);
+    await process.completeWithTerminalOutput(
       this.writeEmitter
     );
+    process.dispose();
   }
 
-  private executeCommand(command: DBTCommand): CommandProcessExecution {
+  private executeCommand(command: DBTCommand, token?: CancellationToken): CommandProcessExecution {
     const { args, cwd } = command.processExecutionParams;
-    if(this.terminal === undefined) {
+    if (this.terminal === undefined) {
       this.terminal = window.createTerminal({
-        name: 'DBT',
+        name: "DBT",
         pty: {
           onDidWrite: this.writeEmitter.event,
-          open: () => this.writeEmitter.fire(''),
-          close: () =>  {
+          open: () => this.writeEmitter.fire(""),
+          close: () => {
             this.terminal?.dispose();
             this.terminal = undefined;
-          }
-        }
+          },
+        },
       });
     }
-    this.writeEmitter.fire(`\r> Executing task:  ${command.commandAsString}\n\r\n\r`);
-    
+    this.writeEmitter.fire(
+      `\r> Executing task:  ${command.commandAsString}\n\r\n\r`
+    );
+
     if (command.focus) {
       this.terminal.show(true);
     }
-    return new CommandProcessExecution(this.pythonPath, args, cwd);
+    return new CommandProcessExecution(this.pythonPath, args, cwd, token);
   }
 
   private raiseDBTInstallationCheckEvent() {

src/dbt_client/dbtCommandQueue.ts

@@ -1,8 +1,9 @@
-import { ProgressLocation, window } from "vscode";
+import { CancellationToken, ProgressLocation, window } from "vscode";
 
 interface Command{
-  command: () => Promise<void>;
+  command: (token: CancellationToken) => Promise<void>;
   statusMessage: string;
+  focus?: boolean;
 }
 export class DBTCommandQueue {
   private queue: Command[] = [];
@@ -16,14 +17,14 @@ export class DBTCommandQueue {
   private async pickCommandToRun(): Promise<void> {
     if (!this.running && this.queue.length > 0) {
       this.running = true;
-      const { command, statusMessage } = this.queue.shift()!;
+      const { command, statusMessage, focus } = this.queue.shift()!;
 
       await window.withProgress({
-          location: ProgressLocation.Window,
-          cancellable: false,
+          location: focus ? ProgressLocation.Notification : ProgressLocation.Window,
+          cancellable: true,
           title: statusMessage,
-      }, async () => {
-          await command();
+      }, async (_, token) => {
+          await command(token);
       });
 
       this.running = false;   

src/manifest/dbtProjectContainer.ts

@@ -19,7 +19,6 @@ export class DbtProjectContainer implements Disposable {
   private manifestCacheChangedHandlers: OnManifestCacheChanged[] = [];
   private dbtInstallationFoundHandlers: OnDBTInstallationFound[] = [];
   private dbtWorkspaceFolders: DBTWorkspaceFolder[] = [];
-  private notYetShownDbtInstalledErrorMessage = true;
 
   constructor() {
     workspace.onDidChangeWorkspaceFolders(async (event) => {
@@ -89,7 +88,6 @@ export class DbtProjectContainer implements Disposable {
       const pythonEnvironment = await PythonEnvironment.getEnvironment();
 
       const pythonPath = pythonEnvironment.getPythonPath();
-      this.notYetShownDbtInstalledErrorMessage = true;
 
       if (pythonPath === undefined) {
         this.dbtClient = undefined;
@@ -114,11 +112,11 @@ export class DbtProjectContainer implements Disposable {
 
   addCommandToQueue(command: DBTCommand) {
     if (this.dbtClient === undefined) {
-      this.notYetShownDbtInstalledErrorMessage &&
+      if (command.focus) {
         window.showErrorMessage(
-          "Please ensure you have selected a Python interpreter with DBT installed."
+          "Can't run the command. Please ensure you have selected a Python interpreter with DBT installed."
         );
-      this.notYetShownDbtInstalledErrorMessage = false;
+      }
       return;
     }
     this.dbtClient.addCommandToQueue(command);
@@ -127,7 +125,7 @@ export class DbtProjectContainer implements Disposable {
   async installDBT() {
     if (this.dbtClient === undefined) {
       window.showErrorMessage(
-        "Please ensure you have selected a Python interpreter before installing DBT."
+        "Can't install DBT. Please ensure you have selected a Python interpreter before installing DBT."
       );
       return;
     }
@@ -140,7 +138,7 @@ export class DbtProjectContainer implements Disposable {
   async updateDBT() {
     if (this.dbtClient === undefined) {
       window.showErrorMessage(
-        "Please ensure you have selected a Python interpreter before updating DBT."
+        "Can't update DBT. Please ensure you have selected a Python interpreter before updating DBT."
       );
       return;
     }