1 files changed, 27 insertions, 1 deletions
diff --git a/tests/test_parser.py b/tests/test_parser.py
index f3e663e..6611b87 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -17,6 +17,13 @@ class TestParser(unittest.TestCase):
         self.assertIsInstance(parse_one("int", into=exp.DataType), exp.DataType)
         self.assertIsInstance(parse_one("array<int>", into=exp.DataType), exp.DataType)
         self.assertIsInstance(parse_one("foo", into=exp.Table), exp.Table)
+        self.assertIsInstance(
+            parse_one(
+                "WHEN MATCHED THEN UPDATE SET target.salary = COALESCE(source.salary, target.salary)",
+                into=exp.When,
+            ),
+            exp.When,
+        )
 
         with self.assertRaises(ParseError) as ctx:
             parse_one("SELECT * FROM tbl", into=exp.Table)
@@ -94,12 +101,31 @@ class TestParser(unittest.TestCase):
         tables = [t.sql() for t in parse_one("select * from a, b.c, .d").find_all(exp.Table)]
         self.assertEqual(set(tables), {"a", "b.c", "d"})
 
-    def test_union_order(self):
+    def test_union(self):
         self.assertIsInstance(parse_one("SELECT * FROM (SELECT 1) UNION SELECT 2"), exp.Union)
         self.assertIsInstance(
             parse_one("SELECT x FROM y HAVING x > (SELECT 1) UNION SELECT 2"), exp.Union
         )
 
+        # Check that modifiers are attached to the topmost union node and not the rightmost query
+        single_union = "SELECT x FROM t1 UNION ALL SELECT x FROM t2 LIMIT 1"
+        expr = parse_one(single_union)
+        limit = expr.assert_is(exp.Union).args.get("limit")
+        self.assertIsInstance(limit, exp.Limit)
+        self.assertEqual(expr.sql(), single_union)
+
+        two_unions = (
+            "SELECT x FROM t1 UNION ALL SELECT x FROM t2 UNION ALL SELECT x FROM t3 LIMIT 1"
+        )
+        expr = parse_one(two_unions)
+        limit = expr.assert_is(exp.Union).args.get("limit")
+        self.assertIsInstance(limit, exp.Limit)
+        self.assertEqual(expr.sql(), two_unions)
+
+        expr = parse_one(single_union, read="clickhouse")
+        self.assertIsNone(expr.args.get("limit"))
+        self.assertEqual(expr.sql(dialect="clickhouse"), single_union)
+
     def test_select(self):
         self.assertIsNotNone(parse_one("select 1 natural"))
         self.assertIsNotNone(parse_one("select * from (select 1) x order by x.y").args["order"])