chore: improve SQL parsing

[betodealmeida]

SUMMARY

This is the first PR in a series to improve our SQL parsing, in effort to clean up the code base and make it more secure.

Currently, we have a class called superset.sql_parse.ParsedQuery that is used as an interface every time we need to parse or manipulate SQL. There are a few problems with the status quo:

1. Even though the class is called ParsedQuery, it's designed to accept a single SQL statement, not necessarily a query. For example, it has a set_or_update_query_limit method that works only on the first statement, if there are multiple.
2. Even though it's designed to accept a single statement, it doesn't raise an error when instantiated with a query that has multiple statements, so in several places the class is used with a full query.
3. There are multiple places that import sqlparse directly to manipulate SQL. For example, when inserting RLS into a query, or when modifying a query to insert a LIMIT or TOP.
4. In many places we manipulate SQL inefficiently, serializing and deserializing the SQL multiple times.

Ideally we'd have a single abstraction that handles SQL parsing. Superset code should never parse code directly by importing a 3rd party library like sqlparse, sqloxide, or sqlglot; instead, it should only call the abstraction. Similarly, the abstraction should wrap library specific exceptions (like sqlglot.errors.ParseError) with its own.

This PR introduces two new classes that aim to be those abstractions: SQLScript and SQLStatement. The SQLStatement strictly represents a single SQL statement. These classes are instantiated with SQL and an optional engine (a BaseEngineSpec.engine attribute), and should provide all the methods needed for manipulating and introspecting SQL queries and statements. The initial implementation uses sqlglot, but the clean interface allows the parser to be changed in the future if we want.

My plan for improving the SQL parsing is first to get rid of code outside of sql_parse.py that using sqlparse directly:

• Introduce SQLScript and SQLStatement, and replace any simple calls to sqlparse with calls to the new classes (this PR).
• Implement a custom SQLStatement class for Kusto, since it doesn't use conventional SQL.
• Implement insert_rls_in_predicate inside SQLStatement.
• Implement get_cte_query inside SQLStatement.
• Implement apply_top_to_sql in SQLStatement.

At this point we should have sqlparse being used only inside sql_parse.py by ParsedQuery. The last steps would be:

• Replace all calls to ParsedQuery with calls to the new classes.
• Get rid of the sqlparse dependency.

Note that despite the big changes, no public interfaces will change. Also, we've been gradually changing the parser due to security issues, having introduced sqlglot as a dependency in #26476. Because of this, I thought it wouldn't be necessary to write a SIP. I'm happy to write one if people think it's necessary.

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

N/A

TESTING INSTRUCTIONS

Updated unit tests, with mostly cosmetic changes because the SQL pretty-printing has changed.

ADDITIONAL INFORMATION

• Has associated issue:
• Required feature flags:
• Changes UI
• Includes DB Migration (follow approval process in SIP-59)
  • Migration is atomic, supports rollback & is backwards-compatible
  • Confirm DB migration upgrade and downgrade tested
  • Runtime estimates and downtime expectations provided
• Introduces new feature or API
• Removes existing feature or API

[betodealmeida]

This function was originally a method in ParsedQuery, I just converted it into a function without any big changes (see below).

[codecov]

Codecov Report

Attention: Patch coverage is 96.25000% with 3 lines in your changes are missing coverage. Please review.

Project coverage is 69.69%. Comparing base (735b895) to head (d1a5b7e).

Files	Patch %	Lines
superset/sql_parse.py	96.42%	2 Missing ⚠️
superset/models/helpers.py	88.88%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #26767      +/-   ##
==========================================
- Coverage   69.73%   69.69%   -0.04%     
==========================================
  Files        1909     1909              
  Lines       74692    74734      +42     
  Branches     8325     8325              
==========================================
  Hits        52086    52086              
- Misses      20556    20598      +42     
  Partials     2050     2050              

Flag	Coverage Δ
hive	?
javascript	57.22% <ø> (ø)
mysql	78.03% <76.25%> (+0.01%)	⬆️
postgres	78.15% <86.25%> (+0.03%)	⬆️
presto	53.69% <67.50%> (+0.03%)	⬆️
python	83.05% <96.25%> (-0.10%)	⬇️
sqlite	77.58% <76.25%> (+0.01%)	⬆️
unit	56.67% <82.50%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mistercrunch]

Few comments:

• Let's remove sqlparse from setup.py and requirements. I noticed pip-compile-multi was SUPER slow last time I tried it
• about the description and now improved support for multi-statements, are all the cases covered by new tests in sql_parse_tests.py?

[villebro]

I think this is a welcome change and a move in the right direction, but I feel a big change like this should probably be discussed in the form of a SIP first. I anticipate the final implementation to end up looking similar to this PR, but it may be beneficial to discuss the change first before jumping directly into the code.

[betodealmeida]

I think this is a welcome change and a move in the right direction, but I feel a big change like this should probably be discussed in the form of a SIP first. I anticipate the final implementation to end up looking similar to this PR, but it may be beneficial to discuss the change first before jumping directly into the code.

SIP here: #26786

[villebro]

Thanks @betodealmeida !

[betodealmeida]

#26786 approved!

[mistercrunch]

Left some comments. At a high level the main thing for me would be to clarify the object model and inheritance scheme for the SQL-related utilities. At a high level thinking about this, it sits close to db_engine_spec or at least relate to it in some ways.

[mistercrunch]

It could be good to rename this module to a better name, the current name can be confused with having a relationship with sqlparse, and overall just isn't a good name. It could be utils/sql.py, or more directly superset/sql/* if we need to grow this into a package with multiple modules. superset/sql_parser.py (?)

[betodealmeida]

Ah, I like this!

@john-bodley, you did some work on reorganizing the codebase, what are your thought here?

[mistercrunch]

[suggesting] it could be good to have some sort of reusable compare_sql(strict=False, case_sensitive=False, disregard_schema_prefix=True) function that could be reused in unit tests. Maybe it's a method of ParseQuery (is_identical or is_similar)

[betodealmeida]

Yeah, good point, I'll add that!

[rafehi]

Where are we at with this work? There's some additional changes that I'd be willing to take on (#19572) but looks like this is a dependency.

We've made and validated significant performance improvements in our Superset fork, but makes sense to consolidate those into SQLScript and SQLStatement rather than directly replace sqlparse calls with sqlglot.

[mistercrunch]

LGTM

[eschutho]

I reviewed just the cypress test portion as a code-owner requirement. That LGTM. 👍

[mistercrunch]

Hey @betodealmeida , this PR might have broken a Hive-specific test -> https://github.com/apache/superset/actions/runs/8272398041/job/22634142747#step:11:608

It appears to fail after this PR was merged (I noticed since it was failing on some other unrelated PR) and trace it back here. Note that this particular GHA test for Hive runs only when superset/** is altered by a PR, that explains why other builds succeeded on master since then.

Hive seems to be complaining about a GROUP BY 1 but from memory it totally supports that. Unclear what's up with the test but I thought you might know more given this PR.