[Refactor] UNION support on Query Parser and Formatter

[HazelYuAhiru]

Overview

Add UNION and UNION ALL support to the query parser and formatter.

Code Changes

• Added COMPOUND_QUERY to NodeType enum in core/ast/enums.py.
• Added CompoundQueryNode to core/ast/node.py. Represents a binary UNION / UNION ALL as a left-associative tree; chains of the same operator (e.g. A UNION B UNION C) are flattened back to flat lists by the formatter to match mo_sql_parsing's wire format.
• Updated core/query_parser.py: added a few helper functions to detect and parse set-operator dicts from mo_sql_parsing. Refactored _build_from_source to handle UNION subqueries in FROM/JOIN sources uniformly.
• Updated core/query_formatter.py: added _collect_union_branches and compound_to_mosql_json to convert the binary tree back to mo_sql_parsing's flat-list JSON. Added a special case in format_from for bare (unaliased) UNION subqueries, which mo_sql_parsing represents as a plain dict rather than a wrapped list entry.

[Copilot]

Pull request overview

Adds AST + parser/formatter support for SQL UNION / UNION ALL by introducing a compound-query node and round-tripping it through mo_sql_parsing’s JSON format, along with enabling previously-skipped UNION-related test cases.

Changes:

• Added NodeType.COMPOUND_QUERY and a new CompoundQueryNode to represent UNION/UNION ALL as a left-associative binary tree.
• Updated QueryParser to detect/parse union / union_all dicts (including in FROM/JOIN sources) and return a Node (query or compound query) at the top level.
• Updated QueryFormatter to convert CompoundQueryNode back into mo_sql_parsing’s union JSON (including a special-case for unaliased UNION subqueries in FROM) and enabled UNION test cases + expected ASTs.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
core/ast/enums.py	Adds a new NodeType for compound queries.
core/ast/node.py	Introduces CompoundQueryNode to model UNION/UNION ALL.
core/query_parser.py	Adds UNION detection + parsing and refactors FROM/JOIN source handling to support UNION subqueries.
core/query_formatter.py	Adds compound-query JSON formatting and a FROM-formatting special case for unaliased UNION subqueries.
data/asts.py	Adds expected ASTs for UNION-related queries (15/29/32) and widens get_ast return type.
tests/ast_util.py	Extends AST visualization to handle compound queries; widens type annotations to Node.
tests/test_query_parser.py	Enables parser tests for UNION queries.
tests/test_query_formatter.py	Enables formatter tests for UNION queries.
tests/test_parser_formatter_e2e.py	Enables parse→format→parse round-trip tests for UNION queries.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[colinthebomb1]

LGTM!

[baiqiushi]

LGTM. Thanks!