feat: support for general message in query validation plugin

pinterest · Feb 9, 2023 · 3308a0f · 3308a0f
1 parent efdbd7b
commit 3308a0f
Show file tree

Hide file tree

Showing 9 changed files with 89 additions and 29 deletions.
diff --git a/querybook/server/lib/query_analysis/validation/base_query_validator.py b/querybook/server/lib/query_analysis/validation/base_query_validator.py
@@ -7,6 +7,11 @@
 )
 
 
+class QueryValidationResultObjectType(Enum):
+    LINT = "lint"
+    GENERAL = "general"
+
+
 class QueryValidationSeverity(Enum):
     ERROR = "error"
     WARNING = "warning"
@@ -15,18 +20,21 @@ class QueryValidationSeverity(Enum):
 class QueryValidationResult(object):
     def __init__(
         self,
+        obj_type: QueryValidationResultObjectType,
         line: int,  # 0 based
         ch: int,  # location of the starting token
         severity: QueryValidationSeverity,
         message: str,
     ):
+        self.type = obj_type
         self.line = line
         self.ch = ch
         self.severity = severity
         self.message = message
 
     def to_dict(self):
         return {
+            "type": self.type.value,
             "line": self.line,
             "ch": self.ch,
             "severity": self.severity.value,

diff --git a/querybook/server/lib/query_analysis/validation/validators/presto_explain_validator.py b/querybook/server/lib/query_analysis/validation/validators/presto_explain_validator.py
@@ -4,6 +4,7 @@
     BaseQueryValidator,
     QueryValidationResult,
     QueryValidationSeverity,
+    QueryValidationResultObjectType,
 )
 from lib.utils.execute_query import ExecuteQuery
 from lib.query_executor.executors.presto import get_presto_error_dict
@@ -55,6 +56,7 @@ def _map_statement_error_to_query(
         )
 
         return QueryValidationResult(
+            QueryValidationResultObjectType.LINT,
             statement_error_line,
             statement_error_ch,
             QueryValidationSeverity.ERROR,

diff --git a/querybook/webapp/components/QueryEditor/QueryEditor.tsx b/querybook/webapp/components/QueryEditor/QueryEditor.tsx
@@ -10,6 +10,7 @@ import React, {
 } from 'react';
 import { Controlled as ReactCodeMirror } from 'react-codemirror2';
 import toast from 'react-hot-toast';
+import styled from 'styled-components';
 
 import { showTooltipFor } from 'components/CodeMirrorTooltip';
 import { ICodeMirrorTooltipProps } from 'components/CodeMirrorTooltip/CodeMirrorTooltip';
@@ -88,6 +89,13 @@ export interface IQueryEditorHandles {
     getEditorSelection: (editor?: CodeMirror.Editor) => IRange;
 }
 
+const StyledQueryValidationMsg = styled.span.attrs({
+    className: 'flex-row mr8',
+})`
+    color: ${(props) =>
+        props.type === 'info' ? 'var(--color-blue-dark)' : 'var(--color-false)'};
+`;
+
 export const QueryEditor: React.FC<
     IQueryEditorProps & {
         ref: React.Ref<IQueryEditorHandles>;
@@ -140,15 +148,23 @@ export const QueryEditor: React.FC<
 
         const [fullScreen, setFullScreen] = useState(false);
 
-        const { getLintAnnotations, lintSummary, isLinting } = useLint({
-            query: value,
-            editorRef,
-            metastoreId,
-            codeAnalysis,
-            getTableByName,
-            getLintErrors,
-            onLintCompletion,
-        });
+        const { getLintAnnotations, lintSummary, isLinting, queryAnnotations } =
+            useLint({
+                query: value,
+                editorRef,
+                metastoreId,
+                codeAnalysis,
+                getTableByName,
+                getLintErrors,
+                onLintCompletion,
+            });
+
+        const generalAnnotation: null | ILinterWarning = useMemo(() => {
+            const list = queryAnnotations.filter(
+                (obj: ILinterWarning) => obj.type === 'general'
+            );
+            return list.length > 0 ? list[0] : null;
+        }, [queryAnnotations]);
 
         const openTableModal = useCallback((tableId: number) => {
             navigateWithinEnv(`/table/${tableId}/`, {
@@ -603,28 +619,18 @@ export const QueryEditor: React.FC<
         );
         /* ---- end of <ReactCodeMirror /> properties ---- */
 
-        const renderLintButtons = () => {
-            if (value.length === 0) {
-                return null;
-            }
-
-            if (isLinting) {
-                return (
-                    <span className="flex-row mr8">
-                        <Icon name="Loading" className="mr4" size={16} />
-                        Linting
-                    </span>
-                );
-            }
-
+        const renderLintButton = () => {
             if (lintSummary.numErrors + lintSummary.numWarnings > 0) {
                 return (
                     <div
                         className="flex-row mr4"
                         title={`${formatNumber(
                             lintSummary.numErrors,
                             'error'
-                        )}, ${formatNumber(lintSummary.numWarnings, 'waring')}`}
+                        )}, ${formatNumber(
+                            lintSummary.numWarnings,
+                            'warning'
+                        )}`}
                     >
                         {lintSummary.numErrors > 0 && (
                             <span className="lint-num-errors flex-row mr4">
@@ -658,10 +664,40 @@ export const QueryEditor: React.FC<
             }
         };
 
+        const renderGeneralValidationMessage = () => {
+            if (!generalAnnotation) {
+                return null;
+            }
+            return (
+                <StyledQueryValidationMsg type={generalAnnotation.severity}>
+                    {generalAnnotation.message}
+                </StyledQueryValidationMsg>
+            );
+        };
+
+        const renderValidationMessages = () => {
+            if (value.length === 0) {
+                return null;
+            }
+            if (isLinting) {
+                return (
+                    <span className="flex-row mr8">
+                        <Icon name="Loading" className="mr4" size={16} />
+                        Linting
+                    </span>
+                );
+            }
+            return (
+                <>
+                    {renderGeneralValidationMessage()}
+                    {renderLintButton()}
+                </>
+            );
+        };
+
         const floatButtons = (
             <div className="query-editor-float-buttons-wrapper flex-row mt8 mr8">
-                {renderLintButtons()}
-
+                {renderValidationMessages()}
                 <IconButton
                     icon={fullScreen ? 'Minimize2' : 'Maximize2'}
                     onClick={toggleFullScreen}

diff --git a/querybook/webapp/const/queryExecution.ts b/querybook/webapp/const/queryExecution.ts
@@ -149,4 +149,5 @@ export interface IQueryValidationResult {
     ch: number;
     message: string;
     severity: string;
+    type: string;
 }
diff --git a/querybook/webapp/hooks/queryEditor/useLint.ts b/querybook/webapp/hooks/queryEditor/useLint.ts
@@ -166,7 +166,12 @@ export function useLint({
     );
     const lintAnnotationsRef = useRef<ILinterWarning[]>([]);
     const lintAnnotations = useMemo(
-        () => tableAnnotations.concat(queryAnnotations),
+        () =>
+            tableAnnotations.concat(
+                queryAnnotations.filter(
+                    (obj: ILinterWarning) => obj.type === 'lint'
+                )
+            ),
         [tableAnnotations, queryAnnotations]
     );
 
@@ -211,5 +216,6 @@ export function useLint({
         getLintAnnotations: getCodeMirrorLintAnnotations,
         isLinting: isLintingQuery || isLintingTable,
         lintSummary,
+        queryAnnotations,
     };
 }
diff --git a/querybook/webapp/lib/codemirror/codemirror-lint.ts b/querybook/webapp/lib/codemirror/codemirror-lint.ts
@@ -17,7 +17,7 @@ export function createSQLLinter(
             );
 
         return validationResults.map((validationError) => {
-            const { line, ch, severity, message } = validationError;
+            const { type, line, ch, severity, message } = validationError;
 
             const errorToken = cm.getTokenAt({
                 line,
@@ -37,6 +37,7 @@ export function createSQLLinter(
                     },
                     severity,
                     message,
+                    type,
                 } as ILinterWarning;
             } else {
                 return {
@@ -50,6 +51,7 @@ export function createSQLLinter(
                     },
                     severity,
                     message,
+                    type,
                 } as ILinterWarning;
             }
         });

diff --git a/querybook/webapp/lib/sql-helper/sql-context-free-linter.ts b/querybook/webapp/lib/sql-helper/sql-context-free-linter.ts
@@ -47,6 +47,7 @@ const contextFreeLinterWarningsByLanguage: Record<
                                 message:
                                     'Do a "count(*) and group by" works better in terms of performance.',
                                 severity: 'warning',
+                                type: 'lint',
                                 from: {
                                     line: token.line,
                                     ch: token.start,
@@ -88,6 +89,7 @@ const contextFreeLinterWarningsByLanguage: Record<
                         warnings.push({
                             message: 'Please add a LIMIT clause',
                             severity: 'warning',
+                            type: 'lint',
                             from: {
                                 line: firstToken.line,
                                 ch: firstToken.start,

diff --git a/querybook/webapp/lib/sql-helper/sql-context-sensitive-linter.ts b/querybook/webapp/lib/sql-helper/sql-context-sensitive-linter.ts
@@ -22,6 +22,7 @@ export function getContextSensitiveWarnings(
                 contextSensitiveWarnings.push({
                     message: `Table ${table.name} is newly created or does not exist`,
                     severity: 'warning',
+                    type: 'lint',
                     from: {
                         line: table.line,
                         ch: implicitSchema
@@ -55,6 +56,7 @@ export function getContextSensitiveWarnings(
 
                 contextSensitiveWarnings.push({
                     message: warningMessage,
+                    type: 'lint',
                     severity:
                         maxSeverity === DataTableWarningSeverity.ERROR
                             ? 'error'

diff --git a/querybook/webapp/lib/sql-helper/sql-lexer.ts b/querybook/webapp/lib/sql-helper/sql-lexer.ts
@@ -174,8 +174,9 @@ export interface IRange {
 }
 
 export interface ILinterWarning extends IRange {
+    type: 'lint' | 'general';
     message: string;
-    severity: 'error' | 'warning';
+    severity: 'error' | 'warning' | 'info';
 }
 
 export interface ILineage {