Implement caching and dynamic data fetching. (#1466)

* Rename rv => o in the decorator.

* Address comments.

* Permissions cleanup: remove none and duplicates. (#1967)

* Updates

* Rename var and dropdown text

* Cleanup

* Resolve comments.

* Add user to the perm check.

superset/__init__.py

@@ -29,13 +29,16 @@
     # In production mode, add log handler to sys.stderr.
     app.logger.addHandler(logging.StreamHandler())
     app.logger.setLevel(logging.INFO)
+logging.getLogger('pyhive.presto').setLevel(logging.INFO)
 
 db = SQLA(app)
 
 
 utils.pessimistic_connection_handling(db.engine.pool)
 
 cache = Cache(app, config=app.config.get('CACHE_CONFIG'))
+tables_cache = Cache(app, config=app.config.get('TABLE_NAMES_CACHE_CONFIG'))
+
 
 migrate = Migrate(app, db, directory=APP_DIR + "/migrations")
 

superset/assets/javascripts/SqlLab/actions.js

@@ -213,9 +213,9 @@ export function mergeTable(table, query) {
   return { type: MERGE_TABLE, table, query };
 }
 
-export function addTable(query, tableName) {
+export function addTable(query, tableName, schemaName) {
   return function (dispatch) {
-    let url = `/superset/table/${query.dbId}/${tableName}/${query.schema}/`;
+    let url = `/superset/table/${query.dbId}/${tableName}/${schemaName}/`;
     $.get(url, (data) => {
       const dataPreviewQuery = {
         id: shortid.generate(),
@@ -232,7 +232,7 @@ export function addTable(query, tableName) {
         Object.assign(data, {
           dbId: query.dbId,
           queryEditorId: query.id,
-          schema: query.schema,
+          schema: schemaName,
           expanded: true,
         }), dataPreviewQuery)
       );
@@ -248,12 +248,12 @@ export function addTable(query, tableName) {
       );
     });
 
-    url = `/superset/extra_table_metadata/${query.dbId}/${tableName}/${query.schema}/`;
+    url = `/superset/extra_table_metadata/${query.dbId}/${tableName}/${schemaName}/`;
     $.get(url, (data) => {
       const table = {
         dbId: query.dbId,
         queryEditorId: query.id,
-        schema: query.schema,
+        schema: schemaName,
         name: tableName,
       };
       Object.assign(table, data);

superset/assets/javascripts/SqlLab/components/SqlEditorLeftBar.jsx

@@ -30,8 +30,8 @@ class SqlEditorLeftBar extends React.PureComponent {
     };
   }
   componentWillMount() {
-    this.fetchSchemas();
-    this.fetchTables();
+    this.fetchSchemas(this.props.queryEditor.dbId);
+    this.fetchTables(this.props.queryEditor.dbId, this.props.queryEditor.schema);
   }
   onChange(db) {
     const val = (db) ? db.value : null;
@@ -58,22 +58,51 @@ class SqlEditorLeftBar extends React.PureComponent {
   resetState() {
     this.props.actions.resetState();
   }
-  fetchTables(dbId, schema) {
-    const actualDbId = dbId || this.props.queryEditor.dbId;
-    if (actualDbId) {
-      const actualSchema = schema || this.props.queryEditor.schema;
-      this.setState({ tableLoading: true });
-      this.setState({ tableOptions: [] });
-      const url = `/superset/tables/${actualDbId}/${actualSchema}`;
+  getTableNamesBySubStr(input) {
+    if (!this.props.queryEditor.dbId || !input) {
+      return Promise.resolve({ options: [] });
+    }
+    const url = `/superset/tables/${this.props.queryEditor.dbId}/\
+${this.props.queryEditor.schema}/${input}`;
+    return $.get(url).then((data) => ({ options: data.options }));
+  }
+  // TODO: move fetching methods to the actions.
+  fetchTables(dbId, schema, substr) {
+    if (dbId) {
+      this.setState({ tableLoading: true, tableOptions: [] });
+      const url = `/superset/tables/${dbId}/${schema}/${substr}/`;
       $.get(url, (data) => {
-        let tableOptions = data.tables.map((s) => ({ value: s, label: s }));
-        const views = data.views.map((s) => ({ value: s, label: '[view] ' + s }));
-        tableOptions = [...tableOptions, ...views];
-        this.setState({ tableOptions });
-        this.setState({ tableLoading: false });
+        this.setState({
+          tableLoading: false,
+          tableOptions: data.options,
+          tableLength: data.tableLength,
+        });
       });
     }
   }
+  changeTable(tableOpt) {
+    if (!tableOpt) {
+      this.setState({ tableName: '' });
+      return;
+    }
+    const namePieces = tableOpt.value.split('.');
+    let tableName = namePieces[0];
+    let schemaName = this.props.queryEditor.schema;
+    if (namePieces.length === 1) {
+      this.setState({ tableName });
+    } else {
+      schemaName = namePieces[0];
+      tableName = namePieces[1];
+      this.setState({ tableName });
+      this.props.actions.queryEditorSetSchema(this.props.queryEditor, schemaName);
+      this.fetchTables(this.props.queryEditor.dbId, schemaName);
+    }
+    this.setState({ tableLoading: true });
+    // TODO: handle setting the tableLoading state depending on success or
+    //       failure of the addTable async call in the action.
+    this.props.actions.addTable(this.props.queryEditor, tableName, schemaName);
+    this.setState({ tableLoading: false });
+  }
   changeSchema(schemaOpt) {
     const schema = (schemaOpt) ? schemaOpt.value : null;
     this.props.actions.queryEditorSetSchema(this.props.queryEditor, schema);
@@ -95,14 +124,6 @@ class SqlEditorLeftBar extends React.PureComponent {
   closePopover(ref) {
     this.refs[ref].hide();
   }
-  changeTable(tableOpt) {
-    const tableName = tableOpt.value;
-    const qe = this.props.queryEditor;
-
-    this.setState({ tableLoading: true });
-    this.props.actions.addTable(qe, tableName);
-    this.setState({ tableLoading: false });
-  }
   render() {
     let networkAlert = null;
     if (!this.props.networkOn) {
@@ -118,6 +139,8 @@ class SqlEditorLeftBar extends React.PureComponent {
               dataEndpoint="/databaseasync/api/read?_flt_0_expose_in_sqllab=1"
               onChange={this.onChange.bind(this)}
               value={this.props.queryEditor.dbId}
+              databaseId={this.props.queryEditor.dbId}
+              actions={this.props.actions}
               valueRenderer={(o) => (
                 <div>
                   <span className="text-muted">Database:</span> {o.label}
@@ -126,8 +149,6 @@ class SqlEditorLeftBar extends React.PureComponent {
               mutator={this.dbMutator.bind(this)}
               placeholder="Select a database"
             />
-          </div>
-          <div className="m-t-5">
             <Select
               name="select-schema"
               placeholder={`Select a schema (${this.state.schemaOptions.length})`}
@@ -144,15 +165,29 @@ class SqlEditorLeftBar extends React.PureComponent {
             />
           </div>
           <div className="m-t-5">
-            <Select
-              name="select-table"
-              ref="selectTable"
-              isLoading={this.state.tableLoading}
-              placeholder={`Add a table (${this.state.tableOptions.length})`}
-              autosize={false}
-              onChange={this.changeTable.bind(this)}
-              options={this.state.tableOptions}
-            />
+            {this.props.queryEditor.schema &&
+              <Select
+                name="select-table"
+                ref="selectTable"
+                isLoading={this.state.tableLoading}
+                value={this.state.tableName}
+                placeholder={`Add a table (${this.state.tableOptions.length})`}
+                autosize={false}
+                onChange={this.changeTable.bind(this)}
+                options={this.state.tableOptions}
+              />
+            }
+            {!this.props.queryEditor.schema &&
+              <Select.Async
+                name="async-select-table"
+                ref="selectTable"
+                value={this.state.tableName}
+                placeholder={"Type to search ..."}
+                autosize={false}
+                onChange={this.changeTable.bind(this)}
+                loadOptions={this.getTableNamesBySubStr.bind(this)}
+              />
+            }
           </div>
           <hr />
           <div className="m-t-5">

superset/cache_util.py

@@ -0,0 +1,27 @@
+from superset import tables_cache
+from flask import request
+
+
+def view_cache_key(*unused_args, **unused_kwargs):
+    args_hash = hash(frozenset(request.args.items()))
+    return 'view/{}/{}'.format(request.path, args_hash)
+
+
+def memoized_func(timeout=5 * 60, key=view_cache_key):
+    """Use this decorator to cache functions that have predefined first arg.
+
+    memoized_func uses simple_cache and stored the data in memory.
+    Key is a callable function that takes function arguments and
+    returns the caching key.
+    """
+    def wrap(f):
+        def wrapped_f(cls, *args, **kwargs):
+            cache_key = key(*args, **kwargs)
+            o = tables_cache.get(cache_key)
+            if o is not None:
+                return o
+            o = f(cls, *args, **kwargs)
+            tables_cache.set(cache_key, o, timeout=timeout)
+            return o
+        return wrapped_f
+    return wrap

superset/config.py

@@ -153,6 +153,7 @@
 
 CACHE_DEFAULT_TIMEOUT = None
 CACHE_CONFIG = {'CACHE_TYPE': 'null'}
+TABLE_NAMES_CACHE_CONFIG = {'CACHE_TYPE': 'null'}
 
 # CORS Options
 ENABLE_CORS = False
@@ -209,6 +210,9 @@
 SQL_MAX_ROW = 1000000
 DISPLAY_SQL_MAX_ROW = 1000
 
+# Maximum number of tables/views displayed in the dropdown window in SQL Lab.
+MAX_TABLE_NAMES = 3000
+
 # If defined, shows this text in an alert-warning box in the navbar
 # one example use case may be "STAGING" to make it clear that this is
 # not the production version of the site.

superset/db_engine_specs.py

@@ -16,12 +16,13 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-from collections import namedtuple
+from collections import namedtuple, defaultdict
+from flask_babel import lazy_gettext as _
 import inspect
 import textwrap
 import time
 
-from flask_babel import lazy_gettext as _
+from superset import cache_util
 
 Grain = namedtuple('Grain', 'name label function')
 
@@ -54,6 +55,33 @@ def extra_table_metadata(cls, database, table_name, schema_name):
     def convert_dttm(cls, target_type, dttm):
         return "'{}'".format(dttm.strftime('%Y-%m-%d %H:%M:%S'))
 
+    @classmethod
+    @cache_util.memoized_func(
+        timeout=600,
+        key=lambda *args, **kwargs: 'db:{}:{}'.format(args[0].id, args[1]))
+    def fetch_result_sets(cls, db, datasource_type):
+        """Returns the dictionary {schema : [result_set_name]}.
+
+        Datasource_type can be 'table' or 'view'.
+        Empty schema corresponds to the list of full names of the all
+        tables or views: <schema>.<result_set_name>.
+        """
+        schemas = db.inspector.get_schema_names()
+        result_sets = {}
+        all_result_sets = []
+        for schema in schemas:
+            if datasource_type == 'table':
+                result_sets[schema] = sorted(
+                    db.inspector.get_table_names(schema))
+            elif datasource_type == 'view':
+                result_sets[schema] = sorted(
+                    db.inspector.get_view_names(schema))
+            all_result_sets += [
+                '{}.{}'.format(schema, t) for t in result_sets[schema]]
+        if all_result_sets:
+            result_sets[""] = all_result_sets
+        return result_sets
+
     @classmethod
     def handle_cursor(cls, cursor, query, session):
         """Handle a live cursor between the execute and fetchall calls
@@ -221,6 +249,28 @@ def show_partition_pql(
         {limit_clause}
         """).format(**locals())
 
+    @classmethod
+    @cache_util.memoized_func(
+        timeout=600,
+        key=lambda *args, **kwargs: 'db:{}:{}'.format(args[0].id, args[1]))
+    def fetch_result_sets(cls, db, datasource_type):
+        """Returns the dictionary {schema : [result_set_name]}.
+
+        Datasource_type can be 'table' or 'view'.
+        Empty schema corresponds to the list of full names of the all
+        tables or views: <schema>.<result_set_name>.
+        """
+        result_set_df = db.get_df(
+            """SELECT table_schema, table_name FROM INFORMATION_SCHEMA.{}S
+               ORDER BY concat(table_schema, '.', table_name)""".format(
+                datasource_type.upper()), None)
+        result_sets = defaultdict(list)
+        for unused, row in result_set_df.iterrows():
+            result_sets[row['table_schema']].append(row['table_name'])
+            result_sets[""].append('{}.{}'.format(
+                row['table_schema'], row['table_name']))
+        return result_sets
+
     @classmethod
     def extra_table_metadata(cls, database, table_name, schema_name):
         indexes = database.get_indexes(table_name, schema_name)

superset/models.py

@@ -845,13 +845,19 @@ def inspector(self):
         return sqla.inspect(engine)
 
     def all_table_names(self, schema=None):
+        if not schema:
+            tables_dict = self.db_engine_spec.fetch_result_sets(self, 'table')
+            return tables_dict.get("", [])
         return sorted(self.inspector.get_table_names(schema))
 
     def all_view_names(self, schema=None):
+        if not schema:
+            views_dict = self.db_engine_spec.fetch_result_sets(self, 'view')
+            return views_dict.get("", [])
         views = []
         try:
             views = self.inspector.get_view_names(schema)
-        except Exception as e:
+        except Exception:
             pass
         return views
 

superset/source_registry.py

@@ -44,25 +44,14 @@ def get_datasource_by_name(cls, session, datasource_type, datasource_name,
         return db_ds[0]
 
     @classmethod
-    def query_datasources_by_name(
-            cls, session, database, datasource_name, schema=None):
+    def query_datasources_by_permissions(cls, session, database, permissions):
         datasource_class = SourceRegistry.sources[database.type]
-        if database.type == 'table':
-            query = (
-                session.query(datasource_class)
-                .filter_by(database_id=database.id)
-                .filter_by(table_name=datasource_name))
-            if schema:
-                query = query.filter_by(schema=schema)
-            return query.all()
-        if database.type == 'druid':
-            return (
-                session.query(datasource_class)
-                .filter_by(cluster_name=database.id)
-                .filter_by(datasource_name=datasource_name)
-                .all()
-            )
-        return None
+        return (
+            session.query(datasource_class)
+            .filter_by(database_id=database.id)
+            .filter(datasource_class.perm.in_(permissions))
+            .all()
+        )
 
     @classmethod
     def get_eager_datasource(cls, session, datasource_type, datasource_id):