feat: Add validate_sql_json endpoint for checking that a given sql query is valid for the chosen database

superset/config.py

@@ -417,6 +417,9 @@ class CeleryConfig(object):
 # Timeout duration for SQL Lab synchronous queries
 SQLLAB_TIMEOUT = 30
 
+# Timeout duration for SQL Lab query validation
+SQLLAB_VALIDATION_TIMEOUT = 10
+
 # SQLLAB_DEFAULT_DBID
 SQLLAB_DEFAULT_DBID = None
 

superset/sql_validators/__init__.py

@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from . import base  # noqa
+from . import presto_db  # noqa
+from .base import SQLValidationAnnotation  # noqa
+
+# TODO: Move this to a config setting
+SQL_VALIDATORS_BY_ENGINE = {
+    'presto': presto_db.PrestoDBSQLValidator
+}

superset/sql_validators/base.py

@@ -0,0 +1,61 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import (
+    Any,
+    List,
+    Optional
+)
+
+class SQLValidationAnnotation:
+    """Represents a single annotation (error/warning) in an SQL querytext"""
+    def __init__(
+        self,
+        message: str,
+        line_number: Optional[int],
+        start_column: Optional[int],
+        end_column: Optional[int]
+    ):
+        self.message = message
+        self.line_number = line_number
+        self.start_column = start_column
+        self.end_column = end_column
+
+    def to_dict(self):
+        return {
+            "line_number": self.line_number,
+            "start_column": self.start_column,
+            "end_column": self.end_column,
+            "message": self.message,
+        }
+
+
+class BaseSQLValidator:
+    """BaseSQLValidator defines the interface for checking that a given sql
+    query is valid for a given database engine."""
+
+    name = 'BaseSQLValidator'
+
+    @classmethod
+    def validate(
+        cls,
+        sql: str,
+        schema: str,
+        database: Any
+    ) -> List[SQLValidationAnnotation]:
+        """Check that the given SQL querystring is valid for the given engine"""
+        raise NotImplementedError

superset/sql_validators/presto_db.py

@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import time
+import json
+import logging
+from contextlib import closing
+from flask import g
+from pyhive.exc import DatabaseError
+from typing import (
+    Any,
+    List,
+    Optional
+)
+
+from superset import app, security_manager
+from superset.utils.core import sources
+from superset.sql_parse import ParsedQuery
+from superset.sql_validators.base import (
+    BaseSQLValidator,
+    SQLValidationAnnotation,
+)
+
+MAX_ERROR_ROWS = 10
+
+config = app.config
+
+class PrestoSQLValidationError(Exception):
+    """Error in the process of asking Presto to validate SQL querytext"""
+    pass
+
+class PrestoDBSQLValidator(BaseSQLValidator):
+    """Validate SQL queries using Presto's built-in EXPLAIN subtype"""
+
+    name = 'PrestoDBSQLValidator'
+
+    @classmethod
+    def validate_statement(
+        cls,
+        statement,
+        database,
+        cursor,
+        user_name
+    ) -> Optional[SQLValidationAnnotation]:
+        db_engine_spec = database.db_engine_spec
+        parsed_query = ParsedQuery(statement)
+        sql = parsed_query.stripped()
+
+        # Hook to allow environment-specific mutation (usually comments) to the SQL
+        SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
+        if SQL_QUERY_MUTATOR:
+            sql = SQL_QUERY_MUTATOR(sql, user_name, security_manager, database)
+
+        # Transform the final statement to an explain call before sending it on
+        # to presto to validate
+        sql = f'EXPLAIN (TYPE VALIDATE) {sql}'
+
+        # Invoke the query against presto. NB this deliberately doesn't use the
+        # engine spec's handle_cursor implementation since we don't record
+        # these EXPLAIN queries done in validation as proper Query objects
+        # in the superset ORM.
+        try:
+            db_engine_spec.execute(cursor, sql)
+            polled = cursor.poll()
+            while polled:
+                logging.info('polling presto for validation progress')
+                stats = polled.get('stats', {})
+                if stats:
+                    state = stats.get('state')
+                    if state == 'FINISHED':
+                        break
+                time.sleep(0.2)
+                polled = cursor.poll()
+            db_engine_spec.fetch_data(cursor, MAX_ERROR_ROWS)
+            return None
+        except DatabaseError as db_error:
+            if not db_error.args or not isinstance(db_error.args[0], dict):
+                raise PrestoSQLValidationError(
+                    "Presto (via pyhive) returned unparseable error text")
+            db_error = db_error.args[0]
+
+            message = db_error.get('message', "unknown prestodb error")
+            err_loc = db_error.get('errorLocation', {})
+            line_number = err_loc.get('lineNumber', None)
+            start_column = err_loc.get('columnNumber', None)
+            end_column = err_loc.get('columnNumber', None)
+
+            return SQLValidationAnnotation(
+                message=message,
+                line_number=line_number,
+                start_column=start_column,
+                end_column=end_column,
+            )
+        except Exception as e:
+            logging.exception(f'Error running validation query: {e}')
+            raise e
+
+    @classmethod
+    def validate(
+        cls,
+        sql: str,
+        schema: str,
+        database: Any
+    ) -> List[SQLValidationAnnotation]:
+        """
+        Presto supports query-validation queries by running them with a
+        prepended explain.
+
+        For example, "SELECT 1 FROM default.mytable" becomes "EXPLAIN (TYPE
+        VALIDATE) SELECT 1 FROM default.mytable.
+        """
+        user_name = g.user.username if g.user else None
+        parsed_query = ParsedQuery(sql)
+        statements = parsed_query.get_statements()
+
+        logging.info(f'Validating {len(statements)} statement(s)')
+        engine = database.get_sqla_engine(
+            schema=schema,
+            nullpool=True,
+            user_name=user_name,
+            source=sources.get('sql_lab', None),
+        )
+        # Sharing a single connection and cursor across the
+        # execution of all statements (if many)
+        annotations: List[SQLValidationAnnotation] = []
+        with closing(engine.raw_connection()) as conn:
+            with closing(conn.cursor()) as cursor:
+                for statement in parsed_query.get_statements():
+                    annotation = cls.validate_statement(
+                        statement,
+                        database,
+                        cursor,
+                        user_name
+                    )
+                    if annotation:
+                        annotations.append(annotation)
+        logging.debug(f'Validation found {len(annotations)} error(s)')
+
+        return annotations

superset/views/core.py

@@ -56,6 +56,7 @@
 from superset.models.sql_lab import Query
 from superset.models.user_attributes import UserAttribute
 from superset.sql_parse import ParsedQuery
+from superset.sql_validators import SQL_VALIDATORS_BY_ENGINE
 from superset.utils import core as utils
 from superset.utils import dashboard_import_export
 from superset.utils.dates import now_as_float
@@ -2503,6 +2504,61 @@ def stop_query(self):
             pass
         return self.json_response('OK')
 
+    @has_access_api
+    @expose('/validate_sql_json/', methods=['POST', 'GET'])
+    @log_this
+    def validate_sql_json(self):
+        """Validates that arbitrary sql is acceptable for the given database.
+        Returns a list of error/warning annotations as json.
+        """
+        sql = request.form.get('sql')
+        database_id = request.form.get('database_id')
+        schema = request.form.get('schema') or None
+        template_params = json.loads(
+            request.form.get('templateParams') or '{}')
+
+        if len(template_params) > 0:
+            # TODO: factor the Database object out of template rendering
+            #       or provide it as mydb so we can render template params
+            #       without having to also persist a Query ORM object.
+            return json_error_response(
+                'SQL validation does not support template parameters')
+
+        session = db.session()
+        mydb = session.query(models.Database).filter_by(id=database_id).first()
+        if not mydb:
+            json_error_response(
+                'Database with id {} is missing.'.format(database_id))
+
+        spec = mydb.db_engine_spec
+        if not spec.engine in SQL_VALIDATORS_BY_ENGINE:
+            return json_error_response(
+                'no SQL validator is configured for {}'.format(spec.engine))
+        validator = SQL_VALIDATORS_BY_ENGINE[spec.engine]
+
+        try:
+            timeout = config.get('SQLLAB_VALIDATION_TIMEOUT')
+            timeout_msg = (
+                f'The query exceeded the {timeout} seconds timeout.')
+            with utils.timeout(seconds=timeout,
+                                error_message=timeout_msg):
+                errors = validator.validate(sql, schema, mydb)
+            payload = json.dumps(
+                [err.to_dict() for err in errors],
+                default=utils.pessimistic_json_iso_dttm_ser,
+                ignore_nan=True,
+                encoding=None,
+            )
+            return json_success(payload)
+        except Exception as e:
+            logging.exception(e)
+            msg = _(
+                'Failed to validate your SQL query text. Please check that '
+                f'you have configured the {validator.name} validator '
+                'correctly and that any services it depends on are up. '
+                f'Exception: {e}')
+            return json_error_response(f'{msg}')
+
     @has_access_api
     @expose('/sql_json/', methods=['POST', 'GET'])
     @log_this

tests/base_tests.py

@@ -190,6 +190,21 @@ def run_sql(self, sql, client_id=None, user_name=None, raise_on_error=False,
             raise Exception('run_sql failed')
         return resp
 
+    def validate_sql(self, sql, client_id=None, user_name=None,
+                     raise_on_error=False):
+        if user_name:
+            self.logout()
+            self.login(username=(user_name if user_name else 'admin'))
+        dbid = get_main_database(db.session).id
+        resp = self.get_json_resp(
+            '/superset/validate_sql_json/',
+            raise_on_error=False,
+            data=dict(database_id=dbid, sql=sql, client_id=client_id),
+        )
+        if raise_on_error and 'error' in resp:
+            raise Exception('validate_sql failed')
+        return resp
+
     @patch.dict('superset._feature_flags', {'FOO': True}, clear=True)
     def test_existing_feature_flags(self):
         self.assertTrue(is_feature_enabled('FOO'))