Repository: reata/sqllineage

Issue:
let user choose whether to filter temp table or not
Currently, If a table is both source table and target table. sqllineage will identify this table as temp table and hide it from user.

We should give use control over the display of temp table

Comments:

Patch:
diff --git a/sqllineage/combiners.py b/sqllineage/combiners.py
index a652340..2887748 100644
--- a/sqllineage/combiners.py
+++ b/sqllineage/combiners.py
@@ -1,7 +1,9 @@
 from functools import reduce
-from operator import add
+from operator import add, or_
+from typing import Set
 
 from sqllineage.core import LineageResult
+from sqllineage.models import Table
 
 
 class LineageCombiner:
@@ -30,10 +32,22 @@ class DefaultLineageCombiner(LineageCombiner):
                         if table_old in st_tables:
                             st_tables.remove(table_old)
                             st_tables.add(table_new)
+            elif lineage_result.with_:
+                combined_result.read |= lineage_result.read - lineage_result.with_
+                combined_result.write |= lineage_result.write
             else:
                 combined_result.read |= lineage_result.read
                 combined_result.write |= lineage_result.write
         tmp_tables = combined_result.read.intersection(combined_result.write)
+        self_depend_tables = reduce(
+            or_,
+            (
+                lineage_result.read.intersection(lineage_result.write)
+                for lineage_result in args
+            ),
+            set(),
+        )  # type: Set[Table]
+        tmp_tables -= self_depend_tables
         combined_result.read -= tmp_tables
         combined_result.write -= tmp_tables
         return combined_result
diff --git a/sqllineage/core.py b/sqllineage/core.py
index 47d3265..3657362 100644
--- a/sqllineage/core.py
+++ b/sqllineage/core.py
@@ -27,12 +27,10 @@ TEMP_TABLE_TOKENS = ("WITH",)
 class LineageResult:
     """Statement(s) Level Lineage Result."""
 
-    __slots__ = ["read", "write", "rename", "drop"]
+    __slots__ = ["read", "write", "rename", "drop", "with_"]
     if TYPE_CHECKING:
-        read = None  # type: Set[Table]
-        write = None  # type: Set[Table]
-        rename = None  # type: Set[Tuple[Table, Table]]
-        drop = None  # type: Set[Table]
+        read = write = drop = with_ = set()  # type: Set[Table]
+        rename = set()  # type: Set[Tuple[Table, Table]]
 
     def __init__(self) -> None:
         for attr in self.__slots__:
@@ -54,6 +52,9 @@ class LineageResult:
             for attr in self.__slots__
         )
 
+    def __repr__(self):
+        return str(self)
+
 
 class LineageAnalyzer:
     """SQL Statement Level Lineage Analyzer."""
@@ -160,8 +161,7 @@ class LineageAnalyzer:
                 else:
                     if not isinstance(sub_token, Identifier):
                         raise SQLLineageException("An Identifier is expected")
-                    self._lineage_result.read.add(Table.create(sub_token))
-                    self._lineage_result.write.add(Table.create(sub_token))
+                    self._lineage_result.with_.add(Table.create(sub_token))
                     self._extract_from_DML(sub_token)
                     temp_table_token_flag = False
 


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index a91e485..58702e9 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -2,5 +2,5 @@ from sqllineage.core import LineageResult
 
 
 def test_dummy():
-    assert str(LineageResult())
+    assert str(LineageResult()) == repr(LineageResult())
     assert LineageResult() + LineageResult() is not None
diff --git a/tests/test_insert.py b/tests/test_insert.py
index f0b1c6c..b206926 100644
--- a/tests/test_insert.py
+++ b/tests/test_insert.py
@@ -47,6 +47,17 @@ def test_insert_overwrite_values():
     )
 
 
+def test_insert_overwrite_from_self():
+    helper(
+        """INSERT OVERWRITE TABLE tab_1
+SELECT tab2.col_a from tab_2
+JOIN tab_1
+ON tab_1.col_a = tab_2.cola""",
+        {"tab_1", "tab_2"},
+        {"tab_1"},
+    )
+
+
 def test_with_insert():
     helper(
         "WITH tab1 AS (SELECT * FROM tab2) INSERT INTO tab3 SELECT * FROM tab1",
diff --git a/tests/test_models.py b/tests/test_models.py
index 47f09bc..9bc2294 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,9 +1,16 @@
+import pytest
+
+from sqllineage.exceptions import SQLLineageException
 from sqllineage.models import Schema, Table
 
 
 def test_repr_dummy():
     assert repr(Schema())
     assert repr(Table(""))
+    with pytest.raises(SQLLineageException):
+        Table("a.b.c")
+    with pytest.warns(Warning):
+        Table("a.b", Schema("c"))
 
 
 def test_hash_eq():
diff --git a/tests/test_others.py b/tests/test_others.py
index 946bf4c..863e1d0 100644
--- a/tests/test_others.py
+++ b/tests/test_others.py
@@ -73,6 +73,12 @@ drop table tab_a;"""
     helper(sql, {"tab_b"}, {"tab_c"})
 
 
+def test_new_create_tab_as_tmp_table():
+    sql = """create table tab_a as select * from tab_b;
+create table tab_c as select * from tab_a;"""
+    helper(sql, {"tab_b"}, {"tab_c"})
+
+
 def test_alter_table_rename():
     helper("alter table tab1 rename to tab2;", None, None)