Repository: reata/sqllineage
Pull request #429 (2023-08-13 05:30:45)

fix: not throwing exception for tsql unparsable batch
making progress in #423 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/429?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#429](https://app.codecov.io/gh/reata/sqllineage/pull/429?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (e685e3c) into [master](https://app.codecov.io/gh/reata/sqllineage/commit/6f9c9fd451669476eb0adad186120b88ee3aed37?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (6f9c9fd) will **decrease** coverage by `0.16%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master     #429      +/-   ##
==========================================
- Coverage   99.54%   99.39%   -0.16%     
==========================================
  Files          41       41              
  Lines        1968     1969       +1     
==========================================
- Hits         1959     1957       -2     
- Misses          9       12       +3     
```


| [Files Changed](https://app.codecov.io/gh/reata/sqllineage/pull/429?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [sqllineage/core/holders.py](https://app.codecov.io/gh/reata/sqllineage/pull/429?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL2hvbGRlcnMucHk=) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlfluff/analyzer.py](https://app.codecov.io/gh/reata/sqllineage/pull/429?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9hbmFseXplci5weQ==) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlfluff/handlers/target.py](https://app.codecov.io/gh/reata/sqllineage/pull/429?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9oYW5kbGVycy90YXJnZXQucHk=) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlfluff/utils.py](https://app.codecov.io/gh/reata/sqllineage/pull/429?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi91dGlscy5weQ==) | `99.23% <100.00%> (-0.77%)` | :arrow_down: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/reata/sqllineage/pull/429/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)



Patch:
diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index 4e63212..e2ec54c 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -108,6 +108,9 @@ class SubQueryLineageHolder(ColumnLineageMixin):
             tgt_tbl = list(self.write)[0]
             for idx, tgt_col in enumerate(tgt_cols):
                 tgt_col.parent = tgt_tbl
+                if tgt_col in self.write_columns:
+                    # for DDL with PARTITIONED BY (col) or CLUSTERED BY (col), column can be added multiple times
+                    break
                 self.graph.add_edge(
                     tgt_tbl, tgt_col, type=EdgeType.HAS_COLUMN, **{EdgeTag.INDEX: idx}
                 )
diff --git a/sqllineage/core/parser/sqlfluff/analyzer.py b/sqllineage/core/parser/sqlfluff/analyzer.py
index 277b28b..10808e4 100644
--- a/sqllineage/core/parser/sqlfluff/analyzer.py
+++ b/sqllineage/core/parser/sqlfluff/analyzer.py
@@ -1,10 +1,7 @@
-from sqlfluff.core import Linter
+from sqlfluff.core import Linter, SQLLexError, SQLParseError
 
 from sqllineage.core.analyzer import LineageAnalyzer
-from sqllineage.core.holders import (
-    StatementLineageHolder,
-    SubQueryLineageHolder,
-)
+from sqllineage.core.holders import StatementLineageHolder
 from sqllineage.core.parser.sqlfluff.extractors.lineage_holder_extractor import (
     LineageHolderExtractor,
 )
@@ -22,8 +19,20 @@ class SqlFluffLineageAnalyzer(LineageAnalyzer):
         self._dialect = dialect
 
     def analyze(self, sql: str) -> StatementLineageHolder:
-        linter = Linter(dialect=self._dialect)
-        parsed = linter.parse_string(sql)
+        parsed = Linter(dialect=self._dialect).parse_string(sql)
+        violations = [
+            str(e)
+            for e in parsed.violations
+            if isinstance(e, (SQLLexError, SQLParseError))
+        ]
+        if violations:
+            violation_msg = "\n".join(violations)
+            raise InvalidSyntaxException(
+                f"This SQL statement is unparsable, please check potential syntax error for SQL:\n"
+                f"{sql}\n"
+                f"{violation_msg}"
+            )
+
         statement_segment = None
         for top_segment in getattr(parsed.tree, "segments", []):
             if top_segment.type == "statement":
@@ -37,37 +46,17 @@ class SqlFluffLineageAnalyzer(LineageAnalyzer):
                 f"SQLLineage cannot parse SQL:" f"{sql}"
             )  # pragma: no cover
         else:
-            extractors = [
+            for extractor in [
                 extractor_cls(self._dialect)
                 for extractor_cls in LineageHolderExtractor.__subclasses__()
-            ]
-            if statement_segment.type == "unparsable":
-                raise InvalidSyntaxException(
-                    f"This SQL statement is unparsable, please check potential syntax error for SQL:"
+            ]:
+                if extractor.can_extract(statement_segment.type):
+                    lineage_holder = extractor.extract(
+                        statement_segment, AnalyzerContext()
+                    )
+                    return StatementLineageHolder.of(lineage_holder)
+            else:
+                raise UnsupportedStatementException(
+                    f"SQLLineage doesn't support analyzing statement type [{statement_segment.type}] for SQL:"
                     f"{sql}"
                 )
-            else:
-                if any(
-                    extractor.can_extract(statement_segment.type)
-                    for extractor in extractors
-                ):
-                    if "unparsable" in statement_segment.descendant_type_set:
-                        raise InvalidSyntaxException(
-                            f"{statement_segment.type} is partially unparsable, "
-                            f"please check potential syntax error for SQL:"
-                            f"{sql}"
-                        )
-                    else:
-                        lineage_holder = SubQueryLineageHolder()
-                        for extractor in extractors:
-                            if extractor.can_extract(statement_segment.type):
-                                lineage_holder = extractor.extract(
-                                    statement_segment, AnalyzerContext()
-                                )
-                                break
-                        return StatementLineageHolder.of(lineage_holder)
-                else:
-                    raise UnsupportedStatementException(
-                        f"SQLLineage doesn't support analyzing statement type [{statement_segment.type}] for SQL:"
-                        f"{sql}"
-                    )
diff --git a/sqllineage/core/parser/sqlfluff/handlers/target.py b/sqllineage/core/parser/sqlfluff/handlers/target.py
index ec41264..7339334 100644
--- a/sqllineage/core/parser/sqlfluff/handlers/target.py
+++ b/sqllineage/core/parser/sqlfluff/handlers/target.py
@@ -89,11 +89,15 @@ class TargetHandler(ConditionalSegmentBaseHandler):
             self._reset_tokens()
 
         elif segment.type in {"literal", "storage_location"}:
-            if self.prev_token_from:
-                holder.add_read(Path(escape_identifier_name(segment.raw)))
+            if segment.raw.isnumeric():
+                # Special Handling for Spark Bucket Table DDL
+                pass
             else:
-                holder.add_write(Path(escape_identifier_name(segment.raw)))
-            self._reset_tokens()
+                if self.prev_token_from:
+                    holder.add_read(Path(escape_identifier_name(segment.raw)))
+                else:
+                    holder.add_write(Path(escape_identifier_name(segment.raw)))
+                self._reset_tokens()
 
         elif segment.type == "into_table_clause":
             obj_ref = get_child(segment, "object_reference")
@@ -145,9 +149,13 @@ class TargetHandler(ConditionalSegmentBaseHandler):
             """
             sub_segments = list_child_segments(segment)
             if all(
-                sub_segment.type == "column_reference" for sub_segment in sub_segments
+                sub_segment.type in ["column_reference", "column_definition"]
+                for sub_segment in sub_segments
             ):
-                # target columns only apply to bracketed column references
-                holder.add_write_column(
-                    *[SqlFluffColumn.of(sub_segment) for sub_segment in sub_segments]
-                )
+                # target columns only apply to bracketed column_reference and column_definition
+                columns = []
+                for sub_segment in sub_segments:
+                    if sub_segment.type == "column_definition":
+                        sub_segment = get_child(sub_segment, "identifier")
+                    columns.append(SqlFluffColumn.of(sub_segment))
+                holder.add_write_column(*columns)
diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
index 053aef2..0697bc1 100644
--- a/sqllineage/core/parser/sqlfluff/utils.py
+++ b/sqllineage/core/parser/sqlfluff/utils.py
@@ -202,7 +202,7 @@ def list_child_segments(
             for seg in segment.iter_segments(
                 expanding=["expression"], pass_through=True
             ):
-                if seg.type == "column_reference":
+                if seg.type in ["column_reference", "column_definition"]:
                     result.append(seg)
                 else:
                     for s in seg.segments:


Test patch:
diff --git a/tests/core/test_exception.py b/tests/core/test_exception.py
index 7081ffd..9edf144 100644
--- a/tests/core/test_exception.py
+++ b/tests/core/test_exception.py
@@ -27,6 +27,15 @@ def test_partial_unparsable_query_in_sqlfluff():
         )._eval()
 
 
+def test_partial_unparsable_query_in_sqlfluff_with_tsql_batch():
+    sql = """SELECT *
+INTO tgt
+FROM tab1 src1 AS src1
+CROSS JOIN tab2 AS src2"""
+    with pytest.raises(InvalidSyntaxException):
+        LineageRunner(sql, dialect="tsql")._eval()
+
+
 def test_unsupported_query_type_in_sqlfluff():
     with pytest.raises(UnsupportedStatementException):
         LineageRunner(
diff --git a/tests/helpers.py b/tests/helpers.py
index 063acfe..cf22424 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -46,6 +46,7 @@ def assert_table_lineage_equal(
     dialect: str = "ansi",
     test_sqlfluff: bool = True,
     test_sqlparse: bool = True,
+    skip_graph_check: bool = False,
 ):
     lr = LineageRunner(sql, dialect=SQLPARSE_DIALECT)
     lr_sqlfluff = LineageRunner(sql, dialect=dialect)
@@ -53,7 +54,7 @@ def assert_table_lineage_equal(
         assert_table_lineage(lr, source_tables, target_tables)
     if test_sqlfluff:
         assert_table_lineage(lr_sqlfluff, source_tables, target_tables)
-    if test_sqlparse and test_sqlfluff:
+    if test_sqlparse and test_sqlfluff and not skip_graph_check:
         assert_lr_graphs_match(lr, lr_sqlfluff)
 
 
@@ -63,6 +64,7 @@ def assert_column_lineage_equal(
     dialect: str = "ansi",
     test_sqlfluff: bool = True,
     test_sqlparse: bool = True,
+    skip_graph_check: bool = False,
 ):
     lr = LineageRunner(sql, dialect=SQLPARSE_DIALECT)
     lr_sqlfluff = LineageRunner(sql, dialect=dialect)
@@ -70,7 +72,7 @@ def assert_column_lineage_equal(
         assert_column_lineage(lr, column_lineages)
     if test_sqlfluff:
         assert_column_lineage(lr_sqlfluff, column_lineages)
-    if test_sqlparse and test_sqlfluff:
+    if test_sqlparse and test_sqlfluff and not skip_graph_check:
         assert_lr_graphs_match(lr, lr_sqlfluff)
 
 
diff --git a/tests/sql/table/multiple_statements/test_tmp_table.py b/tests/sql/table/multiple_statements/test_tmp_table.py
index 5fb9a0b..088ad08 100644
--- a/tests/sql/table/multiple_statements/test_tmp_table.py
+++ b/tests/sql/table/multiple_statements/test_tmp_table.py
@@ -24,7 +24,7 @@ FROM tab_a;"""
 
 def test_create_after_drop():
     assert_table_lineage_equal(
-        "DROP TABLE IF EXISTS tab1; CREATE TABLE IF NOT EXISTS tab1 (col1 STRING)",
+        "DROP TABLE IF EXISTS tab1; CREATE TABLE IF NOT EXISTS tab1 AS SELECT 1",
         None,
         {"tab1"},
     )
@@ -32,7 +32,7 @@ def test_create_after_drop():
 
 def test_drop_after_create():
     assert_table_lineage_equal(
-        "CREATE TABLE IF NOT EXISTS tab1 (col1 STRING);DROP TABLE IF EXISTS tab1",
+        "CREATE TABLE IF NOT EXISTS tab1 AS SELECT 1; DROP TABLE IF EXISTS tab1",
         None,
         None,
     )
diff --git a/tests/sql/table/other_single_statement/test_other_without_lineage_dialect_specific.py b/tests/sql/table/other_single_statement/test_other_without_lineage_dialect_specific.py
index 26a6689..fce59f0 100644
--- a/tests/sql/table/other_single_statement/test_other_without_lineage_dialect_specific.py
+++ b/tests/sql/table/other_single_statement/test_other_without_lineage_dialect_specific.py
@@ -13,7 +13,7 @@ def test_rename_table(dialect: str):
     assert_table_lineage_equal("rename table tab1 to tab2", dialect=dialect)
 
 
-@pytest.mark.parametrize("dialect", ["exasol", "mysql", "teradata"])
+@pytest.mark.parametrize("dialect", ["mysql"])
 def test_rename_tables(dialect: str):
     assert_table_lineage_equal(
         "rename table tab1 to tab2, tab3 to tab4", dialect=dialect
diff --git a/tests/sql/table/test_create.py b/tests/sql/table/test_create.py
index ae5776a..6ffb610 100644
--- a/tests/sql/table/test_create.py
+++ b/tests/sql/table/test_create.py
@@ -2,12 +2,20 @@ from ...helpers import assert_table_lineage_equal
 
 
 def test_create():
-    assert_table_lineage_equal("CREATE TABLE tab1 (col1 STRING)", None, {"tab1"})
+    assert_table_lineage_equal(
+        "CREATE TABLE tab1 (col1 STRING)",
+        None,
+        {"tab1"},
+        skip_graph_check=True,  # sqlfluff graph includes table to column edge
+    )
 
 
 def test_create_if_not_exist():
     assert_table_lineage_equal(
-        "CREATE TABLE IF NOT EXISTS tab1 (col1 STRING)", None, {"tab1"}
+        "CREATE TABLE IF NOT EXISTS tab1 (col1 STRING)",
+        None,
+        {"tab1"},
+        skip_graph_check=True,  # sqlfluff graph includes table to column edge
     )
 
 
@@ -34,12 +42,3 @@ col1""",
         {"tab1"},
         {"view1"},
     )
-
-
-def test_bucket_with_using_parenthesis():
-    assert_table_lineage_equal(
-        """CREATE TABLE tbl1 (col1 VARCHAR)
-  WITH (bucketed_on = array['col1'], bucket_count = 256);""",
-        None,
-        {"tbl1"},
-    )
diff --git a/tests/sql/table/test_create_dialect_specific.py b/tests/sql/table/test_create_dialect_specific.py
index 35f83cf..f9e9615 100644
--- a/tests/sql/table/test_create_dialect_specific.py
+++ b/tests/sql/table/test_create_dialect_specific.py
@@ -3,12 +3,31 @@ import pytest
 from ...helpers import assert_table_lineage_equal
 
 
-@pytest.mark.parametrize("dialect", ["bigquery", "snowflake"])
-def test_create_bucket_table(dialect: str):
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_create_bucket_table_in_spark(dialect: str):
     assert_table_lineage_equal(
-        "CREATE TABLE tab1 USING parquet CLUSTERED BY (col1) INTO 500 BUCKETS",
+        """CREATE TABLE student (id INT, name STRING, age INT)
+USING CSV
+PARTITIONED BY (age)
+CLUSTERED BY (Id) INTO 4 buckets""",
         None,
-        {"tab1"},
+        {"student"},
+        dialect,
+        skip_graph_check=True,  # sqlfluff graph includes table to column edge
+    )
+
+
+@pytest.mark.parametrize("dialect", ["athena"])
+def test_create_bucket_table_in_athena(dialect: str):
+    assert_table_lineage_equal(
+        """CREATE TABLE bar
+WITH (
+  bucketed_by = ARRAY['customer_id'],
+  bucket_count = 8
+)
+AS SELECT * FROM foo""",
+        {"foo"},
+        {"bar"},
         dialect,
     )
 
@@ -62,4 +81,5 @@ STORED AS TEXTFILE""",  # noqa
         None,
         {"apachelog"},
         dialect=dialect,
+        skip_graph_check=True,  # sqlfluff graph includes table to column edge
     )