feat: add column hover tooltip and column value auto completion

pinterest · Dec 6, 2024 · 8fb7da6 · 8fb7da6
1 parent dff25db
commit 8fb7da6
Show file tree

Hide file tree

Showing 16 changed files with 625 additions and 430 deletions.
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.37.1",
+    "version": "3.38.0",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {

diff --git a/querybook/server/models/metastore.py b/querybook/server/models/metastore.py
@@ -316,6 +316,7 @@ def to_dict(self, include_table=False):
             "comment": self.comment,
             "description": self.description,
             "table_id": self.table_id,
+            "stats": self.statistics,
         }
 
         if include_table:

diff --git a/querybook/webapp/components/CodeMirrorTooltip/FunctionDocumentationTooltip.tsx b/querybook/webapp/components/CodeMirrorTooltip/FunctionDocumentationTooltip.tsx
@@ -27,7 +27,7 @@ export const FunctionDocumentationTooltip: React.FunctionComponent<IProps> = ({
             return (
                 <div key={index}>
                     <div className="rich-text-content">
-                        <div className="table-tooltip-header">{signature}</div>
+                        <div className="tooltip-header">{signature}</div>
 
                         <div className="tooltip-title">Returns</div>
                         <div className="tooltip-content">{returnType}</div>

diff --git a/querybook/webapp/components/CodeMirrorTooltip/TableColumnTooltip.tsx b/querybook/webapp/components/CodeMirrorTooltip/TableColumnTooltip.tsx
@@ -0,0 +1,88 @@
+import { ContentState } from 'draft-js';
+import React from 'react';
+
+import { TableTag } from 'components/DataTableTags/DataTableTags';
+import { IDataColumn } from 'const/metastore';
+import { useResource } from 'hooks/useResource';
+import { TableColumnResource } from 'resource/table';
+import { Tag, TagGroup } from 'ui/Tag/Tag';
+
+interface IProps {
+    column: IDataColumn;
+}
+
+export const TableColumnTooltip: React.FunctionComponent<IProps> = ({
+    column,
+}) => {
+    const { data: tags } = useResource(
+        React.useCallback(
+            () => TableColumnResource.getTags(column.id),
+            [column.id]
+        )
+    );
+
+    const tagsDOM = (tags || []).map((tag) => (
+        <TableTag tag={tag} readonly={true} key={tag.id} mini={true} />
+    ));
+
+    const description =
+        typeof column.description === 'string'
+            ? column.description
+            : (column.description as ContentState).getPlainText();
+
+    const statsDOM = (column.stats || []).map((stat, i) => {
+        const formattedValue = Array.isArray(stat.value)
+            ? stat.value.join(', ')
+            : stat.value;
+        return (
+            <TagGroup key={stat.key}>
+                <Tag mini>{stat.key}</Tag>
+                <Tag highlighted mini>
+                    {formattedValue}
+                </Tag>
+            </TagGroup>
+        );
+    });
+
+    const contentDOM = (
+        <>
+            <div className="tooltip-header flex-row">
+                <div>{column.name}</div>
+            </div>
+            {column.type && (
+                <div className="mt4">
+                    <div className="tooltip-title">Type</div>
+                    <div className="tooltip-content">{column.type}</div>
+                </div>
+            )}
+            {tagsDOM.length > 0 && (
+                <div className="mt4">
+                    <div className="tooltip-title">Tags</div>
+                    <div className="tooltip-content">
+                        <div className="DataTableTags flex-row">{tagsDOM}</div>
+                    </div>
+                </div>
+            )}
+            {column.comment && (
+                <div className="mt4">
+                    <div className="tooltip-title">Definition</div>
+                    <div className="tooltip-content">{column.comment}</div>
+                </div>
+            )}
+            {description && (
+                <div className="mt4">
+                    <div className="tooltip-title">Description</div>
+                    <div className="tooltip-content">{description}</div>
+                </div>
+            )}
+            {!!column?.stats?.length && (
+                <div className="mt4">
+                    <div className="tooltip-title">Stats</div>
+                    <div className="tooltip-content">{statsDOM}</div>
+                </div>
+            )}
+        </>
+    );
+
+    return <div className="rich-text-content">{contentDOM}</div>;
+};
diff --git a/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx b/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx
@@ -95,7 +95,7 @@ export const TableTooltip: React.FunctionComponent<IProps> = ({
 
     const contentDOM = (
         <>
-            <div className="table-tooltip-header flex-row">
+            <div className="tooltip-header flex-row">
                 <div>{tableName}</div>
                 <div className="flex-row mt4 ml4">
                     {pinToSidebarButton}
@@ -110,7 +110,7 @@ export const TableTooltip: React.FunctionComponent<IProps> = ({
         </>
     );
 
-    return <div className="rich-text-content ">{contentDOM}</div>;
+    return <div className="rich-text-content">{contentDOM}</div>;
 };
 
 export const TableTooltipByName: React.FunctionComponent<{
@@ -144,7 +144,9 @@ export const TableTooltipByName: React.FunctionComponent<{
                         metastoreId
                     )
                 );
-                setTableId(table.id);
+                if (table) {
+                    setTableId(table.id);
+                }
             } catch (error) {
                 console.error('Error fetching table:', error);
             }

diff --git a/querybook/webapp/components/QueryEditor/QueryEditor.scss b/querybook/webapp/components/QueryEditor/QueryEditor.scss
@@ -74,9 +74,9 @@
                     word-break: break-word;
                     font-size: var(--text-size);
 
-                    .table-tooltip-header {
+                    .tooltip-header {
                         font-size: var(--text-size);
-                        color: var(--text-dark);
+                        color: var(--color-accent-dark);
                         font-weight: var(--bold-font);
                         justify-content: space-between;
                         align-items: start;

diff --git a/querybook/webapp/components/QueryEditor/QueryEditor.tsx b/querybook/webapp/components/QueryEditor/QueryEditor.tsx
@@ -17,7 +17,10 @@ import toast from 'react-hot-toast';
 import { TDataDocMetaVariables } from 'const/datadoc';
 import KeyMap from 'const/keyMap';
 import { IDataTable } from 'const/metastore';
-import { useAutoCompleteExtension } from 'hooks/queryEditor/extensions/useAutoCompleteExtension';
+import {
+    AutoCompleteType,
+    useAutoCompleteExtension,
+} from 'hooks/queryEditor/extensions/useAutoCompleteExtension';
 import { useEventsExtension } from 'hooks/queryEditor/extensions/useEventsExtension';
 import { useHoverTooltipExtension } from 'hooks/queryEditor/extensions/useHoverTooltipExtension';
 import { useKeyMapExtension } from 'hooks/queryEditor/extensions/useKeyMapExtension';
@@ -26,13 +29,12 @@ import { useOptionsExtension } from 'hooks/queryEditor/extensions/useOptionsExte
 import { useSearchExtension } from 'hooks/queryEditor/extensions/useSearchExtension';
 import { useSqlCompleteExtension } from 'hooks/queryEditor/extensions/useSqlCompleteExtension';
 import { useStatusBarExtension } from 'hooks/queryEditor/extensions/useStatusBarExtension';
-import { useAutoComplete } from 'hooks/queryEditor/useAutoComplete';
 import { useCodeAnalysis } from 'hooks/queryEditor/useCodeAnalysis';
 import { useLint } from 'hooks/queryEditor/useLint';
+import { useSqlParser } from 'hooks/queryEditor/useSqlParser';
 import useDeepCompareEffect from 'hooks/useDeepCompareEffect';
 import { CodeMirrorKeyMap } from 'lib/codemirror';
 import { mixedSQL } from 'lib/codemirror/codemirror-mixed';
-import { AutoCompleteType } from 'lib/sql-helper/sql-autocompleter';
 import { format, ISQLFormatOptions } from 'lib/sql-helper/sql-formatter';
 import { TableToken } from 'lib/sql-helper/sql-lexer';
 import { navigateWithinEnv } from 'lib/utils/query-string';
@@ -204,12 +206,7 @@ export const QueryEditor: React.FC<
             language,
             query: value,
         });
-        const autoCompleterRef = useAutoComplete(
-            metastoreId,
-            autoCompleteType,
-            language,
-            codeAnalysis
-        );
+        const sqlParserRef = useSqlParser(metastoreId, language, codeAnalysis);
 
         const tableReferences: TableToken[] = useMemo(
             () =>
@@ -287,7 +284,8 @@ export const QueryEditor: React.FC<
 
         const autoCompleteExtension = useAutoCompleteExtension({
             view: editorRef.current?.view,
-            autoCompleterRef,
+            sqlParserRef,
+            type: autoCompleteType,
         });
 
         const lintExtension = useLintExtension({
@@ -296,7 +294,7 @@ export const QueryEditor: React.FC<
 
         const { extension: hoverTooltipExtension, getTableAtCursor } =
             useHoverTooltipExtension({
-                codeAnalysisRef,
+                sqlParserRef,
                 metastoreId,
                 language,
             });

diff --git a/querybook/webapp/const/metastore.ts b/querybook/webapp/const/metastore.ts
@@ -127,9 +127,9 @@ export interface IDataColumn {
     name: string;
     table_id: number;
     type: string;
+    stats?: ITableColumnStats[];
 }
 export interface IDetailedDataColumn extends IDataColumn {
-    stats?: ITableColumnStats[];
     tags?: ITag[];
     data_element_association?: IDataElementAssociation;
 }

diff --git a/querybook/webapp/hooks/queryEditor/extensions/useAutoCompleteExtension.ts b/querybook/webapp/hooks/queryEditor/extensions/useAutoCompleteExtension.ts
@@ -1,41 +1,131 @@
 import {
     autocompletion,
+    Completion,
     CompletionContext,
     startCompletion,
 } from '@codemirror/autocomplete';
 import { EditorView } from '@uiw/react-codemirror';
 import { useCallback, useEffect, useMemo, useState } from 'react';
 
-import { SqlAutoCompleter } from 'lib/sql-helper/sql-autocompleter';
+import { SqlParser } from 'lib/sql-helper/sql-parser';
+
+export type AutoCompleteType = 'none' | 'schema' | 'all';
 
 // STATIC
 const RESULT_MAX_LENGTH = 10;
 
 export const useAutoCompleteExtension = ({
     view,
-    autoCompleterRef,
+    sqlParserRef,
+    type = 'all',
 }: {
     view: EditorView;
-    autoCompleterRef: React.MutableRefObject<SqlAutoCompleter>;
+    sqlParserRef: React.MutableRefObject<SqlParser>;
+    type: AutoCompleteType;
 }) => {
     const [typing, setTyping] = useState(false);
 
-    const getCompletions = useCallback(async (context: CompletionContext) => {
-        // Get the token before the cursor, token could be schema.table.column
-        const token = context.matchBefore(/(\w+\.){0,2}(\w+)?/);
-        // is no word before the cursor, don't open completions.
-        if (!token || !token.text) return null;
+    const getColumnValueCompletions = useCallback(
+        (cursor, token) => {
+            console.log('getColumnValueCompletions', cursor, token);
+            const [textBeforeEqual, textAfterEqual] = token.text
+                .split('=')
+                .map((s) => s.trim());
+            const columnValues = sqlParserRef.current.getColumnValues(
+                cursor,
+                textBeforeEqual
+            );
+
+            const hasQuote = textAfterEqual.startsWith("'");
+            return {
+                from: token.to - textAfterEqual.length + (hasQuote ? 1 : 0),
+                options: columnValues.map((v) => ({
+                    label: `${v}`,
+                    apply:
+                        typeof v === 'number' || hasQuote ? `${v}` : `'${v}'`,
+                })),
+            };
+        },
+        [sqlParserRef]
+    );
+
+    const getCompletions = useCallback(
+        async (context: CompletionContext) => {
+            if (type === 'none') {
+                return null;
+            }
+
+            // Get the token before the cursor, token could be in below foramts
+            //  - column: schema.table.column, table.column, column
+            //  - table: schema.table, table
+            //  - keyword: any keyword
+            //  - column value: column = value, (value may be quoted)
+            const token = context.matchBefore(
+                /(\w+\.){0,2}\w*|(\w+.)?\s*=\s*'?\w*/
+            );
+
+            // no token before the cursor, don't open completions.
+            if (!token?.text) return null;
+
+            // Get the cursor position in codemirror v5 format
+            const cursorPos = context.pos;
+            const line = context.state.doc.lineAt(cursorPos);
+            const cursor = { line: line.number - 1, ch: cursorPos - line.from };
 
-        const cursorPos = context.pos;
-        const line = context.state.doc.lineAt(cursorPos);
-        const cursor = { line: line.number - 1, ch: cursorPos - line.from };
+            const tokenText = token.text.toLowerCase();
+            const sqlParserContext =
+                sqlParserRef.current.getContextAtPos(cursor);
 
-        const completions = await autoCompleterRef.current.getCompletions(
-            cursor,
-            token
-        );
-        return completions;
-    }, []);
+            // handle the case where the token is a column and the user is trying to type a value in a where clause
+            if (sqlParserContext === 'column' && tokenText.includes('=')) {
+                return getColumnValueCompletions(cursor, token);
+            }
+
+            const options: Completion[] = [];
+            if (sqlParserContext === 'column') {
+                const columns = sqlParserRef.current.getColumnMatches(
+                    cursor,
+                    token.text
+                );
+                options.push(
+                    ...columns.map((column) => ({
+                        label: column.name,
+                        detail: 'column',
+                    }))
+                );
+            } else if (sqlParserContext === 'table') {
+                const tableNames =
+                    await sqlParserRef.current.getTableNameMatches(tokenText);
+                options.push(
+                    ...tableNames.map((tableName) => ({
+                        label: tableName,
+                        detail: 'table',
+                    }))
+                );
+            }
+
+            // keyword may appear in all contexts
+            const keywordMatches =
+                type === 'all'
+                    ? sqlParserRef.current.getKeyWordMatches(tokenText)
+                    : [];
+
+            options.push(
+                ...keywordMatches.map((keyword) => ({
+                    label: keyword,
+                    detail: 'keyword',
+                }))
+            );
+
+            let from = token.from;
+            if (sqlParserContext === 'column') {
+                from += token.text.lastIndexOf('.') + 1;
+            }
+
+            return { from, options };
+        },
+        [sqlParserRef, type, getColumnValueCompletions]
+    );
 
     const triggerCompletionOnType = () => {
         return EditorView.updateListener.of((update) => {
@@ -51,11 +141,11 @@ export const useAutoCompleteExtension = ({
     };
 
     useEffect(() => {
-        if (autoCompleterRef.current.codeAnalysis && typing && view) {
+        if (sqlParserRef.current.codeAnalysis && typing && view) {
             startCompletion(view);
             setTyping(false);
         }
-    }, [autoCompleterRef.current.codeAnalysis, view]);
+    }, [sqlParserRef.current.codeAnalysis, view]);
 
     const extension = useMemo(
         () => [