Repository: reata/sqllineage

Issue:
Enable support of `sqlfluff` context
## Is your feature request related to a problem? Please describe

As soon as dialect-based parsing use `sqlfluff` to parse SQL it will be great to support sqlfluff context (from `.sqlfluff` files) where we can use/set templater and other options

## Describe the solution you'd like

`sqlfluff.core.linter.Linter` have optional `config: FluffConfig` parameter. Need to pass it or ensure that default is used

## Additional context

```verilog
dmytro.kulyk@MAC-570022 datalake % sqllineage -f transform/airflow/dags/queries/aws_transforms/aggregate/pokermatch_payment.partitions_range.sql
Traceback (most recent call last):
  File "/opt/homebrew/bin/sqllineage", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/sqllineage/cli.py", line 125, in main
    runner.print_table_lineage()
  File "/opt/homebrew/lib/python3.11/site-packages/sqllineage/runner.py", line 175, in print_table_lineage
    print(str(self))
          ^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/sqllineage/runner.py", line 26, in wrapper
    self._eval()
  File "/opt/homebrew/lib/python3.11/site-packages/sqllineage/runner.py", line 195, in _eval
    stmt_holder = analyzer.analyze(stmt, session.metadata_provider)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/sqllineage/core/parser/sqlfluff/analyzer.py", line 48, in analyze
    raise UnsupportedStatementException(
sqllineage.exceptions.UnsupportedStatementException: SQLLineage cannot parse SQL:{%- set task_id = 'ph2_transaction.data_processing' %}
{%- set transaction_date = ti.xcom_pull(task_ids=task_id, key='new_partitions') %}
{% if transaction_date|length > 0 %}
    select
        row_num,
        min(as_of_date) as start_date,
        max(as_of_date) as end_date,
        array_agg(as_of_date) as dates
    from (
        select
...
```

running `sqlfluff` directly:
```verilog
dmytro.kulyk@MAC-570022 datalake % sqlfluff parse transform/airflow/dags/queries/aws_transforms/aggregate/pokermatch_payment.partitions_range.sql --verbose | more
==== sqlfluff ====
sqlfluff:                2.3.5 python:                 3.11.6
implementation:        cpython verbosity:                   1
dialect:                athena templater:               jinja
rules:                                all
=== [ path: transform/airflow/dags/queries/aws_transforms/aggregate/pokermatch_payment.partitions_range.sql ] ===

[L:  1, P:  1]      |file:
[L:  1, P:  1]      |    [META] placeholder:                                       [Type: 'templated', Raw: "{%- set task_id = 'ph2_transaction.data_processing' %}"]
[L:  1, P: 55]      |    [META] placeholder:                                       [Type: 'literal', Raw: '\n']
[L:  2, P:  1]      |    [META] placeholder:                                       [Type: 'templated', Raw: "{%- set transaction_date = ti.xcom_pull(task_ids=task_id, key='new_partitions') %}"]
[L:  2, P: 83]      |    newline:                                                  '\n'
[L:  3, P:  1]      |    [META] placeholder:                                       [Type: 'block_start', Raw: '{% if transaction_date|length > 0 %}', Block: 'd42817']
[L:  3, P: 37]      |    [META] indent:                                            [Block: 'd42817']
[L:  3, P: 37]      |    newline:                                                  '\n'
[L:  4, P:  1]      |    whitespace:                                               '    '
[L:  4, P:  5]      |    statement:
[L:  4, P:  5]      |        select_statement:
[L:  4, P:  5]      |            select_clause:
[L:  4, P:  5]      |                keyword:                                      'select'
[L:  4, P: 11]      |                [META] indent:
[L:  4, P: 11]      |                newline:                                      '\n'
[L:  5, P:  1]      |                whitespace:                                   '        '
[L:  5, P:  9]      |                select_clause_element:
[L:  5, P:  9]      |                    column_reference:
[L:  5, P:  9]      |                        naked_identifier:                     'row_num'
...
```

`.sqlfluff` example
```ini
[sqlfluff]
templater = jinja
output_line_length = 180
exclude_rules = AL07, AM02, AM03, AM05, CV10, LT08, ST05, ST06, ST07
warnings = RF01
dialect = athena
large_file_skip_char_limit = 0
large_file_skip_byte_limit = 0
max_line_length = 120

[sqlfluff:rules:references.consistent]
single_table_references = consistent

[sqlfluff:rules:references.keywords]
ignore_words =
    account, domain, language, operator, rank, result, source, type, uid, position, format, state, offset, name,
    location, role, user, comment, method, date, quarter, hour, week, weekday, days

[sqlfluff:templater:jinja]
library_path = operation/deploy/lint
apply_dbt_builtins = false
```



Discussion:
It makes sense to be support analyzing templated SQL so lineage can analyzed based on what's in orchestration system instead of merely on SQL logs.

Quick question: do you require any configuration in sqlfluff other than templating? Because supports all sqlfluff configuration looks like a overkill for me.
we use both `sqlfluff` and `sqllineage` (which is using `sqlfluff`) and looks like templating configuration is ignored
Thanks for your confirmation. I'll see what we can do.



Patch:
diff --git a/sqllineage/core/parser/sqlfluff/analyzer.py b/sqllineage/core/parser/sqlfluff/analyzer.py
index b7acf059..74e0ace8 100644
--- a/sqllineage/core/parser/sqlfluff/analyzer.py
+++ b/sqllineage/core/parser/sqlfluff/analyzer.py
@@ -1,7 +1,13 @@
 import warnings
 from typing import Dict, List
 
-from sqlfluff.core import Linter, SQLLexError, SQLParseError, dialect_readout
+from sqlfluff.core import (
+    FluffConfig,
+    Linter,
+    SQLLexError,
+    SQLParseError,
+    dialect_readout,
+)
 from sqlfluff.core.parser import BaseSegment
 
 from sqllineage.core.analyzer import LineageAnalyzer
@@ -73,7 +79,9 @@ def analyze(
                     )
 
     def _list_specific_statement_segment(self, sql: str):
-        parsed = Linter(dialect=self._dialect).parse_string(sql)
+        parsed = Linter(
+            config=FluffConfig.from_root(overrides={"dialect": self._dialect})
+        ).parse_string(sql)
         violations = [
             str(e)
             for e in parsed.violations


Test patch:
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index e3b96e63..ecde150d 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -11,8 +11,8 @@
 from sqllineage.runner import LineageRunner
 
 
-@patch(
-    "os.environ",
+@patch.dict(
+    os.environ,
     {
         "SQLLINEAGE_DIRECTORY": os.path.join(os.path.dirname(__file__), "data"),
         "SQLLINEAGE_DEFAULT_SCHEMA": "<default>",
diff --git a/tests/core/test_exception.py b/tests/core/test_exception.py
index 50991b40..64cfeabb 100644
--- a/tests/core/test_exception.py
+++ b/tests/core/test_exception.py
@@ -1,3 +1,4 @@
+import os
 from unittest.mock import patch
 
 import pytest
@@ -56,7 +57,7 @@ def test_syntax_warning_no_semicolon_in_tsql():
         )._eval()
 
 
-@patch("os.environ", {"SQLLINEAGE_TSQL_NO_SEMICOLON": "TRUE"})
+@patch.dict(os.environ, {"SQLLINEAGE_TSQL_NO_SEMICOLON": "TRUE"})
 def test_user_warning_enable_tsql_no_semicolon_with_other_dialect():
     with pytest.warns(UserWarning):
         LineageRunner(
diff --git a/tests/core/test_runner.py b/tests/core/test_runner.py
index 96593928..23210cfd 100644
--- a/tests/core/test_runner.py
+++ b/tests/core/test_runner.py
@@ -1,6 +1,10 @@
+import os
+import tempfile
+
 from sqllineage.core.models import SubQuery
 from sqllineage.runner import LineageRunner
 from sqllineage.utils.constant import LineageLevel
+from ..helpers import assert_table_lineage_equal
 
 
 def test_runner_dummy():
@@ -31,3 +35,19 @@ def test_get_column_lineage_exclude_subquery_inpath():
     for col_tuple in parse.get_column_lineage(exclude_subquery_columns=True):
         for col in col_tuple:
             assert not isinstance(col.parent, SubQuery)
+
+
+def test_respect_sqlfluff_configuration_file():
+    sqlfluff_config = """[sqlfluff:templater:jinja:context]
+num_things=456
+tbl_name=my_table"""
+    cwd = os.getcwd()
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        try:
+            os.chdir(tmpdirname)
+            with open(".sqlfluff", "w") as f:
+                f.write(sqlfluff_config)
+            sql = "SELECT {{ num_things }} FROM {{ tbl_name }} WHERE id > 10 LIMIT 5"
+            assert_table_lineage_equal(sql, {"my_table"}, test_sqlparse=False)
+        finally:
+            os.chdir(cwd)
diff --git a/tests/sql/column/test_column_select_lateral_alias_ref.py b/tests/sql/column/test_column_select_lateral_alias_ref.py
index 9a6f0ed9..c22090f6 100644
--- a/tests/sql/column/test_column_select_lateral_alias_ref.py
+++ b/tests/sql/column/test_column_select_lateral_alias_ref.py
@@ -1,3 +1,4 @@
+import os
 from unittest.mock import patch
 
 import pytest
@@ -14,7 +15,7 @@
 
 
 @pytest.mark.parametrize("provider", providers)
-@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+@patch.dict(os.environ, {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
 def test_column_top_level_enable_lateral_ref(
     provider: MetaDataProvider,
 ):
@@ -106,7 +107,7 @@ def test_column_top_level_enable_lateral_ref(
 
 
 @pytest.mark.parametrize("provider", providers)
-@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+@patch.dict(os.environ, {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
 def test_column_top_level_enable_lateral_ref_with_metadata_from_table(
     provider: MetaDataProvider,
 ):
@@ -148,7 +149,7 @@ def test_column_top_level_enable_lateral_ref_with_metadata_from_table(
 
 
 @pytest.mark.parametrize("provider", providers)
-@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+@patch.dict(os.environ, {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
 def test_column_top_level_enable_lateral_ref_with_metadata_from_subquery(
     provider: MetaDataProvider,
 ):
@@ -195,7 +196,7 @@ def test_column_top_level_enable_lateral_ref_with_metadata_from_subquery(
 
 
 @pytest.mark.parametrize("provider", providers)
-@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+@patch.dict(os.environ, {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
 def test_column_top_level_enable_lateral_ref_with_metadata_from_nested_subquery(
     provider: MetaDataProvider,
 ):
@@ -247,7 +248,7 @@ def test_column_top_level_enable_lateral_ref_with_metadata_from_nested_subquery(
 
 
 @pytest.mark.parametrize("provider", providers)
-@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+@patch.dict(os.environ, {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
 def test_column_enable_lateral_ref_within_subquery(
     provider: MetaDataProvider,
 ):
diff --git a/tests/sql/table/multiple_statements/test_tsql_no_semicolon.py b/tests/sql/table/multiple_statements/test_tsql_no_semicolon.py
index c96a7876..dde29572 100644
--- a/tests/sql/table/multiple_statements/test_tsql_no_semicolon.py
+++ b/tests/sql/table/multiple_statements/test_tsql_no_semicolon.py
@@ -1,3 +1,4 @@
+import os
 from unittest.mock import patch
 
 import pytest
@@ -5,7 +6,7 @@
 from ....helpers import assert_table_lineage_equal
 
 
-@patch("os.environ", {"SQLLINEAGE_TSQL_NO_SEMICOLON": "TRUE"})
+@patch.dict(os.environ, {"SQLLINEAGE_TSQL_NO_SEMICOLON": "TRUE"})
 @pytest.mark.parametrize("dialect", ["tsql"])
 def test_tsql_multi_statement_no_semicolon(dialect: str):
     """