docs: add docs for sql complete and fix re-render issue

docs_website/docs/integrations/add_ai_assistant.mdx

@@ -12,7 +12,12 @@ The AI assistant plugin is powered by LLM(Large Language Model), like ChatGPT fr
 
 ## AI Assistant Plugin
 
-The AI Assistant plugin will allow users to do title generation, text to sql and query auto fix.
+The AI Assistant plugin will allow users to do
+
+-   title generation
+-   text to sql
+-   query auto fix
+-   sql completion
 
 Please follow below steps to enable AI assistant plugin:
 

docs_website/docs/user_guide/ai_assistant.mdx

@@ -37,6 +37,13 @@ If your query failed, you will see ‘Auto fix’ button on the right corner of
 
 ![](/img/user_guide/sql_fix.gif)
 
+## SQL Completion
+
+SQL completion offers SQL code suggestions similar to GitHub Copilot while you write queries.
+This feature is disabled by default. You can enable it in the Editor tab of user settings.
+
+![](/img/user_guide/sql_complete.png)
+
 ## Search Table by Natural Language
 
 If [vector store](../integrations/add_ai_assistant.mdx#vector-store) of the AI assistant plugin is also enabled, you'll be able to search the tables by natual language as well as keyword based search.

querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx

@@ -1,4 +1,4 @@
-import React, { useContext, useEffect, useMemo } from 'react';
+import React, { useContext, useEffect, useMemo, useRef } from 'react';
 import { useDispatch } from 'react-redux';
 
 import {
@@ -7,13 +7,9 @@ import {
     QueryEditor,
 } from 'components/QueryEditor/QueryEditor';
 import { IQueryEngine } from 'const/queryEngine';
-import { SearchAndReplaceContext } from 'context/searchAndReplace';
 import { useUserQueryEditorConfig } from 'hooks/redux/useUserQueryEditorConfig';
 import { useForwardedRef } from 'hooks/useForwardedRef';
-import {
-    fetchDataTableByNameIfNeeded,
-    fetchFunctionDocumentationIfNeeded,
-} from 'redux/dataSources/action';
+import { fetchDataTableByNameIfNeeded } from 'redux/dataSources/action';
 
 export const BoundQueryEditor = React.forwardRef<
     IQueryEditorHandles,
@@ -34,7 +30,6 @@ export const BoundQueryEditor = React.forwardRef<
     }
 >(({ options: propOptions, keyMap, engine, cellId, ...otherProps }, ref) => {
     const dispatch = useDispatch();
-    const searchContext = useContext(SearchAndReplaceContext);
     const editorRef = useForwardedRef<IQueryEditorHandles>(ref);
 
     // Code Editor related Props
@@ -81,7 +76,6 @@ export const BoundQueryEditor = React.forwardRef<
             getTableByName={fetchDataTable}
             metastoreId={engine?.metastore_id}
             language={engine?.language}
-            searchContext={searchContext}
             cellId={cellId}
             engineId={engine?.id}
             sqlCompleteEnabled={sqlCompleteEnabled}

querybook/webapp/components/QueryEditor/QueryEditor.scss

@@ -35,19 +35,14 @@
         height: 100%;
         border-radius: var(--border-radius-sm);
         overflow: hidden;
+        background-color: var(--bg-query-editor);
 
         &.cm-theme-light {
             .cm-editor {
                 .cm-scroller {
-                    background-color: var(--bg-query-editor-gutter);
-
                     .cm-gutters {
                         background-color: var(--bg-query-editor-gutter);
                     }
-
-                    .cm-content {
-                        background-color: var(--bg-query-editor);
-                    }
                 }
             }
         }

querybook/webapp/components/QueryEditor/QueryEditor.tsx

@@ -14,7 +14,6 @@ import toast from 'react-hot-toast';
 import { TDataDocMetaVariables } from 'const/datadoc';
 import KeyMap from 'const/keyMap';
 import { IDataTable } from 'const/metastore';
-import { ISearchAndReplaceContextType } from 'context/searchAndReplace';
 import { useAutoCompleteExtension } from 'hooks/queryEditor/extensions/useAutoCompleteExtension';
 import { useEventsExtension } from 'hooks/queryEditor/extensions/useEventsExtension';
 import { useHoverTooltipExtension } from 'hooks/queryEditor/extensions/useHoverTooltipExtension';
@@ -55,7 +54,6 @@ export interface IQueryEditorProps {
     engineId: number;
     templatedVariables?: TDataDocMetaVariables;
     cellId?: number;
-    searchContext?: ISearchAndReplaceContextType;
 
     fontSize?: string;
     height?: 'auto' | 'full' | 'fixed';
@@ -109,7 +107,6 @@ export const QueryEditor: React.FC<
             engineId,
             cellId,
             templatedVariables = [],
-            searchContext,
 
             hasQueryLint,
             height = 'auto',
@@ -270,13 +267,12 @@ export const QueryEditor: React.FC<
 
         const searchExtension = useSearchExtension({
             editorView: editorRef.current?.view,
-            searchContext,
             cellId,
         });
 
         const eventsExtension = useEventsExtension({
-            onFocus: (evt) => onFocus?.(),
-            onBlur: (evt) => onBlur?.(),
+            onFocus,
+            onBlur,
         });
 
         const statusBarExtension = useStatusBarExtension({
@@ -313,9 +309,8 @@ export const QueryEditor: React.FC<
             return true;
         }, []);
 
-        const keyMapExtention = useKeyMapExtension({
-            keyMap,
-            keyBindings: [
+        const keyBindings = useMemo(
+            () => [
                 { key: 'Tab', run: acceptCompletion },
                 {
                     key: KeyMap.queryEditor.autocomplete.key,
@@ -328,18 +323,16 @@ export const QueryEditor: React.FC<
                         return true;
                     },
                 },
-                {
-                    key: 'Cmd-F',
-                    run: () => {
-                        searchContext?.showSearchAndReplace();
-                        return true;
-                    },
-                },
                 {
                     key: KeyMap.queryEditor.openTable.key,
                     run: openTableModalCommand,
                 },
             ],
+            [formatQuery, openTableModalCommand]
+        );
+        const keyMapExtention = useKeyMapExtension({
+            keyMap,
+            keyBindings,
         });
 
         const optionsExtension = useOptionsExtension({
@@ -367,7 +360,6 @@ export const QueryEditor: React.FC<
         const extensions = useMemo(
             () => [
                 mixedSQL(),
-
                 keyMapExtention,
                 statusBarExtension,
                 eventsExtension,
@@ -395,10 +387,11 @@ export const QueryEditor: React.FC<
 
         const basicSetup = useMemo(
             () => ({
-                drawSelection: false,
+                drawSelection: true,
                 highlightSelectionMatches: true,
                 searchKeymap: false,
                 foldGutter: false,
+                allowMultipleSelections: true,
             }),
             []
         );

querybook/webapp/components/SearchAndReplace/SearchAndReplace.tsx

@@ -2,6 +2,7 @@ import React, {
     useCallback,
     useEffect,
     useImperativeHandle,
+    useMemo,
     useRef,
     useState,
 } from 'react';
@@ -68,16 +69,16 @@ export function useSearchAndReplace({
 
     const reset = useCallback(() => setSearchState(initialSearchState), []);
     const performSearch = useCallback(
-        // Active search happens when user is using the search input
-        // Passive search is when the search content is changed
         (isActiveSearch: boolean = false) => {
+            if (!showing) {
+                return;
+            }
+
             setSearchState((oldSearchState) => {
-                const searchResults = showing
-                    ? getSearchResults(
-                          oldSearchState.searchString,
-                          oldSearchState.searchOptions
-                      )
-                    : [];
+                const searchResults = getSearchResults(
+                    oldSearchState.searchString,
+                    oldSearchState.searchOptions
+                );
 
                 if (isActiveSearch) {
                     jumpToResult(
@@ -195,15 +196,26 @@ export function useSearchAndReplace({
         setShowing(false);
     }, []);
 
-    return {
-        searchAndReplaceContext: {
+    const searchAndReplaceContext = useMemo(
+        () => ({
             searchState,
             focusSearchBar,
 
             showSearchAndReplace,
             hideSearchAndReplace,
             showing,
-        },
+        }),
+        [
+            searchState,
+            focusSearchBar,
+            showSearchAndReplace,
+            hideSearchAndReplace,
+            showing,
+        ]
+    );
+
+    return {
+        searchAndReplaceContext,
         searchAndReplaceProps: {
             onSearchStringChange,
             onReplaceStringChange,

querybook/webapp/hooks/queryEditor/extensions/useEventsExtension.ts

@@ -1,12 +1,18 @@
 import * as events from '@uiw/codemirror-extensions-events';
 import { useMemo } from 'react';
 
-export const useEventsExtension = ({ onFocus, onBlur }) => {
+export const useEventsExtension = ({
+    onFocus,
+    onBlur,
+}: {
+    onFocus?: () => void;
+    onBlur?: () => void;
+}) => {
     const extension = useMemo(
         () =>
             events.content({
-                focus: onFocus,
-                blur: onBlur,
+                focus: (evt) => onFocus?.(),
+                blur: (evt) => onBlur?.(),
             }),
         [onFocus, onBlur]
     );

querybook/webapp/hooks/queryEditor/extensions/useSearchExtension.ts

@@ -5,20 +5,24 @@ import {
     SearchQuery,
     setSearchQuery,
 } from '@codemirror/search';
-import { EditorSelection, EditorView } from '@uiw/react-codemirror';
-import { useEffect, useMemo } from 'react';
+import {
+    EditorSelection,
+    EditorView,
+    keymap,
+    Prec,
+} from '@uiw/react-codemirror';
+import { useContext, useEffect, useMemo } from 'react';
 
-import { ISearchAndReplaceContextType } from 'context/searchAndReplace';
+import { SearchAndReplaceContext } from 'context/searchAndReplace';
 
 export const useSearchExtension = ({
     editorView,
     cellId,
-    searchContext,
 }: {
     editorView: EditorView;
     cellId: number;
-    searchContext?: ISearchAndReplaceContextType;
 }) => {
+    const searchContext = useContext(SearchAndReplaceContext);
     useEffect(() => {
         if (editorView && searchContext) {
             if (searchContext.showing) {
@@ -30,7 +34,7 @@ export const useSearchExtension = ({
                 closeSearchPanel(editorView);
             }
         }
-    }, [editorView, searchContext]);
+    }, [editorView, searchContext?.showing]);
 
     const shouldHighlight = useMemo(
         () =>
@@ -80,7 +84,23 @@ export const useSearchExtension = ({
         }
     }, [shouldHighlight, searchContext]);
 
-    const extension = useMemo(() => search(), []);
+    const extension = useMemo(
+        () => [
+            search(),
+            Prec.highest(
+                keymap.of([
+                    {
+                        key: 'Cmd-f',
+                        run: () => {
+                            searchContext?.showSearchAndReplace();
+                            return true;
+                        },
+                    },
+                ])
+            ),
+        ],
+        []
+    );
 
     return extension;
 };

querybook/webapp/hooks/queryEditor/useLint.tsx

@@ -8,7 +8,7 @@ import { TDataDocMetaVariables } from 'const/datadoc';
 import { IQueryValidationResult } from 'const/queryExecution';
 import { useDebounce } from 'hooks/useDebounce';
 import useDeepCompareEffect from 'hooks/useDeepCompareEffect';
-import { posToOffset } from 'lib/codemirror/utils';
+import { getTokenAtOffset, posToOffset } from 'lib/codemirror/utils';
 import { getContextSensitiveWarnings } from 'lib/sql-helper/sql-context-sensitive-linter';
 import { ILinterWarning, TableToken } from 'lib/sql-helper/sql-lexer';
 import { TemplatedQueryResource } from 'resource/queryExecution';
@@ -75,10 +75,7 @@ const queryValidationErrorsToDiagnostics = (
             suggestion,
         } = validationError;
 
-        const startPos = posToOffset(editorView, {
-            line,
-            ch: ch + 1,
-        });
+        const startPos = posToOffset(editorView, { line, ch });
         const endPos =
             endLine !== null && endCh !== null
                 ? posToOffset(editorView, {