replace sqlparse.format with sqlglot.transpile

Author: dvchristianbors

superset/common/query_object.py

@@ -276,7 +276,7 @@ def validate(
             # performance for WHERE ... IN (...) clauses
             # Clauses are anyway checked for their validity in
             # e.g., connectors/sqla/models/get_query_str_extended
-            # self._sanitize_filters()
+            self._sanitize_filters()
             return None
         except QueryObjectValidationError as ex:
             if raise_exceptions:

superset/config.py

@@ -1900,7 +1900,7 @@ class ExtraDynamicQueryFilters(TypedDict, total=False):
 elif importlib.util.find_spec("superset_config") and not is_test():
     try:
         # pylint: disable=import-error,wildcard-import,unused-wildcard-import
-        import superset_config
+        import superset_config as superset_config
         from superset_config import *  # noqa: F403, F401
 
         click.secho(

superset/db_engine_specs/base.py

@@ -38,7 +38,7 @@
 
 import pandas as pd
 import requests
-import sqlparse
+import sqlglot
 from apispec import APISpec
 from apispec.ext.marshmallow import MarshmallowPlugin
 from deprecation import deprecated
@@ -1236,7 +1236,7 @@ def get_cte_query(cls, sql: str) -> str | None:
 
         """
         if not cls.allows_cte_in_subquery:
-            stmt = sqlparse.parse(sql)[0]
+            stmt = sqlglot.tokenize(sql)
 
             # The first meaningful token for CTE will be with WITH
             idx, token = stmt.token_next(-1, skip_ws=True, skip_cm=True)
@@ -2158,7 +2158,8 @@ def cancel_query(  # pylint: disable=unused-argument
 
     @classmethod
     def parse_sql(cls, sql: str) -> list[str]:
-        return [str(s).strip(" ;") for s in sqlparse.parse(sql)]
+        return sqlglot.transpile(sql)
+        # return [str(s).strip(" ;") for s in sqlparse.parse(sql)]
 
     @classmethod
     def get_impersonation_key(cls, user: User | None) -> Any:

superset/models/core.py

@@ -58,7 +58,6 @@
 from sqlalchemy.pool import NullPool
 from sqlalchemy.schema import UniqueConstraint
 from sqlalchemy.sql import ColumnElement, expression, Select
-from sqlglot import parse
 
 from superset import app, db_engine_specs, is_feature_enabled
 from superset.commands.database.exceptions import DatabaseInvalidError
@@ -654,11 +653,7 @@ def get_df(  # pylint: disable=too-many-locals
         schema: str | None = None,
         mutator: Callable[[pd.DataFrame], None] | None = None,
     ) -> pd.DataFrame:
-        # before we split sqls using sql parse, however this core code is only reachable
-        # with single sql queries. Thus, we remove the engine spec parser here
-        # sqls = self.db_engine_spec.parse_sql(sql)
-        sqls = parse(sql)
-
+        sqls = self.db_engine_spec.parse_sql(sql)
         with self.get_sqla_engine(catalog=catalog, schema=schema) as engine:
             engine_url = engine.url
 

superset/sql_parse.py

@@ -24,11 +24,13 @@
 from collections.abc import Iterator
 from typing import Any, cast, TYPE_CHECKING
 
+import sqlglot
 import sqlparse
 from flask_babel import gettext as __
 from jinja2 import nodes
 from sqlalchemy import and_
 from sqlglot.dialects.dialect import Dialects
+from sqlglot.errors import ParseError
 from sqlparse import keywords
 from sqlparse.lexer import Lexer
 from sqlparse.sql import (
@@ -42,7 +44,6 @@
     Where,
 )
 from sqlparse.tokens import (
-    Comment,
     CTE,
     DDL,
     DML,
@@ -257,6 +258,7 @@ def __init__(
         sql_statement: str,
         engine: str = "base",
     ):
+        sql_statement = sqlglot.transpile(sql_statement)
         self.sql: str = sql_statement
         self._engine = engine
         self._dialect = SQLGLOT_DIALECTS.get(engine) if engine else None
@@ -579,30 +581,35 @@ def set_or_update_query_limit(self, new_limit: int, force: bool = False) -> str:
 
 def sanitize_clause(clause: str) -> str:
     # clause = sqlparse.format(clause, strip_comments=True)
-    statements = sqlparse.parse(clause)
+    try:
+        statements = sqlglot.transpile(clause, pretty=True)
+    except Exception as p_err:
+        if isinstance(p_err, ParseError):
+            raise QueryClauseValidationException(str(p_err)) from p_err
+        raise ValueError(str(p_err)) from None
     if len(statements) != 1:
         raise QueryClauseValidationException("Clause contains multiple statements")
-    open_parens = 0
-
-    previous_token = None
-    for token in statements[0]:
-        if token.value == "/" and previous_token and previous_token.value == "*":
-            raise QueryClauseValidationException("Closing unopened multiline comment")
-        if token.value == "*" and previous_token and previous_token.value == "/":
-            raise QueryClauseValidationException("Unclosed multiline comment")
-        if token.value in (")", "("):
-            open_parens += 1 if token.value == "(" else -1
-            if open_parens < 0:
-                raise QueryClauseValidationException(
-                    "Closing unclosed parenthesis in filter clause"
-                )
-        previous_token = token
-    if open_parens > 0:
-        raise QueryClauseValidationException("Unclosed parenthesis in filter clause")
-
-    if previous_token and previous_token.ttype in Comment:
-        if previous_token.value[-1] != "\n":
-            clause = f"{clause}\n"
+    # open_parens = 0
+
+    # previous_token = None
+    # for token in statements[0]:
+    #     if token.value == "/" and previous_token and previous_token.value == "*":
+    #         raise QueryClauseValidationException("Closing unopened multiline comment")
+    #     if token.value == "*" and previous_token and previous_token.value == "/":
+    #         raise QueryClauseValidationException("Unclosed multiline comment")
+    #     if token.value in (")", "("):
+    #         open_parens += 1 if token.value == "(" else -1
+    #         if open_parens < 0:
+    #             raise QueryClauseValidationException(
+    #                 "Closing unclosed parenthesis in filter clause"
+    #             )
+    #     previous_token = token
+    # if open_parens > 0:
+    #     raise QueryClauseValidationException("Unclosed parenthesis in filter clause")
+
+    # if previous_token and previous_token.ttype in Comment:
+    #     if previous_token.value[-1] != "\n":
+    #         clause = f"{clause}\n"
 
     return clause
 

tests/unit_tests/sql_parse_tests.py

@@ -20,6 +20,7 @@
 from unittest.mock import Mock
 
 import pytest
+import sqlglot
 import sqlparse
 from pytest_mock import MockerFixture
 from sqlalchemy import text
@@ -1168,17 +1169,17 @@ def test_messy_breakdown_statements() -> None:
     ]
 
 
-def test_sqlparse_formatting():
+def test_sqlglot_formatting():
     """
     Test that ``from_unixtime`` is formatted correctly.
     """
-    assert sqlparse.format(
+    assert sqlglot.transpile(
         "SELECT extract(HOUR from from_unixtime(hour_ts) "
         "AT TIME ZONE 'America/Los_Angeles') from table",
-        reindent=True,
-    ) == (
-        "SELECT extract(HOUR\n               from from_unixtime(hour_ts) "
-        "AT TIME ZONE 'America/Los_Angeles')\nfrom table"
+        pretty=True,
+    )[0] == (
+        "SELECT\n  EXTRACT(HOUR FROM FROM_UNIXTIME(hour_ts) AT TIME ZONE 'America/Los_Angeles')"
+        "\nFROM table"
     )