fix: query bookmarks feedback changes (#1295)

* fix: show contact us for enabling bookmarks * fix: errors when test added for new column * fix: allow jinja autocompletions for new sql file * fix: open code in editor * fix: cleanup * fix: add title for e2e tests
AltimateAI · Jul 12, 2024 · d2a68ff · d2a68ff
1 parent 28a470d
commit d2a68ff
Show file tree

Hide file tree

Showing 14 changed files with 221 additions and 148 deletions.
diff --git a/src/code_lens_provider/virtualSqlCodeLensProvider.ts b/src/code_lens_provider/virtualSqlCodeLensProvider.ts
@@ -34,7 +34,10 @@ export class VirtualSqlCodeLensProvider implements CodeLensProvider {
     token: CancellationToken,
   ): CodeLens[] | Thenable<CodeLens[]> {
     // Enable this code lens only for adhoc query files created using command: dbtPowerUser.createSqlFile
-    if (document.uri.scheme !== "untitled" && document.languageId !== "sql") {
+    if (
+      document.uri.scheme !== "untitled" &&
+      document.languageId !== "jinja-sql"
+    ) {
       return [];
     }
 

diff --git a/src/commands/index.ts b/src/commands/index.ts
@@ -539,77 +539,81 @@ export class VSCodeCommands implements Disposable {
           this.dbtTerminal.logLine(`Error=${e}`);
         }
       }),
-      commands.registerCommand("dbtPowerUser.createSqlFile", async () => {
-        try {
-          const project =
-            await this.queryManifestService.getOrPickProjectFromWorkspace();
-          if (!project) {
-            window.showErrorMessage("No dbt project selected.");
-            return;
-          }
+      commands.registerCommand(
+        "dbtPowerUser.createSqlFile",
+        async ({ code, fileName }: { code?: string; fileName?: string }) => {
+          try {
+            const project =
+              await this.queryManifestService.getOrPickProjectFromWorkspace();
+            if (!project) {
+              window.showErrorMessage("No dbt project selected.");
+              return;
+            }
 
-          const uri = Uri.parse(
-            `${project.projectRoot}/poweruser-${getFormattedDateTime()}.sql`,
-          ).with({ scheme: "untitled" });
-          const annotationDecoration: TextEditorDecorationType =
-            window.createTextEditorDecorationType({
-              rangeBehavior: DecorationRangeBehavior.OpenOpen,
-            });
+            const fileNamePrefix = fileName || "poweruser";
+            const uri = Uri.parse(
+              `${project.projectRoot}/${fileNamePrefix}-${getFormattedDateTime()}.sql`,
+            ).with({ scheme: "untitled" });
+            const annotationDecoration: TextEditorDecorationType =
+              window.createTextEditorDecorationType({
+                rangeBehavior: DecorationRangeBehavior.OpenOpen,
+              });
 
-          const contentText =
-            "Enter your query here and execute it just like any dbt model file. This file is unsaved, you can either save it to your project or save it as a bookmark for later usage or share it with your team members.";
+            const contentText =
+              "Enter your query here and execute it just like any dbt model file. This file is unsaved, you can either save it to your project or save it as a bookmark for later usage or share it with your team members.";
 
-          const decorations = [
-            {
-              renderOptions: {
-                before: {
-                  color: "#666666",
-                  contentText,
-                  // hacking to add more css properties
-                  width: "90%;display: block;white-space: pre-line;",
+            const decorations = [
+              {
+                renderOptions: {
+                  before: {
+                    color: "#666666",
+                    contentText,
+                    // hacking to add more css properties
+                    width: "90%;display: block;white-space: pre-line;",
+                  },
                 },
+                range: new Range(2, 0, 2, 0),
               },
-              range: new Range(2, 0, 2, 0),
-            },
-          ];
+            ];
 
-          workspace.openTextDocument(uri).then((doc) => {
-            // set this to sql language so we can bind codelens and other features
-            languages.setTextDocumentLanguage(doc, "sql");
-            window.showTextDocument(doc).then((editor) => {
-              editor.edit((editBuilder) => {
-                const entireDocumentRange = new Range(
-                  doc.positionAt(0),
-                  doc.positionAt(doc.getText().length),
-                );
-                editBuilder.replace(entireDocumentRange, "\n");
+            workspace.openTextDocument(uri).then((doc) => {
+              // set this to sql language so we can bind codelens and other features
+              languages.setTextDocumentLanguage(doc, "jinja-sql");
+              window.showTextDocument(doc).then((editor) => {
+                editor.edit((editBuilder) => {
+                  const entireDocumentRange = new Range(
+                    doc.positionAt(0),
+                    doc.positionAt(doc.getText().length),
+                  );
+                  editBuilder.replace(entireDocumentRange, code || "\n");
 
-                editor.setDecorations(annotationDecoration, decorations);
-                setTimeout(() => {
-                  commands.executeCommand("cursorMove", {
-                    to: "up",
-                    by: "line",
-                    value: 1,
-                  });
-                }, 0);
-                const disposable = workspace.onDidChangeTextDocument((e) => {
-                  const activeEditor = window.activeTextEditor;
-                  if (activeEditor && e.document === editor.document) {
-                    if (activeEditor.document.getText().trim()) {
-                      activeEditor.setDecorations(annotationDecoration, []);
-                      disposable.dispose();
+                  editor.setDecorations(annotationDecoration, decorations);
+                  setTimeout(() => {
+                    commands.executeCommand("cursorMove", {
+                      to: "up",
+                      by: "line",
+                      value: 1,
+                    });
+                  }, 0);
+                  const disposable = workspace.onDidChangeTextDocument((e) => {
+                    const activeEditor = window.activeTextEditor;
+                    if (activeEditor && e.document === editor.document) {
+                      if (activeEditor.document.getText().trim()) {
+                        activeEditor.setDecorations(annotationDecoration, []);
+                        disposable.dispose();
+                      }
                     }
-                  }
+                  });
                 });
               });
             });
-          });
-        } catch (e) {
-          const message = (e as Error).message;
-          this.dbtTerminal.error("createSqlFile", message, e, true);
-          window.showErrorMessage(message);
-        }
-      }),
+          } catch (e) {
+            const message = (e as Error).message;
+            this.dbtTerminal.error("createSqlFile", message, e, true);
+            window.showErrorMessage(message);
+          }
+        },
+      ),
       commands.registerCommand("dbtPowerUser.sqlLineage", async () => {
         window.withProgress(
           {

diff --git a/src/dbtPowerUserExtension.ts b/src/dbtPowerUserExtension.ts
@@ -32,7 +32,7 @@ export class DBTPowerUserExtension implements Disposable {
   static DBT_SQL_SELECTOR = [
     { language: "jinja-sql", scheme: "file" },
     { language: "sql", scheme: "file" },
-    { language: "sql", scheme: "untitled" },
+    { language: "jinja-sql", scheme: "untitled" },
   ];
   static DBT_YAML_SELECTOR = [
     { language: "yaml", scheme: "file" },

diff --git a/src/webview_provider/docsEditPanel.ts b/src/webview_provider/docsEditPanel.ts
@@ -339,7 +339,7 @@ export class DocsEditViewPanel implements WebviewViewProvider {
       const testFullName: string = namespace ? `${namespace}.${name}` : name;
 
       const columnTestConfigFromYml = getColumnTestConfigFromYml(
-        existingColumn.tests,
+        existingColumn?.tests,
         kwargs,
         testFullName,
       );

diff --git a/src/webview_provider/queryResultPanel.ts b/src/webview_provider/queryResultPanel.ts
@@ -92,6 +92,7 @@ enum InboundCommand {
   GetQueryTabData = "getQueryTabData",
   RunAdhocQuery = "runAdhocQuery",
   ViewResultSet = "viewResultSet",
+  OpenCodeInEditor = "openCodeInEditor",
 }
 
 interface RecInfo {
@@ -323,11 +324,24 @@ export class QueryResultPanel extends AltimateWebviewProvider {
     }
   }
 
+  private async handleOpenCodeInEditor(message: {
+    code: string;
+    name: string;
+  }) {
+    commands.executeCommand("dbtPowerUser.createSqlFile", {
+      code: message.code,
+      name: message.name,
+    });
+  }
+
   /** Primary interface for WebviewView inbound communication */
   private setupWebviewHooks() {
     this._panel!.webview.onDidReceiveMessage(
       async (message) => {
         switch (message.command) {
+          case InboundCommand.OpenCodeInEditor:
+            this.handleOpenCodeInEditor(message);
+            break;
           case InboundCommand.ViewResultSet:
             const queryHistoryData = message.queryHistory;
             this._queryTabData = {

diff --git a/webview_panels/src/assets/icons/index.tsx b/webview_panels/src/assets/icons/index.tsx
@@ -155,3 +155,7 @@ export const SearchIcon = (props: HTMLAttributes<HTMLElement>): JSX.Element => (
 export const FileCodeIcon = (
   props: HTMLAttributes<HTMLElement>,
 ): JSX.Element => <Icon icon="file-code" {...props} />;
+
+export const OpenNewIcon = (
+  props: HTMLAttributes<HTMLElement>,
+): JSX.Element => <Icon icon="link-external" {...props} />;
diff --git a/webview_panels/src/lib/altimate/altimate-components.d.ts b/webview_panels/src/lib/altimate/altimate-components.d.ts
@@ -1,6 +1,5 @@
-/// <reference types="react" />
-
 import { JSX as JSX_2 } from "react/jsx-runtime";
+import { ReactNode } from "react";
 
 export declare const ApiHelper: {
   get: <T>(
@@ -21,6 +20,7 @@ export declare const CodeBlock: ({
   fileName,
   theme,
   showLineNumbers,
+  titleActions,
 }: Props_4) => JSX.Element;
 
 export declare interface Conversation {
@@ -131,6 +131,7 @@ declare interface Props_4 {
   fileName?: string;
   showLineNumbers?: boolean;
   theme?: "vs" | "vsc-dark-plus" | "solarizedLight";
+  titleActions?: ReactNode;
 }
 
 declare interface User {