Support Clause groups and statement separators

These allow for unions to be parsed correctly and for queries to be
written with surrounding parenthesis.

Author: pgjones

src/sql_tstring/__init__.py

@@ -9,6 +9,7 @@
 
 from sql_tstring.parser import (
     Clause,
+    ClauseGroup,
     Element,
     Expression,
     ExpressionGroup,
@@ -102,10 +103,10 @@ def sql(
     for raw_parsed_query in parsed_queries:
         parsed_query = deepcopy(raw_parsed_query)
         new_values = _replace_placeholders(parsed_query, 0)
-        result_str += _print_node(parsed_query, [None] * len(result_values), ctx.dialect)
+        result_str += " " + _print_node(parsed_query, [None] * len(result_values), ctx.dialect)
         result_values.extend(new_values)
 
-    return result_str, result_values
+    return result_str.strip(), result_values
 
 
 class _ContextManager:
@@ -157,7 +158,12 @@ def _print_node(
 
     match node:
         case Statement():
-            result = " ".join(_print_node(clause, placeholders, dialect) for clause in node.clauses)
+            addition = " ".join(
+                _print_node(clause, placeholders, dialect) for clause in node.clauses
+            )
+            result = f"{node.separator} {addition}"
+        case ClauseGroup():
+            result = f"({" ".join(_print_node(clause, placeholders, dialect) for clause in node.clauses)})"  # noqa: E501
         case Clause() | ExpressionGroup():
             result = ""
 

src/sql_tstring/parser.py

@@ -274,13 +274,14 @@ class ClauseProperties:
 
 @dataclass
 class Statement:
-    clauses: list[Clause] = field(default_factory=list)
+    clauses: list[Clause | ClauseGroup] = field(default_factory=list)
     parent: ExpressionGroup | Function | Group | None = None
+    separator: str = ""
 
 
 @dataclass
 class Clause:
-    parent: Statement
+    parent: ClauseGroup | Statement
     properties: ClauseProperties
     text: str
     expressions: list[Expression] = field(init=False)
@@ -290,6 +291,12 @@ def __post_init__(self) -> None:
         self.expressions = [Expression(self)]
 
 
+@dataclass
+class ClauseGroup:
+    parent: Statement
+    clauses: list[Clause] = field(default_factory=list)
+
+
 @dataclass
 class Expression:
     parent: Clause | ExpressionGroup
@@ -351,7 +358,7 @@ class Operator:
 
 
 type ParentNode = Clause | Expression | ExpressionGroup | Function | Group
-type Node = ParentNode | Literal | Statement
+type Node = ParentNode | ClauseGroup | Literal | Statement
 type Element = Node | Operator | Part | Placeholder
 
 
@@ -387,7 +394,7 @@ def _parse_placeholder(
 ) -> None:
     if isinstance(current_node, (Expression, Function, Group, Literal)):
         parent = current_node
-    elif isinstance(current_node, Statement):
+    elif isinstance(current_node, (Statement, ClauseGroup)):
         raise ValueError("Invalid syntax")
     else:  # Clause | ExpressionGroup
         parent = current_node.expressions[-1]
@@ -461,7 +468,7 @@ def _parse_string(
                 current_node, consumed = _parse_token(
                     current_node, raw_current_token, current_token, tokens[index:], statements
                 )
-        else:  # Statement
+        else:  # ClauseGroup | Statement
             current_node, consumed = _parse_token(
                 current_node, raw_current_token, current_token, tokens[index:], statements
             )
@@ -472,7 +479,7 @@ def _parse_string(
 
 
 def _parse_token(
-    current_node: ParentNode | Statement,
+    current_node: ParentNode | ClauseGroup | Statement,
     raw_current_token: str,
     current_token: str,
     tokens: list[str],
@@ -482,8 +489,17 @@ def _parse_token(
         return _parse_clause(current_node, tokens)
     elif current_token == ";":
         statements.append(Statement())
+        statements[-1].separator = ";"
         return statements[-1], 1
-    elif not isinstance(current_node, Statement):
+    elif current_token == "union":
+        statements.append(Statement())
+        statements[-1].separator = raw_current_token
+        consumed = 1
+        if tokens[1].lower() == "all":
+            statements[-1].separator += f" {tokens[1]}"
+            consumed = 2
+        return statements[-1], consumed
+    elif not isinstance(current_node, (ClauseGroup, Statement)):
         if current_token in OPERATORS:
             return _parse_operator(current_node, tokens)
         elif current_token == "'":
@@ -494,17 +510,21 @@ def _parse_token(
             return _parse_function(current_node, raw_current_token[:-1])
         elif current_token == ")":
             current_node = _find_node(  # type: ignore[assignment]
-                current_node, (ExpressionGroup, Function, Group)
+                current_node, (ExpressionGroup, Function, Group, ClauseGroup)
             )
             return current_node.parent, 1
         else:
             return _parse_part(current_node, raw_current_token)
+    elif isinstance(current_node, Statement) and current_token == "(":
+        statement_group = ClauseGroup(current_node)
+        current_node.clauses.append(statement_group)
+        return statement_group, 1
     else:
         raise ValueError("Invalid syntax")
 
 
 def _parse_clause(
-    current_node: ParentNode | Statement,
+    current_node: ParentNode | ClauseGroup | Statement,
     tokens: list[str],
 ) -> tuple[Clause, int]:
     index = 0
@@ -523,16 +543,16 @@ def _parse_clause(
         statement = Statement(parent=current_node)
         current_node.expressions[-1].parts.append(statement)
         current_node = statement
-    else:  # Clause | Expression | Statement
-        current_node = _find_node(current_node, Statement)
+    else:  # Clause | Expression | Statement | ClauseGroup
+        current_node = _find_node(current_node, (Statement, ClauseGroup))  # type: ignore[assignment] # noqa: E501
 
     clause_properties = cast(ClauseProperties, clause_entry[""])
     clause = Clause(
-        parent=current_node,
+        parent=current_node,  # type: ignore[arg-type]
         properties=clause_properties,
         text=text,
     )
-    current_node.clauses.append(clause)
+    current_node.clauses.append(clause)  # type: ignore[union-attr]
     return clause, index
 
 

tests/test_parsing.py

@@ -105,3 +105,8 @@ def test_nested() -> None:
     inner = t("x = 'a'", locals())
     query, _ = sql("SELECT x FROM y WHERE {inner}", locals())
     assert query == "SELECT x FROM y WHERE x = 'a'"
+
+
+def test_opening_parenthesis() -> None:
+    query, _ = sql("(SELECT x FROM y) UNION ALL (SELECT z FROM y)", locals())
+    assert query == "(SELECT x FROM y) UNION ALL (SELECT z FROM y)"