diff --git a/superset/sql/parse.py b/superset/sql/parse.py index ec46b5850dc..f5f93d34d06 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -32,9 +32,11 @@ from sqlglot import exp from sqlglot.dialects.dialect import ( Dialect, Dialects, + DialectType, ) from sqlglot.dialects.singlestore import SingleStore from sqlglot.errors import ParseError +from sqlglot.generator import Generator from sqlglot.optimizer.pushdown_predicates import ( pushdown_predicates, ) @@ -137,6 +139,32 @@ class CTASMethod(enum.Enum): VIEW = enum.auto() +def _normalized_generator( + dialect_name: DialectType, + *, + pretty: bool, + comments: bool, +) -> Generator: + """ + Generator that preserves multi-argument DISTINCT expressions. + + Build a sqlglot generator that preserves user-written multi-argument + DISTINCT expressions verbatim. Postgres, Presto, Trino, and DuckDB + set ``MULTI_ARG_DISTINCT = False`` to emulate the unsupported + ``COUNT(DISTINCT a, b)`` idiom via a ``CASE WHEN`` row-expression, which + silently corrupts user-defined aggregates that natively accept multiple + arguments. Superset's sanitize / format paths normalize user SQL — they + do not transpile — so the emulation is undesirable here. + """ + dialect = Dialect.get_or_raise(dialect_name) + normalized_cls = type( + f"Normalized{dialect.generator_class.__name__}", + (dialect.generator_class,), + {"MULTI_ARG_DISTINCT": True}, + ) + return normalized_cls(dialect=dialect, pretty=pretty, comments=comments) + + class RLSMethod(enum.Enum): """ Methods for enforcing RLS. @@ -911,12 +939,11 @@ class SQLStatement(BaseSQLStatement[exp.Expression]): """ Pretty-format the SQL statement. """ - return Dialect.get_or_raise(self._dialect).generate( - self._parsed, - copy=True, - comments=comments, + return _normalized_generator( + self._dialect, pretty=True, - ) + comments=comments, + ).generate(self._parsed, copy=True) def get_settings(self) -> dict[str, str | bool]: """ @@ -1817,14 +1844,13 @@ def sanitize_clause(clause: str, engine: str) -> str: """ try: statement = SQLStatement(clause, engine) - dialect = SQLGLOT_DIALECTS.get(engine) - from sqlglot.dialects.dialect import Dialect - - return Dialect.get_or_raise(dialect).generate( + return _normalized_generator( + SQLGLOT_DIALECTS.get(engine), + pretty=False, + comments=True, + ).generate( statement._parsed, # pylint: disable=protected-access copy=True, - comments=True, - pretty=False, ) except SupersetParseError as ex: raise QueryClauseValidationException(f"Invalid SQL clause: {clause}") from ex diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 6d58550f028..2f3b72c5b38 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -3153,6 +3153,48 @@ def test_is_valid_cvas(sql: str, engine: str, expected: bool) -> None: ("col1 = 1) AND (col2 = 2", QueryClauseValidationException, "base"), ("(col1 = 1)) AND ((col2 = 2)", QueryClauseValidationException, "base"), ("TRUE; SELECT 1", QueryClauseValidationException, "base"), + # Regression test for https://github.com/apache/superset/issues/39223: + # dialects with `MULTI_ARG_DISTINCT=False` (Postgres, Presto, Trino, + # DuckDB) must not rewrite user-defined multi-argument DISTINCT + # aggregates into row-expression null guards. Dremio is included + # below as a defensive regression guard even though its generator + # does not currently set `MULTI_ARG_DISTINCT=False`. + ( + "DISTINCT_AVG(DISTINCT report_id, time_to_accept / 86400)", + "DISTINCT_AVG(DISTINCT report_id, time_to_accept / 86400)", + "postgresql", + ), + ( + "DISTINCT_SUM(DISTINCT report_id, total_bounty_reward_amount)", + "DISTINCT_SUM(DISTINCT report_id, total_bounty_reward_amount)", + "postgresql", + ), + ( + "DISTINCT_AVG(DISTINCT k, v)", + "DISTINCT_AVG(DISTINCT k, v)", + "presto", + ), + ( + "DISTINCT_AVG(DISTINCT k, v)", + "DISTINCT_AVG(DISTINCT k, v)", + "trino", + ), + ( + "DISTINCT_AVG(DISTINCT k, v)", + "DISTINCT_AVG(DISTINCT k, v)", + "duckdb", + ), + ( + "DISTINCT_AVG(DISTINCT k, v)", + "DISTINCT_AVG(DISTINCT k, v)", + "dremio", + ), + # Single-argument DISTINCT must still round-trip cleanly. + ( + "COUNT(DISTINCT x)", + "COUNT(DISTINCT x)", + "postgresql", + ), ], ) def test_sanitize_clause(sql: str, expected: str | Exception, engine: str) -> None: @@ -3166,6 +3208,30 @@ def test_sanitize_clause(sql: str, expected: str | Exception, engine: str) -> No sanitize_clause(sql, engine) +@pytest.mark.parametrize( + "engine", + [ + "postgresql", + "presto", + "trino", + "duckdb", + "dremio", + ], +) +def test_sqlstatement_format_preserves_multi_arg_distinct(engine: str) -> None: + """ + Regression guard for https://github.com/apache/superset/issues/39223: + ``SQLStatement.format()`` must not rewrite user-defined multi-argument + DISTINCT aggregates into row-expression null guards. This is the SQL Lab / + executor path; the metric-expression path is covered by + ``test_sanitize_clause``. + """ + sql = "SELECT DISTINCT_AVG(DISTINCT a, b) FROM t" + formatted = SQLScript(sql, engine).format() + assert "DISTINCT_AVG(DISTINCT a, b)" in formatted + assert "CASE WHEN" not in formatted + + @pytest.mark.parametrize( "engine", [