Repository: sqlfluff/sqlfluff

Issue:
Lint and fix throws exception when having jinja for loop inside set
### Search before asking

- [X] I searched the [issues](https://github.com/sqlfluff/sqlfluff/issues) and found no similar issues.


### What Happened

To reproduce the error, create test.template.sql
```
{% set whitelisted= [
    {'name': 'COL_1'},
    {'name': 'COL_2'},
    {'name': 'COL_3'}
] %}

{% set some_part_of_the_query %}
    {% for col in whitelisted %}
    {{col.name}}{{ ", " if not loop.last }}
    {% endfor %}
{% endset %}

SELECT {{some_part_of_the_query}}
FROM SOME_TABLE

```

when running lint i get this error:
```
==== sqlfluff ====
sqlfluff:               0.11.0 python:                 3.8.12
implementation:        cpython dialect:             snowflake
verbosity:                   1 templater:               jinja

==== readout ====

=== [ path: test.template.sql ] ===

WARNING    Unable to lint test.template.sql due to an internal error. Please report this as an issue with your query's contents and stacktrace below!
To hide this warning, add the failing file to .sqlfluffignore
Traceback (most recent call last):
  File "lib/python3.8/site-packages/sqlfluff/core/linter/runner.py", line 103, in run
    yield partial()
  File "lib/python3.8/site-packages/sqlfluff/core/linter/linter.py", line 666, in lint_rendered
    parsed = cls.parse_rendered(rendered)
  File "lib/python3.8/site-packages/sqlfluff/core/linter/linter.py", line 352, in parse_rendered
    tokens, lvs, config = cls._lex_templated_file(
  File "lib/python3.8/site-packages/sqlfluff/core/linter/linter.py", line 139, in _lex_templated_file
    tokens, lex_vs = lexer.lex(templated_file)
  File "lib/python3.8/site-packages/sqlfluff/core/parser/lexer.py", line 321, in lex
    segments: Tuple[RawSegment, ...] = self.elements_to_segments(
  File "lib/python3.8/site-packages/sqlfluff/core/parser/lexer.py", line 348, in elements_to_segments
    source_slice = templated_file.templated_slice_to_source_slice(
  File "lib/python3.8/site-packages/sqlfluff/core/templaters/base.py", line 258, in templated_slice_to_source_slice
    ts_stop_sf_start, ts_stop_sf_stop = self._find_slice_indices_of_templated_pos(
  File "lib/python3.8/site-packages/sqlfluff/core/templaters/base.py", line 177, in _find_slice_indices_of_templated_pos
    raise ValueError("Position Not Found")
ValueError: Position Not Found
 
==== summary ====
violations:        0 status:         PASS
All Finished 📜 🎉!

```

This is the rendered query:
```
 SELECT

    COL_1,

    COL_2,

    COL_3


FROM SOME_TABLE

```

And when trying around to make this work i removed the new lines between the selected columns like this:
```
{% set whitelisted= [
    {'name': 'COL_1'},
    {'name': 'COL_2'},
    {'name': 'COL_3'}
] %}

{% set some_part_of_the_query %}
    {% for col in whitelisted -%}
    {{col.name}}{{ ", " if not loop.last }}
    {% endfor -%}
{% endset %}

SELECT {{some_part_of_the_query}}
FROM SOME_TABLE

```

which renders:
```
SELECT
    COL_1,
    COL_2,
    COL_3

FROM SOME_TABLE

```

And this will make the linter pass:

```
==== sqlfluff ====
sqlfluff:               0.11.0 python:                 3.8.12
implementation:        cpython dialect:             snowflake
verbosity:                   1 templater:               jinja

==== readout ====

=== [ path: test.template.sql ] ===

== [test.template.sql] PASS                                                                                                                          
==== summary ====
violations:        0 status:         PASS
All Finished 📜 🎉!

```



### Expected Behaviour

My expectations is that the linter and fix should pass.

### Observed Behaviour

Right now lint and fix throws exception (see "What Happened" section)

### How to reproduce

Mentioned above.

### Dialect

snowflake

### Version

sqlfluff, version 0.11.0

### Configuration

[sqlfluff]
verbose = 1
dialect = snowflake
templater = jinja
exclude_rules = L027,L031,L032,L036,L044,L046,L034,L050
output_line_length = 121
sql_file_exts=.sql

[sqlfluff:rules]
tab_space_size = 4
max_line_length = 250
indent_unit = space
comma_style = trailing
allow_scalar = True
single_table_references = consistent
unquoted_identifiers_policy = aliases

[sqlfluff:rules:L042]
forbid_subquery_in = both

[sqlfluff:rules:L010]  # Keywords
capitalisation_policy = upper

[sqlfluff:rules:L014]
extended_capitalisation_policy = lower

[sqlfluff:rules:L030]  # function names
extended_capitalisation_policy = upper

### Are you willing to work on and submit a PR to address the issue?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/sqlfluff/sqlfluff/blob/main/CODE_OF_CONDUCT.md)



Discussion:

Patch:
diff --git a/plugins/sqlfluff-templater-dbt/sqlfluff_templater_dbt/templater.py b/plugins/sqlfluff-templater-dbt/sqlfluff_templater_dbt/templater.py
--- a/plugins/sqlfluff-templater-dbt/sqlfluff_templater_dbt/templater.py
+++ b/plugins/sqlfluff-templater-dbt/sqlfluff_templater_dbt/templater.py
@@ -522,17 +522,21 @@ def make_template(in_str):
             # sliced_file to reflect the mapping of the added character(s) back
             # to the raw SQL.
             templated_sql = templated_sql + "\n" * n_trailing_newlines
-            sliced_file.append(
-                TemplatedFileSlice(
-                    slice_type="literal",
-                    source_slice=slice(
-                        len(source_dbt_sql) - n_trailing_newlines, len(source_dbt_sql)
-                    ),
-                    templated_slice=slice(
-                        len(templated_sql) - n_trailing_newlines, len(templated_sql)
-                    ),
+            if sliced_file and sliced_file[-1].templated_slice.stop != len(
+                templated_sql
+            ):
+                sliced_file.append(
+                    TemplatedFileSlice(
+                        slice_type="literal",
+                        source_slice=slice(
+                            len(source_dbt_sql) - n_trailing_newlines,
+                            len(source_dbt_sql),
+                        ),
+                        templated_slice=slice(
+                            len(templated_sql) - n_trailing_newlines, len(templated_sql)
+                        ),
+                    )
                 )
-            )
         return (
             TemplatedFile(
                 source_str=source_dbt_sql,
diff --git a/src/sqlfluff/core/templaters/base.py b/src/sqlfluff/core/templaters/base.py
--- a/src/sqlfluff/core/templaters/base.py
+++ b/src/sqlfluff/core/templaters/base.py
@@ -73,6 +73,7 @@ def __init__(
         templated_str: Optional[str] = None,
         sliced_file: Optional[List[TemplatedFileSlice]] = None,
         raw_sliced: Optional[List[RawFileSlice]] = None,
+        check_consistency=True,
     ):
         """Initialise the TemplatedFile.
 
@@ -104,6 +105,36 @@ def __init__(
         self._source_newlines = list(iter_indices_of_newlines(self.source_str))
         self._templated_newlines = list(iter_indices_of_newlines(self.templated_str))
 
+        # NOTE: The "check_consistency" flag should always be True when using
+        # SQLFluff in real life. This flag was only added because some legacy
+        # templater tests in test/core/templaters/jinja_test.py use hardcoded
+        # test data with issues that will trigger errors here. It would be cool
+        # to fix that data someday. I (Barry H.) started looking into it, but
+        # it was much trickier than I expected, because bits of the same data
+        # are shared across multiple tests.
+        if check_consistency:
+            # Sanity check raw string and slices.
+            pos = 0
+            rfs: RawFileSlice
+            for idx, rfs in enumerate(self.raw_sliced):
+                assert rfs.source_idx == pos
+                pos += len(rfs.raw)
+            assert pos == len(self.source_str)
+
+            # Sanity check templated string and slices.
+            previous_slice = None
+            tfs: Optional[TemplatedFileSlice] = None
+            for idx, tfs in enumerate(self.sliced_file):
+                if previous_slice:
+                    assert (
+                        tfs.templated_slice.start == previous_slice.templated_slice.stop
+                    )
+                else:
+                    assert tfs.templated_slice.start == 0
+                previous_slice = tfs
+            if self.sliced_file and templated_str is not None:
+                assert tfs.templated_slice.stop == len(templated_str)
+
     @classmethod
     def from_string(cls, raw):
         """Create TemplatedFile from a string."""
diff --git a/src/sqlfluff/core/templaters/slicers/tracer.py b/src/sqlfluff/core/templaters/slicers/tracer.py
--- a/src/sqlfluff/core/templaters/slicers/tracer.py
+++ b/src/sqlfluff/core/templaters/slicers/tracer.py
@@ -77,9 +77,6 @@ def trace(self) -> JinjaTrace:
             except IndexError:
                 pos2 = len(trace_template_output)
             p = trace_template_output[pos1 + 1 : pos2]
-            is_set_or_macro = p[:3] == "set"
-            if is_set_or_macro:
-                p = p[3:]
             m_id = regex.match(r"^([0-9a-f]+)(_(\d+))?", p)
             if not m_id:
                 raise ValueError(  # pragma: no cover
@@ -98,18 +95,7 @@ def trace(self) -> JinjaTrace:
             alt_id, content_info, literal = value
             target_slice_idx = self.find_slice_index(alt_id)
             slice_length = content_info if literal else len(str(content_info))
-            if not is_set_or_macro:
-                self.move_to_slice(target_slice_idx, slice_length)
-            else:
-                # If we find output from a {% set %} directive or a macro,
-                # record a trace without reading or updating the program
-                # counter. Such slices are always treated as "templated"
-                # because they are inserted during expansion of templated
-                # code (i.e. {% set %} variable or macro defined within the
-                # file).
-                self.record_trace(
-                    slice_length, target_slice_idx, slice_type="templated"
-                )
+            self.move_to_slice(target_slice_idx, slice_length)
         return JinjaTrace(
             self.make_template(self.raw_str).render(), self.raw_sliced, self.sliced_file
         )
@@ -241,9 +227,17 @@ def _slice_template(self) -> List[RawFileSlice]:
                         idx,
                     )
                 )
-                self.raw_slice_info[result[-1]] = self.slice_info_for_literal(
-                    len(raw), "" if set_idx is None else "set"
-                )
+                if set_idx is None:
+                    rsi = self.slice_info_for_literal(
+                        len(raw), "" if set_idx is None else "set"
+                    )
+                else:
+                    # For "set" blocks, don't generate alternate ID or code.
+                    # Sometimes, dbt users use {% set %} blocks to generate
+                    # queries that get sent to actual databases, thus causing
+                    # errors if we tamper with it.
+                    rsi = RawSliceInfo(None, None, [])
+                self.raw_slice_info[result[-1]] = rsi
                 idx += len(raw)
                 continue
             str_buff += raw
@@ -326,15 +320,20 @@ def _slice_template(self) -> List[RawFileSlice]:
                         # effects, but return a unique slice ID.
                         if trimmed_content:
                             assert m_open and m_close
-                            unique_id = self.next_slice_id()
-                            unique_alternate_id = unique_id
-                            prefix = "set" if set_idx is not None else ""
-                            open_ = m_open.group(1)
-                            close_ = m_close.group(1)
-                            alternate_code = (
-                                f"\0{prefix}{unique_alternate_id} {open_} "
-                                f"{trimmed_content} {close_}"
-                            )
+                            # For "set" blocks, don't generate alternate ID or
+                            # code. Sometimes, dbt users use {% set %} blocks to
+                            # generate queries that get sent to actual
+                            # databases, thus causing errors if we tamper with
+                            # it.
+                            if set_idx is None:
+                                unique_id = self.next_slice_id()
+                                unique_alternate_id = unique_id
+                                open_ = m_open.group(1)
+                                close_ = m_close.group(1)
+                                alternate_code = (
+                                    f"\0{unique_alternate_id} {open_} "
+                                    f"{trimmed_content} {close_}"
+                                )
                 if block_type == "block_start" and trimmed_content.split()[0] in (
                     "macro",
                     "set",
@@ -343,16 +342,24 @@ def _slice_template(self) -> List[RawFileSlice]:
                     # - {% set variable = value %}
                     # - {% set variable %}value{% endset %}
                     # https://jinja.palletsprojects.com/en/2.10.x/templates/#block-assignments
-                    # When the second format is used, set the variable 'is_set'
+                    # When the second format is used, set the variable 'set_idx'
                     # to a non-None value. This info is used elsewhere, as
                     # literals inside a {% set %} block require special handling
                     # during the trace.
                     trimmed_content_parts = trimmed_content.split(maxsplit=2)
-                    if len(trimmed_content_parts) <= 2 or not trimmed_content_parts[
-                        2
-                    ].startswith("="):
+                    if len(trimmed_content_parts) <= 2 or (
+                        not trimmed_content_parts[1].endswith("=")
+                        and not trimmed_content_parts[2].startswith("=")
+                    ):
                         set_idx = len(result)
-                elif block_type == "block_end" and set_idx is not None:
+                elif (
+                    block_type == "block_end"
+                    and set_idx is not None
+                    and (
+                        trimmed_content.startswith("endset")
+                        or trimmed_content.startswith("endmacro")
+                    )
+                ):
                     # Exiting a {% set %} block. Clear the indicator variable.
                     set_idx = None
                 m = regex.search(r"\s+$", raw, regex.MULTILINE | regex.DOTALL)


Test patch:
diff --git a/test/core/templaters/base_test.py b/test/core/templaters/base_test.py
--- a/test/core/templaters/base_test.py
+++ b/test/core/templaters/base_test.py
@@ -134,6 +134,7 @@ def test__templated_file_get_line_pos_of_char_pos(
         templated_str=templated_str,
         sliced_file=file_slices,
         fname="test",
+        check_consistency=False,
     )
     res_line_no, res_line_pos = file.get_line_pos_of_char_pos(in_charpos)
     assert res_line_no == out_line_no
@@ -287,6 +288,7 @@ def test__templated_file_templated_slice_to_source_slice(
             for rs in raw_slices
         ],
         fname="test",
+        check_consistency=False,
     )
     source_slice = file.templated_slice_to_source_slice(in_slice)
     literal_test = file.is_source_slice_literal(source_slice)
@@ -303,5 +305,6 @@ def test__templated_file_source_only_slices():
             RawFileSlice("b" * 7, "comment", 10),
             RawFileSlice("a" * 10, "literal", 17),
         ],
+        check_consistency=False,
     )
     assert file.source_only_slices() == [RawFileSlice("b" * 7, "comment", 10)]
diff --git a/test/core/templaters/jinja_test.py b/test/core/templaters/jinja_test.py
--- a/test/core/templaters/jinja_test.py
+++ b/test/core/templaters/jinja_test.py
@@ -370,6 +370,35 @@ def test__templater_jinja_slices(case: RawTemplatedTestCase):
     assert actual_rs_source_list == case.expected_raw_sliced__source_list
 
 
+def test_templater_set_block_handling():
+    """Test handling of literals in {% set %} blocks.
+
+    Specifically, verify they are not modified in the alternate template.
+    """
+
+    def run_query(sql):
+        # Prior to the bug fix, this assertion failed. This was bad because,
+        # inside JinjaTracer, dbt templates similar to the one in this test
+        # would call the database with funky SQL (including weird strings it
+        # uses internally like: 00000000000000000000000000000002.
+        assert sql == "\n\nselect 1 from foobarfoobarfoobarfoobar_dev\n\n"
+        return sql
+
+    t = JinjaTemplater(override_context=dict(run_query=run_query))
+    instr = """{% set my_query1 %}
+select 1 from foobarfoobarfoobarfoobar_{{ "dev" }}
+{% endset %}
+{% set my_query2 %}
+{{ my_query1 }}
+{% endset %}
+
+{{ run_query(my_query2) }}
+"""
+    outstr, vs = t.process(in_str=instr, fname="test", config=FluffConfig())
+    assert str(outstr) == "\n\n\n\n\nselect 1 from foobarfoobarfoobarfoobar_dev\n\n\n"
+    assert len(vs) == 0
+
+
 def test__templater_jinja_error_variable():
     """Test missing variable error handling in the jinja templater."""
     t = JinjaTemplater(override_context=dict(blah="foo"))
@@ -846,6 +875,45 @@ def test__templater_jinja_slice_template(test, result):
                 ("literal", slice(312, 327, None), slice(27, 42, None)),
             ],
         ),
+        (
+            # Test for issue 2835. There's no space between "col" and "="
+            """{% set col= "col1" %}
+SELECT {{ col }}
+""",
+            None,
+            [
+                ("block_start", slice(0, 21, None), slice(0, 0, None)),
+                ("literal", slice(21, 29, None), slice(0, 8, None)),
+                ("templated", slice(29, 38, None), slice(8, 12, None)),
+                ("literal", slice(38, 39, None), slice(12, 13, None)),
+            ],
+        ),
+        (
+            # Another test for issue 2835. The {% for %} loop inside the
+            # {% set %} caused JinjaTracer to think the {% set %} ended
+            # at the {% endfor %}
+            """{% set some_part_of_the_query %}
+    {% for col in ["col1"] %}
+    {{col}}
+    {% endfor %}
+{% endset %}
+
+SELECT {{some_part_of_the_query}}
+FROM SOME_TABLE
+""",
+            None,
+            [
+                ("block_start", slice(0, 32, None), slice(0, 0, None)),
+                ("literal", slice(32, 37, None), slice(0, 0, None)),
+                ("block_start", slice(37, 62, None), slice(0, 0, None)),
+                ("block_end", slice(79, 91, None), slice(0, 0, None)),
+                ("literal", slice(91, 92, None), slice(0, 0, None)),
+                ("block_end", slice(92, 104, None), slice(0, 0, None)),
+                ("literal", slice(104, 113, None), slice(0, 9, None)),
+                ("templated", slice(113, 139, None), slice(9, 29, None)),
+                ("literal", slice(139, 156, None), slice(29, 46, None)),
+            ],
+        ),
     ],
 )
 def test__templater_jinja_slice_file(raw_file, override_context, result, caplog):