From cae1361c32a8a506cc5bb228d614b26fb64d9f55 Mon Sep 17 00:00:00 2001
From: mcbrewster <37322608+mcbrewster@users.noreply.github.com>
Date: Sat, 17 Aug 2019 09:34:36 -0700
Subject: [PATCH] Web-console: gate auto complete on current table and schema
 (#8322)

* gate auto complete on current table and schema

* reset defaults

* add static functions

* move completions to state
---
 .../__snapshots__/query-view.spec.tsx.snap    |   1 +
 .../query-view/query-input/query-input.tsx    | 124 +++++++++++-------
 .../src/views/query-view/query-view.tsx       |  17 +++
 3 files changed, 93 insertions(+), 49 deletions(-)

diff --git a/web-console/src/views/query-view/__snapshots__/query-view.spec.tsx.snap b/web-console/src/views/query-view/__snapshots__/query-view.spec.tsx.snap
index 22338e01b92..1090f255c82 100644
--- a/web-console/src/views/query-view/__snapshots__/query-view.spec.tsx.snap
+++ b/web-console/src/views/query-view/__snapshots__/query-view.spec.tsx.snap
@@ -32,6 +32,7 @@ exports[`sql view matches snapshot 1`] = `
       className="control-pane"
     >
       <QueryInput
+        currentSchema="druid"
         onQueryStringChange={[Function]}
         queryString="test"
         runeMode={false}
diff --git a/web-console/src/views/query-view/query-input/query-input.tsx b/web-console/src/views/query-view/query-input/query-input.tsx
index 4e761b1d37c..a8806db2230 100644
--- a/web-console/src/views/query-view/query-input/query-input.tsx
+++ b/web-console/src/views/query-view/query-input/query-input.tsx
@@ -25,7 +25,6 @@ import AceEditor from 'react-ace';
 import { SQL_DATE_TYPES, SQL_FUNCTIONS, SyntaxDescription } from '../../../../lib/sql-function-doc';
 import { uniq } from '../../../utils';
 import { ColumnMetadata } from '../../../utils/column-metadata';
-import { ColumnTreeProps, ColumnTreeState } from '../column-tree/column-tree';
 
 import { SQL_CONSTANTS, SQL_DYNAMICS, SQL_EXPRESSION_PARTS, SQL_KEYWORDS } from './keywords';
 
@@ -38,59 +37,22 @@ export interface QueryInputProps {
   onQueryStringChange: (newQueryString: string) => void;
   runeMode: boolean;
   columnMetadata?: ColumnMetadata[];
+  currentSchema?: string;
+  currentTable?: string;
 }
 
 export interface QueryInputState {
   // For reasons (https://github.com/securingsincity/react-ace/issues/415) react ace editor needs an explicit height
   // Since this component will grown and shrink dynamically we will measure its height and then set it.
   editorHeight: number;
+  completions: any[];
   prevColumnMetadata?: ColumnMetadata[];
+  prevCurrentTable?: string;
+  prevCurrentSchema?: string;
 }
 
 export class QueryInput extends React.PureComponent<QueryInputProps, QueryInputState> {
-  static getDerivedStateFromProps(props: ColumnTreeProps, state: ColumnTreeState) {
-    const { columnMetadata } = props;
-
-    if (columnMetadata && columnMetadata !== state.prevColumnMetadata) {
-      const completions = ([] as any[]).concat(
-        uniq(columnMetadata.map(d => d.TABLE_SCHEMA)).map(v => ({
-          value: v,
-          score: 10,
-          meta: 'schema',
-        })),
-        uniq(columnMetadata.map(d => d.TABLE_NAME)).map(v => ({
-          value: v,
-          score: 49,
-          meta: 'datasource',
-        })),
-        uniq(columnMetadata.map(d => d.COLUMN_NAME)).map(v => ({
-          value: v,
-          score: 50,
-          meta: 'column',
-        })),
-      );
-
-      langTools.addCompleter({
-        getCompletions: (_editor: any, _session: any, _pos: any, _prefix: any, callback: any) => {
-          callback(null, completions);
-        },
-      });
-
-      return {
-        prevColumnMetadata: columnMetadata,
-      };
-    }
-    return null;
-  }
-
-  constructor(props: QueryInputProps, context: any) {
-    super(props, context);
-    this.state = {
-      editorHeight: 200,
-    };
-  }
-
-  private replaceDefaultAutoCompleter = () => {
+  static replaceDefaultAutoCompleter(): void {
     if (!langTools) return;
 
     const keywordList = ([] as any[]).concat(
@@ -112,9 +74,9 @@ export class QueryInput extends React.PureComponent<QueryInputProps, QueryInputS
       langTools.textCompleter,
       keywordCompleter,
     ]);
-  };
+  }
 
-  private addFunctionAutoCompleter = (): void => {
+  static addFunctionAutoCompleter(): void {
     if (!langTools) return;
 
     const functionList: any[] = SQL_FUNCTIONS.map((entry: SyntaxDescription) => {
@@ -161,11 +123,75 @@ export class QueryInput extends React.PureComponent<QueryInputProps, QueryInputS
         }
       },
     });
-  };
+  }
+
+  static getDerivedStateFromProps(props: QueryInputProps, state: QueryInputState) {
+    const { columnMetadata, currentSchema, currentTable } = props;
+
+    if (
+      columnMetadata &&
+      (columnMetadata !== state.prevColumnMetadata ||
+        currentSchema !== state.prevCurrentSchema ||
+        currentTable !== state.prevCurrentTable)
+    ) {
+      const completions = ([] as any[]).concat(
+        uniq(columnMetadata.map(d => d.TABLE_SCHEMA)).map(v => ({
+          value: v,
+          score: 10,
+          meta: 'schema',
+        })),
+        uniq(
+          columnMetadata
+            .filter(d => (currentSchema ? d.TABLE_SCHEMA === currentSchema : true))
+            .map(d => d.TABLE_NAME),
+        ).map(v => ({
+          value: v,
+          score: 49,
+          meta: 'datasource',
+        })),
+        uniq(
+          columnMetadata
+            .filter(d =>
+              currentTable && currentSchema
+                ? d.TABLE_NAME === currentTable && d.TABLE_SCHEMA === currentSchema
+                : true,
+            )
+            .map(d => d.COLUMN_NAME),
+        ).map(v => ({
+          value: v,
+          score: 50,
+          meta: 'column',
+        })),
+      );
+
+      return {
+        completions,
+        prevColumnMetadata: columnMetadata,
+        prevCurrentSchema: currentSchema,
+        prevCurrentTable: currentTable,
+      };
+    }
+    return null;
+  }
+
+  constructor(props: QueryInputProps, context: any) {
+    super(props, context);
+    this.state = {
+      editorHeight: 200,
+      completions: [],
+    };
+  }
 
   componentDidMount(): void {
-    this.replaceDefaultAutoCompleter();
-    this.addFunctionAutoCompleter();
+    QueryInput.replaceDefaultAutoCompleter();
+    QueryInput.addFunctionAutoCompleter();
+    if (langTools) {
+      langTools.addCompleter({
+        getCompletions: (_editor: any, _session: any, _pos: any, _prefix: any, callback: any) => {
+          callback(null, this.state.completions);
+        },
+      });
+    }
   }
 
   private handleAceContainerResize = (entries: IResizeEntry[]) => {
diff --git a/web-console/src/views/query-view/query-view.tsx b/web-console/src/views/query-view/query-view.tsx
index 732f6970889..5d3a6c6e9ca 100644
--- a/web-console/src/views/query-view/query-view.tsx
+++ b/web-console/src/views/query-view/query-view.tsx
@@ -477,6 +477,21 @@ export class QueryView extends React.PureComponent<QueryViewProps, QueryViewStat
       wrapQuery,
     } = this.state;
 
+    let currentSchema;
+    let currentTable;
+
+    if (result && result.parsedQuery instanceof SqlQuery) {
+      currentSchema = result.parsedQuery.getSchema();
+      currentTable = result.parsedQuery.getTableName();
+    } else if (localStorageGet(LocalStorageKeys.QUERY_KEY)) {
+      const defaultQueryString = localStorageGet(LocalStorageKeys.QUERY_KEY);
+      const tempAst = defaultQueryString ? parser(defaultQueryString) : undefined;
+      if (tempAst) {
+        currentSchema = tempAst.getSchema();
+        currentTable = tempAst.getTableName();
+      }
+    }
+
     const runeMode = QueryView.isJsonLike(queryString);
     return (
       <SplitterLayout
@@ -491,6 +506,8 @@ export class QueryView extends React.PureComponent<QueryViewProps, QueryViewStat
       >
         <div className="control-pane">
           <QueryInput
+            currentSchema={currentSchema ? currentSchema : 'druid'}
+            currentTable={currentTable}
             queryString={queryString}
             onQueryStringChange={this.handleQueryStringChange}
             runeMode={runeMode}