mirror of https://github.com/apache/superset.git
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
This commit is contained in:
parent
e1bb7f4364
commit
9568985b7b
|
@ -153,6 +153,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
|
||||||
allows_joins = True
|
allows_joins = True
|
||||||
allows_subqueries = True
|
allows_subqueries = True
|
||||||
allows_column_aliases = True
|
allows_column_aliases = True
|
||||||
|
allows_sql_comments = True
|
||||||
force_column_alias_quotes = False
|
force_column_alias_quotes = False
|
||||||
arraysize = 0
|
arraysize = 0
|
||||||
max_column_name_length = 0
|
max_column_name_length = 0
|
||||||
|
@ -857,7 +858,6 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
|
||||||
"""
|
"""
|
||||||
parsed_query = ParsedQuery(statement)
|
parsed_query = ParsedQuery(statement)
|
||||||
sql = parsed_query.stripped()
|
sql = parsed_query.stripped()
|
||||||
|
|
||||||
sql_query_mutator = config["SQL_QUERY_MUTATOR"]
|
sql_query_mutator = config["SQL_QUERY_MUTATOR"]
|
||||||
if sql_query_mutator:
|
if sql_query_mutator:
|
||||||
sql = sql_query_mutator(sql, user_name, security_manager, database)
|
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()
|
:param kwargs: kwargs to be passed to cursor.execute()
|
||||||
:return:
|
:return:
|
||||||
"""
|
"""
|
||||||
|
if not cls.allows_sql_comments:
|
||||||
|
query = sql_parse.strip_comments_from_sql(query)
|
||||||
|
|
||||||
if cls.arraysize:
|
if cls.arraysize:
|
||||||
cursor.arraysize = cls.arraysize
|
cursor.arraysize = cls.arraysize
|
||||||
try:
|
try:
|
||||||
|
|
|
@ -33,6 +33,7 @@ class ElasticSearchEngineSpec(BaseEngineSpec): # pylint: disable=abstract-metho
|
||||||
time_secondary_columns = True
|
time_secondary_columns = True
|
||||||
allows_joins = False
|
allows_joins = False
|
||||||
allows_subqueries = True
|
allows_subqueries = True
|
||||||
|
allows_sql_comments = False
|
||||||
|
|
||||||
_time_grain_expressions = {
|
_time_grain_expressions = {
|
||||||
None: "{col}",
|
None: "{col}",
|
||||||
|
@ -69,6 +70,7 @@ class OpenDistroEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method
|
||||||
time_secondary_columns = True
|
time_secondary_columns = True
|
||||||
allows_joins = False
|
allows_joins = False
|
||||||
allows_subqueries = True
|
allows_subqueries = True
|
||||||
|
allows_sql_comments = False
|
||||||
|
|
||||||
_time_grain_expressions = {
|
_time_grain_expressions = {
|
||||||
None: "{col}",
|
None: "{col}",
|
||||||
|
|
|
@ -58,6 +58,19 @@ def _extract_limit_from_query(statement: TokenList) -> Optional[int]:
|
||||||
return None
|
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)
|
@dataclass(eq=True, frozen=True)
|
||||||
class Table: # pylint: disable=too-few-public-methods
|
class Table: # pylint: disable=too-few-public-methods
|
||||||
"""
|
"""
|
||||||
|
@ -150,6 +163,9 @@ class ParsedQuery:
|
||||||
def stripped(self) -> str:
|
def stripped(self) -> str:
|
||||||
return self.sql.strip(" \t\n;")
|
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]:
|
def get_statements(self) -> List[str]:
|
||||||
"""Returns a list of SQL statements as strings, stripped"""
|
"""Returns a list of SQL statements as strings, stripped"""
|
||||||
statements = []
|
statements = []
|
||||||
|
|
|
@ -14,6 +14,8 @@
|
||||||
# KIND, either express or implied. See the License for the
|
# KIND, either express or implied. See the License for the
|
||||||
# specific language governing permissions and limitations
|
# specific language governing permissions and limitations
|
||||||
# under the License.
|
# under the License.
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
from sqlalchemy import column
|
from sqlalchemy import column
|
||||||
|
|
||||||
from superset.db_engine_specs.elasticsearch import (
|
from superset.db_engine_specs.elasticsearch import (
|
||||||
|
@ -54,3 +56,15 @@ class TestElasticSearchDbEngineSpec(TestDbEngineSpec):
|
||||||
for original, expected in test_cases.items():
|
for original, expected in test_cases.items():
|
||||||
actual = OpenDistroEngineSpec.make_label_compatible(column(original).name)
|
actual = OpenDistroEngineSpec.make_label_compatible(column(original).name)
|
||||||
self.assertEqual(actual, expected)
|
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")
|
||||||
|
|
|
@ -18,7 +18,7 @@ import unittest
|
||||||
|
|
||||||
import sqlparse
|
import sqlparse
|
||||||
|
|
||||||
from superset.sql_parse import ParsedQuery, Table
|
from superset.sql_parse import ParsedQuery, strip_comments_from_sql, Table
|
||||||
|
|
||||||
|
|
||||||
class TestSupersetSqlParse(unittest.TestCase):
|
class TestSupersetSqlParse(unittest.TestCase):
|
||||||
|
@ -732,3 +732,19 @@ class TestSupersetSqlParse(unittest.TestCase):
|
||||||
"""
|
"""
|
||||||
parsed = ParsedQuery(query, strip_comments=True)
|
parsed = ParsedQuery(query, strip_comments=True)
|
||||||
assert not parsed.is_valid_ctas()
|
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"
|
||||||
|
)
|
||||||
|
|
Loading…
Reference in New Issue