From 02b2fe5ec12f66399b140738de49796f9b1d4ba2 Mon Sep 17 00:00:00 2001
From: Ronan Dunklau <ronan.dunklau@people-doc.com>
Date: Thu, 4 Apr 2019 14:46:08 +0200
Subject: [PATCH] Support ORDER / LIMIT / OFFSET for set operations

---
 pglast/printers/dml.py                       | 32 ++++++++++----------
 tests/test_printers_roundtrip/dml/select.sql | 11 +++++++
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/pglast/printers/dml.py b/pglast/printers/dml.py
index 2119e505..e7e19a35 100644
--- a/pglast/printers/dml.py
+++ b/pglast/printers/dml.py
@@ -759,22 +759,22 @@ def select_stmt(node, output):
                 output.newline()
                 output.write('WINDOW ')
                 output.print_list(node.windowClause)
-            if node.sortClause:
-                output.newline()
-                output.write('ORDER BY ')
-                output.print_list(node.sortClause)
-            if node.limitCount:
-                output.newline()
-                output.write('LIMIT ')
-                output.print_node(node.limitCount)
-            if node.limitOffset:
-                output.newline()
-                output.write('OFFSET ')
-                output.print_node(node.limitOffset)
-            if node.lockingClause:
-                output.newline()
-                output.write('FOR ')
-                output.print_list(node.lockingClause)
+        if node.sortClause:
+            output.newline()
+            output.write('ORDER BY ')
+            output.print_list(node.sortClause)
+        if node.limitCount:
+            output.newline()
+            output.write('LIMIT ')
+            output.print_node(node.limitCount)
+        if node.limitOffset:
+            output.newline()
+            output.write('OFFSET ')
+            output.print_node(node.limitOffset)
+        if node.lockingClause:
+            output.newline()
+            output.write('FOR ')
+            output.print_list(node.lockingClause)
 
         if node.withClause:
             output.dedent()
diff --git a/tests/test_printers_roundtrip/dml/select.sql b/tests/test_printers_roundtrip/dml/select.sql
index 23cfb459..4e9e70ef 100644
--- a/tests/test_printers_roundtrip/dml/select.sql
+++ b/tests/test_printers_roundtrip/dml/select.sql
@@ -302,6 +302,17 @@ SELECT * FROM "T"
 UNION ALL
 SELECT * FROM "T"
 
+SELECT * FROM "T"
+UNION ALL
+SELECT * FROM "T"
+ORDER BY 1, 2, 3
+
+SELECT * FROM "T"
+UNION ALL
+SELECT * FROM "T"
+ORDER BY 1, 2, 3
+LIMIT 2 OFFSET 1
+
 WITH RECURSIVE employee_recursive("Distance", employee_name, manager_name) AS (
     SELECT 1, employee_name, manager_name
     FROM employee