Merge pull request #120 from barrywhart/bhart-add_sqlfluff_format

Update "format document" and "format selection to use dbt-core-interface

src/extension.ts

@@ -93,6 +93,18 @@ export const activate = (context: vscode.ExtensionContext) => {
   };
   context.subscriptions.push(vscode.commands.registerCommand(fixCommand, fixCommandHandler));
 
+  if (Configuration.dbtInterfaceEnabled()) {
+      // When dbt-core-interface is enabled, adds a "Format document with
+      // sqlfluff" button to the lower right corner of the VS Code window. Use
+      // of the word "format" (vs "fix") is deliberate, as the button hits the
+      // dbt-core-interface "/format" endpoint, equivalent to "sqlfluff format".
+      const customStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left);
+      customStatusBarItem.text = "$(list-selection) Format SQL";
+      customStatusBarItem.tooltip = "Format document with sqlfluff";
+      customStatusBarItem.command = fixCommand;
+      customStatusBarItem.show();
+  }
+
   const formatSelection = "sqlfluff.format.selection";
   const formatSelectionHandler = async () => {
     if (vscode.window.activeTextEditor) {

src/features/helper/dbtInterface.ts

@@ -64,6 +64,22 @@ export class DbtInterface {
     return url;
   }
 
+  public getFormatURL(): string {
+    // This endpoint is equivalent to "sqlfluff format". The behavior is
+    // _similar_ to "sqlfluff fix", but it applies a different set of rules.
+    // https://docs.sqlfluff.com/en/stable/cli.html#sqlfluff-format
+    let url = `http://${Configuration.dbtInterfaceHost()}:${Configuration.dbtInterfacePort()}/format?sql_path=${this.sql_path}`;
+    if (this.sql !== undefined) {
+      url = `http://${Configuration.dbtInterfaceHost()}:${Configuration.dbtInterfacePort()}/format?`;
+    }
+
+    if (this.extra_config_path) {
+      url += `&extra_config_path=${this.extra_config_path}`;
+    }
+
+    return url;
+  }
+
   public async healthCheck(): Promise<any> {
     const abortController = new AbortController();
     const timeoutHandler = setTimeout(() => {
@@ -134,7 +150,64 @@ export class DbtInterface {
       );
     } catch (error) {
       Utilities.appendHyphenatedLine();
-      Utilities.outputChannel.appendLine("Raw dbt-omsosis /lint error response:");
+      Utilities.outputChannel.appendLine("Raw dbt-core-interface /lint error response:");
+      Utilities.appendHyphenatedLine();
+      Utilities.outputChannel.appendLine(error as string);
+      Utilities.appendHyphenatedLine();
+
+      clearTimeout(timeoutHandler);
+      return failedToReachServerError;
+    }
+    clearTimeout(timeoutHandler);
+    return await response.json() as T;
+  }
+
+  public async format<T>(timeout = 25000) {
+    const failedToReachServerError: DbtInterfaceErrorContainer = {
+      error: {
+        code: DbtInterfaceErrorCode.FailedToReachServer,
+        message: "Query failed to reach dbt sync server.",
+        data: {
+          "error": `Is the server listening on the http://${Configuration.dbtInterfaceHost()}:${Configuration.dbtInterfacePort()} address?`,
+        },
+      },
+    };
+
+    const projectNotRegisteredError: DbtInterfaceErrorContainer = {
+      error: {
+        code: DbtInterfaceErrorCode.FailedToReachServer,
+        message: "dbt project not registered",
+        data: {
+          "error": "",
+        },
+      },
+    };
+
+    if (!await this.healthCheck()) {
+      Utilities.appendHyphenatedLine();
+      Utilities.outputChannel.appendLine("Unhealthy dbt project:");
+      Utilities.appendHyphenatedLine();
+      return projectNotRegisteredError;
+    }
+
+    const abortController = new AbortController();
+    const timeoutHandler = setTimeout(() => {
+      abortController.abort();
+    }, timeout);
+    let response: Response;
+
+    try {
+      response = await fetch(
+        encodeURI(this.getFormatURL()),
+        {
+          method: "POST",
+          signal: abortController.signal as AbortSignal,
+          body: this.sql,
+        },
+      );
+    } catch (error) {
+      Utilities.appendHyphenatedLine();
+      Utilities.outputChannel.appendLine("Raw dbt-core-interface /format error response:");
       Utilities.appendHyphenatedLine();
       Utilities.outputChannel.appendLine(error as string);
       Utilities.appendHyphenatedLine();

src/features/providers/formatter/rangeFormat.ts

@@ -2,6 +2,7 @@ import * as fs from "fs";
 import * as vscode from "vscode";
 
 import Configuration from "../../helper/configuration";
+import { DbtInterface } from "../../helper/dbtInterface";
 import Utilities from "../../helper/utilities";
 import SQLFluff from "../sqlfluff";
 import CommandOptions from "../types/commandOptions";
@@ -72,24 +73,55 @@ export class FormatSelectionProvider {
       return [];
     }
 
-    const commandOptions: CommandOptions = {
-      filePath: filePath,
-      fileContents: document.getText(lineRange),
-    };
-
-    const result = await SQLFluff.run(
-      workingDirectory,
-      CommandType.FIX,
-      Configuration.formatFileArguments(),
-      commandOptions,
-    );
+    let lines = undefined;
+    if (!Configuration.dbtInterfaceEnabled()) {
+      // Format the selection using sqlfluff CLI
+      const commandOptions: CommandOptions = {
+        filePath: filePath,
+        fileContents: document.getText(lineRange),
+      };
+
+      const result = await SQLFluff.run(
+        workingDirectory,
+        CommandType.FIX,
+        Configuration.formatFileArguments(),
+        commandOptions,
+      );
+
+      if (!result.succeeded) {
+        throw new Error("Command failed to execute, check logs for details");
+      }
 
-    if (!result.succeeded) {
-      throw new Error("Command failed to execute, check logs for details");
+      lines = FormatHelper.parseLines(result.lines);
+    } else {
+       // Format the selection using dbt-core-interface
+      const dbtInterface = new DbtInterface(
+        document.getText(lineRange),
+        workingDirectory,
+        Configuration.config(),
+      );
+
+      Utilities.outputChannel.appendLine("\n--------------------Executing Command--------------------\n");
+      Utilities.outputChannel.appendLine(dbtInterface.getFormatURL());
+      Utilities.appendHyphenatedLine();
+
+      const response: any = await dbtInterface.format();
+      Utilities.outputChannel.appendLine("Raw DBT-Interface /format output:");
+      Utilities.appendHyphenatedLine();
+      Utilities.outputChannel.appendLine(JSON.stringify(response, undefined, 2));
+      Utilities.appendHyphenatedLine();
+
+      const code = response?.error?.code ?? 0;
+      const succeeded = code === 0;
+      if (succeeded) {
+        // response.sql is a multiline string. Split it into an array of lines.
+        // This is similar to FormatHelper.parseLines(), but it does not look
+        // for SQLFluff messages in the text because, given this is API output,
+        // there won't be any.
+        lines = response.sql.split(/\r?\n|\r|\n/g);
+      }
     }
 
-    let lines = FormatHelper.parseLines(result.lines);
-
     const leadingWhitespace = document.lineAt(range.start.line).firstNonWhitespaceCharacterIndex + 1;
     lines = lines ? FormatHelper.addLeadingWhitespace(lines, document.languageId, leadingWhitespace) : undefined;
 

src/features/providers/sqlfluff.ts

@@ -25,6 +25,58 @@ export default class SQLFluff {
       throw new Error("You must supply either a target file path or the file contents to scan");
     }
 
+    if (Configuration.dbtInterfaceEnabled() && command === CommandType.FIX) {
+      // Handles CommandType.FIX when dbt-core-interface is enabled.
+      // TRICKY: Note that this actually hits the dbt-core-interface /format
+      // endpoint. This is a deliberate choice, but may look odd to readers of
+      // the code.
+      const dbtInterface = new DbtInterface(
+        undefined,
+        options.workspacePath ?? options.filePath,
+        Configuration.config(),
+      );
+
+      Utilities.outputChannel.appendLine("\n--------------------Executing Command--------------------\n");
+      Utilities.outputChannel.appendLine(dbtInterface.getFormatURL());
+      Utilities.appendHyphenatedLine();
+
+      const response: any = await dbtInterface.format();
+      const output: FilePath[] = [
+        {
+          filepath: options.filePath,
+          // The server returns a message field which contains any errors.
+          // Should we display this to the user in the error handling block
+          // below?
+          //message: response.message ?? "",
+          // The "FilePath" interface requires a "violations" field, but /format
+          // doesn't return any violations. We'll just return an empty array.
+          violations: [],
+        },
+      ];
+
+      Utilities.outputChannel.appendLine("Raw DBT-Interface /format output:");
+      Utilities.appendHyphenatedLine();
+      Utilities.outputChannel.appendLine(JSON.stringify(response, undefined, 2));
+      Utilities.appendHyphenatedLine();
+
+      return new Promise<CommandOutput>((resolve) => {
+        const code = response?.error?.code ?? 0;
+        const succeeded = code === 0;
+        if (!succeeded && !Configuration.suppressNotifications()) {
+          const message = response?.error?.message ?? "DBT-Interface formatting error.";
+          const detail = response?.error?.data?.error ?? "";
+
+          vscode.window.showErrorMessage([message, detail].join("\n"));
+        }
+
+        resolve({
+          // 0 = all good, 1 = format passed but contains unfixable linting violations, 65 = lint passed but found errors
+          succeeded: succeeded,
+          lines: [],
+        });
+      });
+    }
+
     // This is an unlikely scenario, but we should limit the amount of processes happening at once.
     while (SQLFluff.childProcesses.length > 10) {
       const process = SQLFluff.childProcesses.shift();