[SQL Lab] Add function names to autocomplete (#9012)

superset/assets/src/SqlLab/components/AceEditorWrapper.jsx

@@ -29,6 +29,7 @@ import {
   SCHEMA_AUTOCOMPLETE_SCORE,
   TABLE_AUTOCOMPLETE_SCORE,
   COLUMN_AUTOCOMPLETE_SCORE,
+  SQL_FUNCTIONS_AUTOCOMPLETE_SCORE,
 } from '../constants';
 
 const langTools = ace.acequire('ace/ext/language_tools');
@@ -39,6 +40,7 @@ const propTypes = {
   sql: PropTypes.string.isRequired,
   schemas: PropTypes.array,
   tables: PropTypes.array,
+  functionNames: PropTypes.array,
   extendedTables: PropTypes.array,
   queryEditor: PropTypes.object.isRequired,
   height: PropTypes.string,
@@ -57,6 +59,7 @@ const defaultProps = {
   onChange: () => {},
   schemas: [],
   tables: [],
+  functionNames: [],
   extendedTables: [],
 };
 
@@ -145,7 +148,9 @@ class AceEditorWrapper extends React.PureComponent {
             this.props.queryEditor.schema,
           );
         }
-        editor.completer.insertMatch({ value: data.caption + ' ' });
+        editor.completer.insertMatch({
+          value: `${data.caption}${data.meta === 'function' ? '' : ' '}`,
+        });
       },
     };
     const words = this.state.words.map(word => ({ ...word, completer }));
@@ -185,9 +190,17 @@ class AceEditorWrapper extends React.PureComponent {
       meta: 'column',
     }));
 
+    const functionWords = props.functionNames.map(func => ({
+      name: func,
+      value: func,
+      score: SQL_FUNCTIONS_AUTOCOMPLETE_SCORE,
+      meta: 'function',
+    }));
+
     const words = schemaWords
       .concat(tableWords)
       .concat(columnWords)
+      .concat(functionWords)
       .concat(sqlKeywords);
 
     this.setState({ words }, () => {

superset/assets/src/SqlLab/components/SqlEditor.jsx

@@ -343,6 +343,9 @@ class SqlEditor extends React.PureComponent {
             sql={this.props.queryEditor.sql}
             schemas={this.props.queryEditor.schemaOptions}
             tables={this.props.queryEditor.tableOptions}
+            functionNames={
+              this.props.database ? this.props.database.function_names : []
+            }
             extendedTables={this.props.tables}
             height={`${aceEditorHeight}px`}
             hotkeys={hotkeys}

superset/assets/src/SqlLab/constants.js

@@ -61,6 +61,7 @@ export const LOCALSTORAGE_WARNING_MESSAGE_THROTTLE_MS = 8000; // danger type toa
 
 // autocomplete score weights
 export const SQL_KEYWORD_AUTOCOMPLETE_SCORE = 100;
+export const SQL_FUNCTIONS_AUTOCOMPLETE_SCORE = 90;
 export const SCHEMA_AUTOCOMPLETE_SCORE = 60;
 export const TABLE_AUTOCOMPLETE_SCORE = 55;
 export const COLUMN_AUTOCOMPLETE_SCORE = 50;

superset/db_engine_specs/base.py

@@ -846,6 +846,17 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
         """
         return sqla_column_type.compile(dialect=dialect).upper()
 
+    @classmethod
+    def get_function_names(cls, database: "Database") -> List[str]:
+        """
+        Get a list of function names that are able to be called on the database.
+        Used for SQL Lab autocomplete.
+
+        :param database: The database to get functions for
+        :return: A list of function names useable in the database
+        """
+        return []
+
     @staticmethod
     def pyodbc_rows_to_tuples(data: List[Any]) -> List[Tuple]:
         """

superset/db_engine_specs/hive.py

@@ -19,7 +19,7 @@ import os
 import re
 import time
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
 from urllib import parse
 
 from sqlalchemy import Column
@@ -34,6 +34,10 @@ from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.presto import PrestoEngineSpec
 from superset.utils import core as utils
 
+if TYPE_CHECKING:
+    # prevent circular imports
+    from superset.models.core import Database  # pylint: disable=unused-import
+
 QueryStatus = utils.QueryStatus
 config = app.config
 
@@ -422,3 +426,14 @@ class HiveEngineSpec(PrestoEngineSpec):
     ):  # pylint: disable=arguments-differ
         kwargs = {"async": async_}
         cursor.execute(query, **kwargs)
+
+    @classmethod
+    def get_function_names(cls, database: "Database") -> List[str]:
+        """
+        Get a list of function names that are able to be called on the database.
+        Used for SQL Lab autocomplete.
+
+        :param database: The database to get functions for
+        :return: A list of function names useable in the database
+        """
+        return database.get_df("SHOW FUNCTIONS")["tab_name"].tolist()

superset/db_engine_specs/presto.py

@@ -945,3 +945,14 @@ class PrestoEngineSpec(BaseEngineSpec):
         if df.empty:
             return ""
         return df.to_dict()[field_to_return][0]
+
+    @classmethod
+    def get_function_names(cls, database: "Database") -> List[str]:
+        """
+        Get a list of function names that are able to be called on the database.
+        Used for SQL Lab autocomplete.
+
+        :param database: The database to get functions for
+        :return: A list of function names useable in the database
+        """
+        return database.get_df("SHOW FUNCTIONS")["Function"].tolist()

superset/models/core.py

@@ -161,6 +161,10 @@ class Database(
     def allows_subquery(self) -> bool:
         return self.db_engine_spec.allows_subqueries
 
+    @property
+    def function_names(self) -> List[str]:
+        return self.db_engine_spec.get_function_names(self)
+
     @property
     def allows_cost_estimate(self) -> bool:
         extra = self.get_extra()
@@ -320,7 +324,7 @@ class Database(
         return self.get_dialect().identifier_preparer.quote
 
     def get_df(  # pylint: disable=too-many-locals
-        self, sql: str, schema: str, mutator: Optional[Callable] = None
+        self, sql: str, schema: Optional[str] = None, mutator: Optional[Callable] = None
     ) -> pd.DataFrame:
         sqls = [str(s).strip(" ;") for s in sqlparse.parse(sql)]
         source_key = None

superset/views/database/api.py

@@ -44,6 +44,7 @@ class DatabaseRestApi(DatabaseMixin, BaseSupersetModelRestApi):
         "allows_subquery",
         "allows_cost_estimate",
         "backend",
+        "function_names",
     ]
     show_columns = list_columns