chore: improve SQL parsing (#26767)

tests/unit_tests/sql_parse_tests.py

@@ -35,6 +35,8 @@ from superset.sql_parse import (
     insert_rls_in_predicate,
     ParsedQuery,
     sanitize_clause,
+    SQLScript,
+    SQLStatement,
     strip_comments_from_sql,
     Table,
 )
@@ -1850,3 +1852,36 @@ WITH t AS (
 )
 SELECT * FROM t"""
     ).is_select()
+
+
+def test_sqlquery() -> None:
+    """
+    Test the `SQLScript` class.
+    """
+    script = SQLScript("SELECT 1; SELECT 2;")
+
+    assert len(script.statements) == 2
+    assert script.format() == "SELECT\n  1;\nSELECT\n  2"
+    assert script.statements[0].format() == "SELECT\n  1"
+
+    script = SQLScript("SET a=1; SET a=2; SELECT 3;")
+    assert script.get_settings() == {"a": "2"}
+
+
+def test_sqlstatement() -> None:
+    """
+    Test the `SQLStatement` class.
+    """
+    statement = SQLStatement("SELECT * FROM table1 UNION ALL SELECT * FROM table2")
+
+    assert statement.tables == {
+        Table(table="table1", schema=None, catalog=None),
+        Table(table="table2", schema=None, catalog=None),
+    }
+    assert (
+        statement.format()
+        == "SELECT\n  *\nFROM table1\nUNION ALL\nSELECT\n  *\nFROM table2"
+    )
+
+    statement = SQLStatement("SET a=1")
+    assert statement.get_settings() == {"a": "1"}