diff --git a/superset/sql/dialects/pinot.py b/superset/sql/dialects/pinot.py index 2e7cbe9ed5c..417d5fd1110 100644 --- a/superset/sql/dialects/pinot.py +++ b/superset/sql/dialects/pinot.py @@ -59,6 +59,11 @@ class Pinot(MySQL): expression=seq_get(args, 1), unit=seq_get(args, 0), ), + "DATE_SUB": lambda args: exp.DateSub( + this=seq_get(args, 2), + expression=seq_get(args, 1), + unit=seq_get(args, 0), + ), } class Generator(MySQL.Generator): @@ -97,6 +102,12 @@ class Pinot(MySQL): e.args.get("expression"), e.this, ), + exp.DateSub: lambda self, e: self.func( + "DATE_SUB", + exp.Literal.string(str(e.args.get("unit").name)), + e.args.get("expression"), + e.this, + ), } # Remove DATE_TRUNC transformation - Pinot supports standard SQL DATE_TRUNC TRANSFORMS.pop(exp.DateTrunc, None) diff --git a/tests/unit_tests/sql/dialects/pinot_tests.py b/tests/unit_tests/sql/dialects/pinot_tests.py index bd2c0003325..88dcfb93844 100644 --- a/tests/unit_tests/sql/dialects/pinot_tests.py +++ b/tests/unit_tests/sql/dialects/pinot_tests.py @@ -546,3 +546,47 @@ def test_pinot_date_add_unit_quoted() -> None: # The unit should be quoted: 'DAY' not DAY assert "DATE_ADD('DAY', -180, NOW())" in result assert "DATE_ADD(DAY," not in result + + +def test_pinot_date_sub_parsing() -> None: + """ + Test that Pinot's DATE_SUB function with Presto-like syntax can be parsed. + """ + from superset.sql.parse import SQLScript + + sql = "SELECT * FROM my_table WHERE dt >= date_sub('day', 7, now())" + script = SQLScript(sql, "pinot") + assert len(script.statements) == 1 + assert not script.has_mutation() + + +def test_pinot_date_sub_simple() -> None: + """ + Test parsing of simple DATE_SUB expressions. + """ + test_cases = [ + "date_sub('day', 7, now())", + "DATE_SUB('month', 3, current_timestamp())", + "date_sub('hour', 24, my_date_column)", + ] + + for sql in test_cases: + parsed = sqlglot.parse_one(sql, Pinot) + assert parsed is not None + # Verify that it generates valid SQL + generated = parsed.sql(dialect=Pinot) + assert "DATE_SUB" in generated.upper() + + +def test_pinot_date_sub_unit_quoted() -> None: + """ + Test that DATE_SUB preserves quotes around the unit argument. + + Pinot requires the unit to be a quoted string, not an identifier. + """ + sql = "dt_epoch_ms >= date_sub('day', -180, now())" + result = sqlglot.parse_one(sql, Pinot).sql(Pinot) + + # The unit should be quoted: 'DAY' not DAY + assert "DATE_SUB('DAY', -180, NOW())" in result + assert "DATE_SUB(DAY," not in result