Repository: sqlfluff/sqlfluff

Issue:
Add a fix routine to L036: Select targets should be on a new line unless there is only one select target
The idea arose from this [Slack thread](https://sqlfluff.slack.com/archives/C01GX4L213J/p1613017651043300).

Discussion:

Patch:
diff --git a/src/sqlfluff/core/rules/std/L036.py b/src/sqlfluff/core/rules/std/L036.py
index 4602f21bc..477040d82 100644
--- a/src/sqlfluff/core/rules/std/L036.py
+++ b/src/sqlfluff/core/rules/std/L036.py
@@ -1,10 +1,23 @@
 """Implementation of Rule L036."""
 
-from collections import namedtuple
+from typing import List, NamedTuple
 
-from sqlfluff.core.rules.base import BaseCrawler, LintResult
+from sqlfluff.core.parser import BaseSegment
+from sqlfluff.core.rules.base import BaseCrawler, LintFix, LintResult
+from sqlfluff.core.rules.doc_decorators import document_fix_compatible
 
 
+class SelectTargetsInfo(NamedTuple):
+    """Info about select targets and nearby whitespace."""
+
+    select_idx: int
+    first_new_line_idx: int
+    first_select_target_idx: int
+    first_whitespace_idx: int
+    select_targets: List[BaseSegment]
+
+
+@document_fix_compatible
 class Rule_L036(BaseCrawler):
     """Select targets should be on a new line unless there is only one select target.
 
@@ -31,32 +44,27 @@ class Rule_L036(BaseCrawler):
 
     def _eval(self, segment, raw_stack, **kwargs):
         if segment.is_type("select_clause"):
-            eval_result = self._get_indexes(segment)
-            if eval_result.cnt_select_targets == 1:
-                return self._eval_single_select_target_element(eval_result, segment)
-            if eval_result.cnt_select_targets > 1:
-                return self._eval_multiple_select_target_elements(eval_result, segment)
+            select_targets_info = self._get_indexes(segment)
+            if len(select_targets_info.select_targets) == 1:
+                parent_stack = kwargs.get("parent_stack")
+                return self._eval_single_select_target_element(
+                    select_targets_info, segment, parent_stack
+                )
+            elif len(select_targets_info.select_targets) > 1:
+                return self._eval_multiple_select_target_elements(
+                    select_targets_info, segment
+                )
 
     @staticmethod
     def _get_indexes(segment):
-        EvalResult = namedtuple(
-            "EvalResults",
-            [
-                "cnt_select_targets",
-                "select_idx",
-                "first_new_line_idx",
-                "first_select_target_idx",
-                "first_whitespace_idx",
-            ],
-        )
-        cnt_select_targets = 0
         select_idx = -1
         first_new_line_idx = -1
         first_select_target_idx = -1
         first_whitespace_idx = -1
+        select_targets = []
         for fname_idx, seg in enumerate(segment.segments):
             if seg.is_type("select_target_element"):
-                cnt_select_targets += 1
+                select_targets.append(seg)
                 if first_select_target_idx == -1:
                     first_select_target_idx = fname_idx
             if seg.is_type("keyword") and seg.name == "SELECT" and select_idx == -1:
@@ -73,34 +81,35 @@ class Rule_L036(BaseCrawler):
             ):
                 first_whitespace_idx = fname_idx
 
-        eval_result = EvalResult(
-            cnt_select_targets,
+        return SelectTargetsInfo(
             select_idx,
             first_new_line_idx,
             first_select_target_idx,
             first_whitespace_idx,
+            select_targets,
         )
 
-        return eval_result
-
-    @staticmethod
-    def _eval_multiple_select_target_elements(eval_result, segment):
-        if eval_result.first_new_line_idx == -1:
+    def _eval_multiple_select_target_elements(self, select_targets_info, segment):
+        if select_targets_info.first_new_line_idx == -1:
             # there are multiple select targets but no new lines
-            return LintResult(anchor=segment)
-        else:
-            # ensure newline before select target and whitespace segment
-            if (
-                eval_result.first_new_line_idx
-                < eval_result.first_whitespace_idx
-                < eval_result.first_select_target_idx
-            ):
-                return None
-            else:
-                return LintResult(anchor=segment)
 
-    @staticmethod
-    def _eval_single_select_target_element(eval_result, select_clause):
+            # Find and delete any whitespace between "SELECT" and its targets.
+            ws_to_delete = segment.select_children(
+                start_seg=segment.segments[select_targets_info.select_idx],
+                select_if=lambda s: s.is_type("whitespace"),
+                loop_while=lambda s: s.is_type("whitespace") or s.is_meta,
+            )
+            fixes = [LintFix("delete", ws) for ws in ws_to_delete]
+            # Insert newline before the first select target.
+            ins = self.make_newline(
+                pos_marker=segment.pos_marker.advance_by(segment.raw)
+            )
+            fixes.append(LintFix("create", select_targets_info.select_targets[0], ins))
+            return LintResult(anchor=segment, fixes=fixes)
+
+    def _eval_single_select_target_element(
+        self, select_targets_info, select_clause, parent_stack
+    ):
         is_wildcard = False
         for segment in select_clause.segments:
             if segment.is_type("select_target_element"):
@@ -111,11 +120,55 @@ class Rule_L036(BaseCrawler):
         if is_wildcard:
             return None
         elif (
-            eval_result.select_idx
-            < eval_result.first_new_line_idx
-            < eval_result.first_select_target_idx
+            select_targets_info.select_idx
+            < select_targets_info.first_new_line_idx
+            < select_targets_info.first_select_target_idx
         ):
             # there is a newline between select and select target
-            return LintResult(anchor=select_clause)
+            insert_buff = [
+                self.make_whitespace(
+                    raw=" ",
+                    pos_marker=select_clause.segments[
+                        select_targets_info.first_new_line_idx
+                    ].pos_marker,
+                ),
+                select_clause.segments[select_targets_info.first_select_target_idx],
+                self.make_newline(
+                    pos_marker=select_clause.segments[
+                        select_targets_info.first_new_line_idx
+                    ].pos_marker
+                ),
+            ]
+            fixes = [
+                # Replace "newline" with <<select_target>>, "newline".
+                LintFix(
+                    "edit",
+                    select_clause.segments[select_targets_info.first_new_line_idx],
+                    insert_buff,
+                ),
+                # Delete the first select target from its original location.
+                LintFix(
+                    "delete",
+                    select_clause.segments[select_targets_info.first_select_target_idx],
+                ),
+            ]
+            if parent_stack and parent_stack[-1].type == "select_statement":
+                select_stmt = parent_stack[-1]
+                select_clause_idx = select_stmt.segments.index(select_clause)
+                after_select_clause_idx = select_clause_idx + 1
+                if len(select_stmt.segments) > after_select_clause_idx:
+                    if select_stmt.segments[after_select_clause_idx].type == "newline":
+                        # The select_clause is immediately followed by a
+                        # newline. Delete the newline in order to avoid leaving
+                        # behind an empty line after fix.
+                        fixes.append(
+                            LintFix(
+                                "delete", select_stmt.segments[after_select_clause_idx]
+                            )
+                        )
+            return LintResult(
+                anchor=select_clause,
+                fixes=fixes,
+            )
         else:
             return None


Test patch:
diff --git a/test/api/simple_test.py b/test/api/simple_test.py
index 2bfdba4b2..789b03bab 100644
--- a/test/api/simple_test.py
+++ b/test/api/simple_test.py
@@ -103,7 +103,7 @@ def test__api__fix_string():
     # Check return types.
     assert isinstance(result, str)
     # Check actual result
-    assert result == "SELECT *, 1, blah AS foo FROM mytable\n"
+    assert result == "SELECT\n    *, 1, blah AS foo FROM mytable\n"
 
 
 def test__api__fix_string_specific():
diff --git a/test/core/rules/test_cases/L036.yml b/test/core/rules/test_cases/L036.yml
index 8c77206c3..bfce26f1f 100644
--- a/test/core/rules/test_cases/L036.yml
+++ b/test/core/rules/test_cases/L036.yml
@@ -8,6 +8,17 @@ test_single_select_target_and_new_line_after_select_target:
     select *
     from x
 
+test_single_select_target_and_new_line_before_select_target:
+  fail_str: |
+      select
+          a
+      from x
+  # In a normal "sqlfluff fix" run, the extra spaces in fix_str would be removed
+  # by a different rule, L039.
+  fix_str: |
+    select a
+        from x
+
 test_multiple_select_targets_on_new_lines_and_new_line_after_select:
   pass_str: |
     select
@@ -26,7 +37,23 @@ test_multiple_select_targets_all_on_the_same_line:
   fail_str: |
     select a, b, c
     from x
+  fix_str: "select\na, b, c\nfrom x\n"
 
 test_multiple_select_targets_trailing_whitespace_after_select:
   # TRICKY: Use explicit newlines to preserve the trailing space after "SELECT".
   pass_str: "SELECT \n    a,\n    b\nFROM t\n"
+
+test_comment_between_select_and_single_select_target:
+  fail_str: |
+    SELECT
+        -- This is the user's ID.
+        user_id
+    FROM
+        safe_user
+  # In a normal "sqlfluff fix" run, the weird "FROM" indenting left over after the fix
+  # would be fixed by L003.
+  fix_str: |
+    SELECT user_id
+        -- This is the user's ID.
+        FROM
+        safe_user