From 9568985b7b92a4ba922c261b370fea1db3a57a3b Mon Sep 17 00:00:00 2001 From: Daniel Vaz Gaspar Date: Wed, 17 Feb 2021 18:01:34 +0000 Subject: [PATCH] fix: engines that don't support comments (#13153) * fix: engines that don't support comments * fix: engines that don't support comments * add quick inexpensive test * add test --- superset/db_engine_specs/base.py | 5 ++++- superset/db_engine_specs/elasticsearch.py | 2 ++ superset/sql_parse.py | 16 ++++++++++++++++ tests/db_engine_specs/elasticsearch_tests.py | 14 ++++++++++++++ tests/sql_parse_tests.py | 18 +++++++++++++++++- 5 files changed, 53 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 1e560e6ed5..1fb9bd5f5e 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -153,6 +153,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods allows_joins = True allows_subqueries = True allows_column_aliases = True + allows_sql_comments = True force_column_alias_quotes = False arraysize = 0 max_column_name_length = 0 @@ -857,7 +858,6 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods """ parsed_query = ParsedQuery(statement) sql = parsed_query.stripped() - sql_query_mutator = config["SQL_QUERY_MUTATOR"] if sql_query_mutator: sql = sql_query_mutator(sql, user_name, security_manager, database) @@ -933,6 +933,9 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods :param kwargs: kwargs to be passed to cursor.execute() :return: """ + if not cls.allows_sql_comments: + query = sql_parse.strip_comments_from_sql(query) + if cls.arraysize: cursor.arraysize = cls.arraysize try: diff --git a/superset/db_engine_specs/elasticsearch.py b/superset/db_engine_specs/elasticsearch.py index dc3651256e..31c15a40e6 100644 --- a/superset/db_engine_specs/elasticsearch.py +++ b/superset/db_engine_specs/elasticsearch.py @@ -33,6 +33,7 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-metho time_secondary_columns = True allows_joins = False allows_subqueries = True + allows_sql_comments = False _time_grain_expressions = { None: "{col}", @@ -69,6 +70,7 @@ class OpenDistroEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method time_secondary_columns = True allows_joins = False allows_subqueries = True + allows_sql_comments = False _time_grain_expressions = { None: "{col}", diff --git a/superset/sql_parse.py b/superset/sql_parse.py index dd345dbd1f..ddf27895dc 100644 --- a/superset/sql_parse.py +++ b/superset/sql_parse.py @@ -58,6 +58,19 @@ def _extract_limit_from_query(statement: TokenList) -> Optional[int]: return None +def strip_comments_from_sql(statement: str) -> str: + """ + Strips comments from a SQL statement, does a simple test first + to avoid always instantiating the expensive ParsedQuery constructor + + This is useful for engines that don't support comments + + :param statement: A string with the SQL statement + :return: SQL statement without comments + """ + return ParsedQuery(statement).strip_comments() if "--" in statement else statement + + @dataclass(eq=True, frozen=True) class Table: # pylint: disable=too-few-public-methods """ @@ -150,6 +163,9 @@ class ParsedQuery: def stripped(self) -> str: return self.sql.strip(" \t\n;") + def strip_comments(self) -> str: + return sqlparse.format(self.stripped(), strip_comments=True) + def get_statements(self) -> List[str]: """Returns a list of SQL statements as strings, stripped""" statements = [] diff --git a/tests/db_engine_specs/elasticsearch_tests.py b/tests/db_engine_specs/elasticsearch_tests.py index 5694e0418f..fd0dcc7e97 100644 --- a/tests/db_engine_specs/elasticsearch_tests.py +++ b/tests/db_engine_specs/elasticsearch_tests.py @@ -14,6 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from unittest.mock import MagicMock + from sqlalchemy import column from superset.db_engine_specs.elasticsearch import ( @@ -54,3 +56,15 @@ class TestElasticSearchDbEngineSpec(TestDbEngineSpec): for original, expected in test_cases.items(): actual = OpenDistroEngineSpec.make_label_compatible(column(original).name) self.assertEqual(actual, expected) + + def test_opendistro_strip_comments(self): + """ + DB Eng Specs (opendistro): Test execute sql strip comments + """ + mock_cursor = MagicMock() + mock_cursor.execute.return_value = [] + + OpenDistroEngineSpec.execute( + mock_cursor, "-- some comment \nSELECT 1\n --other comment" + ) + mock_cursor.execute.assert_called_once_with("SELECT 1\n") diff --git a/tests/sql_parse_tests.py b/tests/sql_parse_tests.py index b54a9ef905..01a54e7309 100644 --- a/tests/sql_parse_tests.py +++ b/tests/sql_parse_tests.py @@ -18,7 +18,7 @@ import unittest import sqlparse -from superset.sql_parse import ParsedQuery, Table +from superset.sql_parse import ParsedQuery, strip_comments_from_sql, Table class TestSupersetSqlParse(unittest.TestCase): @@ -732,3 +732,19 @@ class TestSupersetSqlParse(unittest.TestCase): """ parsed = ParsedQuery(query, strip_comments=True) assert not parsed.is_valid_ctas() + + def test_strip_comments_from_sql(self): + """Test that we are able to strip comments out of SQL stmts""" + + assert ( + strip_comments_from_sql("SELECT col1, col2 FROM table1") + == "SELECT col1, col2 FROM table1" + ) + assert ( + strip_comments_from_sql("SELECT col1, col2 FROM table1\n-- comment") + == "SELECT col1, col2 FROM table1\n" + ) + assert ( + strip_comments_from_sql("SELECT '--abc' as abc, col2 FROM table1\n") + == "SELECT '--abc' as abc, col2 FROM table1" + )