[sqllab] add support for results backends (#1377)

* [sqllab] add support for results backends

Long running SQL queries (beyond the scope of a web request) can now use
a k/v store to hold their result sets.

* Addressing comments, fixed js tests

* Fixing mysql has gone away

* Adressing more comments

* Touchups

caravel/__init__.py

@@ -10,12 +10,12 @@
 
 from flask import Flask, redirect
 from flask_appbuilder import SQLA, AppBuilder, IndexView
-from sqlalchemy import event, exc
 from flask_appbuilder.baseviews import expose
 from flask_cache import Cache
 from flask_migrate import Migrate
 from caravel.source_registry import SourceRegistry
 from werkzeug.contrib.fixers import ProxyFix
+from caravel import utils
 
 
 APP_DIR = os.path.dirname(__file__)
@@ -31,33 +31,7 @@
 db = SQLA(app)
 
 
-@event.listens_for(db.engine, 'checkout')
-def checkout(dbapi_con, con_record, con_proxy):
-    """
-    Making sure the connection is live, and preventing against:
-    'MySQL server has gone away'
-
-    Copied from:
-    http://stackoverflow.com/questions/30630120/mysql-keeps-losing-connection-during-celery-tasks
-    """
-    try:
-        try:
-            if hasattr(dbapi_con, 'ping'):
-                dbapi_con.ping(False)
-            else:
-                cursor = dbapi_con.cursor()
-                cursor.execute("SELECT 1")
-        except TypeError:
-            app.logger.debug('MySQL connection died. Restoring...')
-            dbapi_con.ping()
-    except dbapi_con.OperationalError as e:
-        app.logger.warning(e)
-        if e.args[0] in (2006, 2013, 2014, 2045, 2055):
-            raise exc.DisconnectionError()
-        else:
-            raise
-    return db
-
+utils.pessimistic_connection_handling(db.engine.pool)
 
 cache = Cache(app, config=app.config.get('CACHE_CONFIG'))
 
@@ -103,6 +77,7 @@ def index(self):
 sm = appbuilder.sm
 
 get_session = appbuilder.get_session
+results_backend = app.config.get("RESULTS_BACKEND")
 
 # Registering sources
 module_datasource_map = app.config.get("DEFAULT_MODULE_DS_MAP")

caravel/assets/javascripts/SqlLab/actions.js

@@ -1,17 +1,17 @@
+import shortid from 'shortid';
+import { now } from '../modules/dates';
+const $ = require('jquery');
+
 export const RESET_STATE = 'RESET_STATE';
 export const ADD_QUERY_EDITOR = 'ADD_QUERY_EDITOR';
 export const CLONE_QUERY_TO_NEW_TAB = 'CLONE_QUERY_TO_NEW_TAB';
 export const REMOVE_QUERY_EDITOR = 'REMOVE_QUERY_EDITOR';
 export const MERGE_TABLE = 'MERGE_TABLE';
 export const REMOVE_TABLE = 'REMOVE_TABLE';
-export const START_QUERY = 'START_QUERY';
-export const STOP_QUERY = 'STOP_QUERY';
 export const END_QUERY = 'END_QUERY';
 export const REMOVE_QUERY = 'REMOVE_QUERY';
 export const EXPAND_TABLE = 'EXPAND_TABLE';
 export const COLLAPSE_TABLE = 'COLLAPSE_TABLE';
-export const QUERY_SUCCESS = 'QUERY_SUCCESS';
-export const QUERY_FAILED = 'QUERY_FAILED';
 export const QUERY_EDITOR_SETDB = 'QUERY_EDITOR_SETDB';
 export const QUERY_EDITOR_SET_SCHEMA = 'QUERY_EDITOR_SET_SCHEMA';
 export const QUERY_EDITOR_SET_TITLE = 'QUERY_EDITOR_SET_TITLE';
@@ -25,11 +25,117 @@ export const ADD_ALERT = 'ADD_ALERT';
 export const REMOVE_ALERT = 'REMOVE_ALERT';
 export const REFRESH_QUERIES = 'REFRESH_QUERIES';
 export const SET_NETWORK_STATUS = 'SET_NETWORK_STATUS';
+export const RUN_QUERY = 'RUN_QUERY';
+export const START_QUERY = 'START_QUERY';
+export const STOP_QUERY = 'STOP_QUERY';
+export const REQUEST_QUERY_RESULTS = 'REQUEST_QUERY_RESULTS';
+export const QUERY_SUCCESS = 'QUERY_SUCCESS';
+export const QUERY_FAILED = 'QUERY_FAILED';
+export const CLEAR_QUERY_RESULTS = 'CLEAR_QUERY_RESULTS';
+export const HIDE_DATA_PREVIEW = 'HIDE_DATA_PREVIEW';
 
 export function resetState() {
   return { type: RESET_STATE };
 }
 
+export function startQuery(query) {
+  Object.assign(query, {
+    id: shortid.generate(),
+    progress: 0,
+    startDttm: now(),
+    state: (query.runAsync) ? 'pending' : 'running',
+  });
+  return { type: START_QUERY, query };
+}
+
+export function querySuccess(query, results) {
+  return { type: QUERY_SUCCESS, query, results };
+}
+
+export function queryFailed(query, msg) {
+  return { type: QUERY_FAILED, query, msg };
+}
+
+export function stopQuery(query) {
+  return { type: STOP_QUERY, query };
+}
+
+export function clearQueryResults(query) {
+  return { type: CLEAR_QUERY_RESULTS, query };
+}
+
+export function hideDataPreview() {
+  return { type: HIDE_DATA_PREVIEW };
+}
+
+export function requestQueryResults(query) {
+  return { type: REQUEST_QUERY_RESULTS, query };
+}
+
+export function fetchQueryResults(query) {
+  return function (dispatch) {
+    dispatch(requestQueryResults(query));
+    const sqlJsonUrl = `/caravel/results/${query.resultsKey}/`;
+    $.ajax({
+      type: 'GET',
+      dataType: 'json',
+      url: sqlJsonUrl,
+      success(results) {
+        dispatch(querySuccess(query, results));
+      },
+      error() {
+        dispatch(queryFailed(query, 'Failed at retrieving results from the results backend'));
+      },
+    });
+  };
+}
+
+export function runQuery(query) {
+  return function (dispatch) {
+    dispatch(startQuery(query));
+    const sqlJsonUrl = '/caravel/sql_json/';
+    const sqlJsonRequest = {
+      client_id: query.id,
+      database_id: query.dbId,
+      json: true,
+      runAsync: query.runAsync,
+      schema: query.schema,
+      sql: query.sql,
+      sql_editor_id: query.sqlEditorId,
+      tab: query.tab,
+      tmp_table_name: query.tempTableName,
+      select_as_cta: query.ctas,
+    };
+    $.ajax({
+      type: 'POST',
+      dataType: 'json',
+      url: sqlJsonUrl,
+      data: sqlJsonRequest,
+      success(results) {
+        if (!query.runAsync) {
+          dispatch(querySuccess(query, results));
+        }
+      },
+      error(err, textStatus, errorThrown) {
+        let msg;
+        try {
+          msg = err.responseJSON.error;
+        } catch (e) {
+          if (err.responseText !== undefined) {
+            msg = err.responseText;
+          }
+        }
+        if (textStatus === 'error' && errorThrown === '') {
+          msg = 'Could not connect to server';
+        } else if (msg === null) {
+          msg = `[${textStatus}] ${errorThrown}`;
+        }
+        dispatch(queryFailed(query, msg));
+      },
+    });
+  };
+}
+
 export function setDatabases(databases) {
   return { type: SET_DATABASES, databases };
 }
@@ -102,22 +208,6 @@ export function removeTable(table) {
   return { type: REMOVE_TABLE, table };
 }
 
-export function startQuery(query) {
-  return { type: START_QUERY, query };
-}
-
-export function stopQuery(query) {
-  return { type: STOP_QUERY, query };
-}
-
-export function querySuccess(query, results) {
-  return { type: QUERY_SUCCESS, query, results };
-}
-
-export function queryFailed(query, msg) {
-  return { type: QUERY_FAILED, query, msg };
-}
-
 export function addWorkspaceQuery(query) {
   return { type: ADD_WORKSPACE_QUERY, query };
 }

caravel/assets/javascripts/SqlLab/common.js

@@ -1,8 +1,12 @@
 export const STATE_BSSTYLE_MAP = {
   failed: 'danger',
   pending: 'info',
+  fetching: 'info',
   running: 'warning',
+  stopped: 'danger',
   success: 'success',
 };
 
+export const DATA_PREVIEW_ROW_COUNT = 100;
+
 export const STATUS_OPTIONS = ['success', 'failed', 'running'];

caravel/assets/javascripts/SqlLab/components/App.jsx

@@ -5,6 +5,7 @@ import TabbedSqlEditors from './TabbedSqlEditors';
 import QueryAutoRefresh from './QueryAutoRefresh';
 import QuerySearch from './QuerySearch';
 import Alerts from './Alerts';
+import DataPreviewModal from './DataPreviewModal';
 
 import { bindActionCreators } from 'redux';
 import { connect } from 'react-redux';
@@ -26,8 +27,9 @@ class App extends React.Component {
     this.setState({ hash: window.location.hash });
   }
   render() {
+    let content;
     if (this.state.hash) {
-      return (
+      content = (
         <div className="container-fluid">
           <div className="row">
             <div className="col-md-12">
@@ -37,16 +39,18 @@ class App extends React.Component {
         </div>
       );
     }
+    content = (
+      <div>
+        <QueryAutoRefresh />
+        <TabbedSqlEditors />
+      </div>
+    );
     return (
       <div className="App SqlLab">
+        <Alerts alerts={this.props.alerts} />
+        <DataPreviewModal />
         <div className="container-fluid">
-          <QueryAutoRefresh />
-          <Alerts alerts={this.props.alerts} />
-          <div className="row">
-            <div className="col-md-12">
-              <TabbedSqlEditors />
-            </div>
-          </div>
+          {content}
         </div>
       </div>
     );

caravel/assets/javascripts/SqlLab/components/DataPreviewModal.jsx

@@ -0,0 +1,58 @@
+import * as Actions from '../actions';
+import React from 'react';
+import { bindActionCreators } from 'redux';
+import { connect } from 'react-redux';
+import { Modal } from 'react-bootstrap';
+
+import ResultSet from './ResultSet';
+
+const propTypes = {
+  queries: React.PropTypes.object,
+  actions: React.PropTypes.object,
+  showDataPreviewModal: React.PropTypes.bool,
+  dataPreviewQueryId: React.PropTypes.string,
+};
+
+class DataPreviewModal extends React.Component {
+  hide() {
+    this.props.actions.hideDataPreview();
+  }
+  render() {
+    if (this.props.showDataPreviewModal && this.props.dataPreviewQueryId) {
+      const query = this.props.queries[this.props.dataPreviewQueryId];
+      return (
+        <Modal
+          show={this.props.showDataPreviewModal}
+          onHide={this.hide.bind(this)}
+          bsStyle="lg"
+        >
+          <Modal.Header closeButton>
+            <Modal.Title>
+              Data preview for <strong>{query.tableName}</strong>
+            </Modal.Title>
+          </Modal.Header>
+          <Modal.Body>
+            <ResultSet query={query} visualize={false} csv={false} />
+          </Modal.Body>
+        </Modal>
+      );
+    }
+    return null;
+  }
+}
+DataPreviewModal.propTypes = propTypes;
+
+function mapStateToProps(state) {
+  return {
+    queries: state.queries,
+    showDataPreviewModal: state.showDataPreviewModal,
+    dataPreviewQueryId: state.dataPreviewQueryId,
+  };
+}
+function mapDispatchToProps(dispatch) {
+  return {
+    actions: bindActionCreators(Actions, dispatch),
+  };
+}
+
+export default connect(mapStateToProps, mapDispatchToProps)(DataPreviewModal);

caravel/assets/javascripts/SqlLab/components/SqlShrink.jsx → caravel/assets/javascripts/SqlLab/components/HighlightedSql.jsx

@@ -2,38 +2,43 @@ import React from 'react';
 import SyntaxHighlighter from 'react-syntax-highlighter';
 import { github } from 'react-syntax-highlighter/dist/styles';
 
-const SqlShrink = (props) => {
+const HighlightedSql = (props) => {
   const sql = props.sql || '';
   let lines = sql.split('\n');
   if (lines.length >= props.maxLines) {
     lines = lines.slice(0, props.maxLines);
     lines.push('{...}');
   }
-  const shrunk = lines.map((line) => {
-    if (line.length > props.maxWidth) {
-      return line.slice(0, props.maxWidth) + '{...}';
-    }
-    return line;
-  })
-  .join('\n');
+  let shownSql = sql;
+  if (props.shrink) {
+    shownSql = lines.map((line) => {
+      if (line.length > props.maxWidth) {
+        return line.slice(0, props.maxWidth) + '{...}';
+      }
+      return line;
+    })
+    .join('\n');
+  }
   return (
     <div>
       <SyntaxHighlighter language="sql" style={github}>
-        {shrunk}
+        {shownSql}
       </SyntaxHighlighter>
     </div>
   );
 };
 
-SqlShrink.defaultProps = {
+HighlightedSql.defaultProps = {
   maxWidth: 60,
   maxLines: 6,
+  shrink: false,
 };
 
-SqlShrink.propTypes = {
+HighlightedSql.propTypes = {
   sql: React.PropTypes.string,
   maxWidth: React.PropTypes.number,
   maxLines: React.PropTypes.number,
+  shrink: React.PropTypes.bool,
 };
 
-export default SqlShrink;
+export default HighlightedSql;