Incorrect Source Tables recognition when using SqlFluffLineageAnalyzer

[Aka-shi]

Sample query

sql = """
Insert into target_table

WITH pl AS (
  Select id, school
From table1)

,fb AS (
Select id, school
FROM table2)

SELECT
table3.id, table3.name, table1.school
FROM table3
LEFT JOIN table1 ON table3.id = table1.id

UNION

SELECT
table4.id, table4.name, table2.school
FROM table4
LEFT JOIN table2 ON table4.id = table2.id
"""

When I run result = LineageRunner(sql) I get the following result which is the correct one.

Result is  Statements(#): 1
Source Tables:
    <default>.table1
    <default>.table2
    <default>.table3
    <default>.table4
Target Tables:
    <default>.target_table

But when I provide a dialect result = LineageRunner(sql, 'ansi'), I get the following result

Result is  Statements(#): 1
Source Tables:
    <default>.table1
    <default>.table2
Target Tables:
    <default>.target_table

It doesn't seem to parse the query after CTEs. This looks like a bug. If anyone can guide me on how I can fix this or take this up it would be really helpful since sqlparse is being deprecated in future versions.

[Aka-shi]

It gives correct result when I modify the query to this

sql = """
Insert into target_table

WITH pl AS (
  Select id, school
From table1)

,fb AS (
Select id, school
FROM table2)

SELECT
table3.id, table3.name, table1.school
FROM table3
LEFT JOIN table1 ON table3.id = table1.id

[reata]

Yes, it's a bug. It's because we didn't handle UNION in DmlCteExtractor

UNION is of type set_expression in sqlfluff AST, and we should hanlde set_expression the same way as select_statement. This should be a simple fix, looking forward to your PR.