fix: revert to use searchConcise for mentioning a table (#1438)

* fix: revert to use searchConcise for mentioning a table * remove unused code * remove unused imports * comments
pinterest · Apr 12, 2024 · 4c74ac8 · 4c74ac8
1 parent cd11be9
commit 4c74ac8
Show file tree

Hide file tree

Showing 9 changed files with 83 additions and 74 deletions.
diff --git a/querybook/config/elasticsearch.yaml b/querybook/config/elasticsearch.yaml
@@ -212,7 +212,7 @@ tables:
                 full_name_ngram:
                     type: text
                     analyzer: edge_ngram_lowercase
-                completion_name: # to be deprecated
+                completion_name:
                     type: completion
                     analyzer: keyword
                     contexts:

diff --git a/querybook/server/datasources/search.py b/querybook/server/datasources/search.py
@@ -130,14 +130,20 @@ def vector_search_tables(
 
 
 @register("/suggest/<int:metastore_id>/tables/", methods=["GET"])
-def suggest_tables(metastore_id, keyword, limit=10):
+def suggest_tables(metastore_id, prefix, limit=10):
     api_assert(limit is None or limit <= 100, "Requesting too many tables")
     verify_metastore_permission(metastore_id)
 
-    query = construct_suggest_table_query(keyword, limit, metastore_id)
+    query = construct_suggest_table_query(prefix, limit, metastore_id)
     options = get_matching_suggestions(query, ES_CONFIG["tables"]["index_name"])
-    tables = [option.get("_source", {}).get("full_name") for option in options]
-    return tables
+    texts = [
+        "{}.{}".format(
+            option.get("_source", {}).get("schema", ""),
+            option.get("_source", {}).get("name", ""),
+        )
+        for option in options
+    ]
+    return texts
 
 
 # /search/ but it is actually suggest

diff --git a/querybook/server/lib/elasticsearch/search_utils.py b/querybook/server/lib/elasticsearch/search_utils.py
@@ -79,6 +79,8 @@ def get_matching_suggestions(query: Union[str, Dict], index_name: str):
         if result is None:
             result = {}
 
-    options = result.get("hits", {}).get("hits", [])
+    options = next(iter(result.get("suggest", {}).get("suggest", [])), {}).get(
+        "options", []
+    )
 
     return options
diff --git a/querybook/server/lib/elasticsearch/suggest_table.py b/querybook/server/lib/elasticsearch/suggest_table.py
@@ -1,16 +1,17 @@
 def construct_suggest_table_query(
-    keyword: str,
+    prefix: str,
     limit: int,
     metastore_id: int,
 ):
     return {
-        "from": 0,
-        "size": limit,
-        "query": {
-            "bool": {
-                "must": [{"match_phrase_prefix": {"full_name": {"query": keyword}}}],
-                "filter": {"match": {"metastore_id": metastore_id}},
+        "suggest": {
+            "suggest": {
+                "text": prefix,
+                "completion": {
+                    "field": "completion_name",
+                    "size": limit,
+                    "contexts": {"metastore_id": metastore_id},
+                },
             }
         },
-        "_source": ["id", "full_name"],
     }
diff --git a/querybook/server/logic/elasticsearch.py b/querybook/server/logic/elasticsearch.py
@@ -601,7 +601,6 @@ def get_completion_name():
         "name": table_name,
         "full_name": full_name,
         "full_name_ngram": full_name,
-        # To be deprecated: completion_name is not used anymore
         "completion_name": get_completion_name,
         "description": get_table_description,
         "created_at": lambda: DATETIME_TO_UTC(table.created_at),

diff --git a/querybook/webapp/components/AIAssistant/AICommandBar.tsx b/querybook/webapp/components/AIAssistant/AICommandBar.tsx
@@ -19,8 +19,6 @@ import { IQueryEngine } from 'const/queryEngine';
 import { CommandRunner, useCommand } from 'hooks/useCommand';
 import { useForwardedRef } from 'hooks/useForwardedRef';
 import { trackClick } from 'lib/analytics';
-import { TableToken } from 'lib/sql-helper/sql-lexer';
-import { analyzeCode } from 'lib/web-worker';
 import { Button } from 'ui/Button/Button';
 import { Message } from 'ui/Message/Message';
 import { IResizableTextareaHandles } from 'ui/ResizableTextArea/ResizableTextArea';
@@ -31,44 +29,28 @@ import './AICommandBar.scss';
 interface IQueryCellCommandBarProps {
     query: string;
     queryEngine: IQueryEngine;
-    engineId: number;
-    queryEngines: IQueryEngine[];
-    queryEngineById: Record<number, IQueryEngine>;
+    tablesInQuery: string[];
     onUpdateQuery: (query: string, run: boolean) => void;
     onUpdateEngineId: (engineId: number) => void;
     onFormatQuery: () => void;
     ref: React.Ref<IResizableTextareaHandles>;
 }
 
-const useTablesInQuery = (query: string, language: string) => {
-    const [tables, setTables] = useState<string[]>([]);
-
-    useEffect(() => {
-        if (!query) {
-            return;
-        }
-
-        analyzeCode(query, 'autocomplete', language).then((codeAnalysis) => {
-            const tableReferences: TableToken[] = [].concat.apply(
-                [],
-                Object.values(codeAnalysis?.lineage.references ?? {})
-            );
-            setTables(
-                tableReferences.map(({ schema, name }) => `${schema}.${name}`)
-            );
-        });
-    }, [query, language]);
-
-    return tables;
-};
-
 export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
-    ({ query = '', queryEngine, onUpdateQuery, onFormatQuery }, ref) => {
+    (
+        {
+            query = '',
+            queryEngine,
+            tablesInQuery = [],
+            onUpdateQuery,
+            onFormatQuery,
+        },
+        ref
+    ) => {
         const defaultCommand = QUERY_CELL_COMMANDS.find(
             (cmd) => cmd.name === (query ? 'edit' : 'generate')
         );
-        const tablesInQuery = useTablesInQuery(query, queryEngine.language);
-        const [tables, setTables] = useState<string[]>([]);
+        const [tables, setTables] = useState<string[]>(tablesInQuery);
         const [mentionedTables, setMentionedTables] = useState<string[]>([]);
         const commandInputRef = useForwardedRef<IResizableTextareaHandles>(ref);
         const [showPopupView, setShowPopupView] = useState(false);
@@ -81,13 +63,25 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
         );
         const [showConfirm, setShowConfirm] = useState(false);
 
-        const finalTablesToUse = useMemo(() => {
-            if (mentionedTables.length > 0) {
-                return mentionedTables;
-            } else {
-                return uniq([...tablesInQuery, ...tables]);
-            }
-        }, [tablesInQuery, mentionedTables, tables]);
+        /**
+         * There are three types of tables:
+         * 1. Tables used in the query.
+         * 2. Tables mentioned in the command.
+         * 3. Tables selected in the table selector.
+         *
+         * Regarding which tables to use:
+         * - If there are mentioned tables in the command, use those.
+         * - If there are no mentioned tables, use the tables in the query (type 1) merged with the selected tables (type 3).
+         * - However, selected tables will override the tables in the query. If a table is deselected, it will be removed from the list, including the tables in the query.
+         */
+        useEffect(() => {
+            setTables(uniq([...tables, ...tablesInQuery]));
+        }, [tablesInQuery]);
+
+        const finalTablesToUse = useMemo(
+            () => (mentionedTables.length ? mentionedTables : tables),
+            [mentionedTables, tables]
+        );
 
         const {
             runCommand,
@@ -154,7 +148,7 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
                         commandKwargs={commandKwargs}
                         metastoreId={queryEngine.metastore_id}
                         commandResult={commandResult}
-                        tables={finalTablesToUse}
+                        tables={tables}
                         hasMentionedTables={mentionedTables.length > 0}
                         originalQuery={query}
                         isStreaming={isRunning}

diff --git a/querybook/webapp/components/AIAssistant/AICommandInput.tsx b/querybook/webapp/components/AIAssistant/AICommandInput.tsx
@@ -207,20 +207,24 @@ export const AICommandInput: React.FC<AICommandInputProps> = forwardRef(
                     return;
                 }
 
-                SearchTableResource.suggest(metastoreId, keyword).then(
-                    ({ data }) => {
-                        const filteredTableNames = data.filter(
-                            (table) => !mentionedTables.includes(table)
-                        );
-                        const tableNameOptions = filteredTableNames.map(
-                            (table) => ({
-                                id: table,
-                                display: table,
-                            })
-                        );
-                        callback(tableNameOptions);
-                    }
-                );
+                SearchTableResource.searchConcise({
+                    metastore_id: metastoreId,
+                    keywords: keyword,
+                }).then(({ data }) => {
+                    const filteredTableNames = data.results.filter(
+                        (result) =>
+                            !mentionedTables.includes(
+                                `${result.schema}.${result.name}`
+                            )
+                    );
+                    const tableNameOptions = filteredTableNames.map(
+                        ({ schema, name }) => ({
+                            id: `${schema}.${name}`,
+                            display: `${schema}.${name}`,
+                        })
+                    );
+                    callback(tableNameOptions);
+                });
             },
             [metastoreId, mentionedTables]
         );

diff --git a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx
@@ -124,6 +124,7 @@ interface IState {
     showRenderedTemplateModal: boolean;
     showUDFModal: boolean;
     hasLintError: boolean;
+    tableNamesInQuery: string[];
     samplingTables: ISamplingTables;
 
     transpilerConfig?: {
@@ -151,6 +152,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
             showRenderedTemplateModal: false,
             showUDFModal: false,
             hasLintError: false,
+            tableNamesInQuery: [],
             samplingTables: {},
         };
     }
@@ -692,7 +694,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
     }
 
     @bind
-    public async handleTablesChange(tablesByName: Record<string, IDataTable>) {
+    public handleTablesChange(tablesByName: Record<string, IDataTable>) {
         const samplingTables = {};
         Object.keys(tablesByName).forEach((tableName) => {
             const table = tablesByName[tableName];
@@ -702,7 +704,10 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                 };
             }
         });
-        this.setState({ samplingTables });
+        this.setState({
+            samplingTables,
+            tableNamesInQuery: Object.keys(tablesByName),
+        });
     }
 
     public componentDidMount() {
@@ -805,10 +810,8 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                     <AICommandBar
                         query={query}
                         queryEngine={queryEngineById[this.engineId]}
-                        engineId={this.engineId}
+                        tablesInQuery={this.state.tableNamesInQuery}
                         onUpdateQuery={this.handleChange}
-                        queryEngineById={queryEngineById}
-                        queryEngines={this.props.queryEngines}
                         onUpdateEngineId={this.handleMetaChange.bind(
                             this,
                             'engine'

diff --git a/querybook/webapp/resource/search.ts b/querybook/webapp/resource/search.ts
@@ -33,9 +33,9 @@ export const SearchTableResource = {
             count: number;
         }>('/search/tables/vector/', { ...params }),
 
-    suggest: (metastoreId: number, keyword: string) =>
+    suggest: (metastoreId: number, prefix: string) =>
         ds.fetch<string[]>(`/suggest/${metastoreId}/tables/`, {
-            keyword,
+            prefix,
         }),
 };