Repository: reata/sqllineage

Issue:
statement granularity lineage result
In command line interface, we currently show file-wise lineage result. It causes problem like #23 . Some self-dependent sql shouldn't identify related table as temp table. This pattern could easily be identified if we look at statement rather than the whole file.
On the other hand, if a table was created/inserted into, and later queried from, it's perfectly reasonable to be identified as temp table.
So this statement grandularity lineage result may be the right choice also concerning #29 . 
We could try the operator reloading in design.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/reata/sqllineage/pull/54?src=pr&el=h1) Report
> Merging [#54](https://codecov.io/gh/reata/sqllineage/pull/54?src=pr&el=desc) into [master](https://codecov.io/gh/reata/sqllineage/commit/9e68f0ca25a1310cc7335002f6c57685fc815c2b&el=desc) will **decrease** coverage by `4.46%`.
> The diff coverage is `69.85%`.

[![Impacted file tree graph](https://codecov.io/gh/reata/sqllineage/pull/54/graphs/tree.svg?width=650&height=150&src=pr&token=latkaw2Lx6)](https://codecov.io/gh/reata/sqllineage/pull/54?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master      #54      +/-   ##
==========================================
- Coverage   85.05%   80.59%   -4.47%     
==========================================
  Files           3        6       +3     
  Lines         174      237      +63     
==========================================
+ Hits          148      191      +43     
- Misses         26       46      +20     
```


| [Impacted Files](https://codecov.io/gh/reata/sqllineage/pull/54?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [sqllineage/runner.py](https://codecov.io/gh/reata/sqllineage/pull/54/diff?src=pr&el=tree#diff-c3FsbGluZWFnZS9ydW5uZXIucHk=) | `49.18% <49.18%> (ø)` | |
| [sqllineage/core.py](https://codecov.io/gh/reata/sqllineage/pull/54/diff?src=pr&el=tree#diff-c3FsbGluZWFnZS9jb3JlLnB5) | `88.65% <81.81%> (+7.40%)` | :arrow_up: |
| [sqllineage/combiners.py](https://codecov.io/gh/reata/sqllineage/pull/54/diff?src=pr&el=tree#diff-c3FsbGluZWFnZS9jb21iaW5lcnMucHk=) | `93.54% <93.54%> (ø)` | |
| [sqllineage/\_\_init\_\_.py](https://codecov.io/gh/reata/sqllineage/pull/54/diff?src=pr&el=tree#diff-c3FsbGluZWFnZS9fX2luaXRfXy5weQ==) | `100.00% <0.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/reata/sqllineage/pull/54?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/reata/sqllineage/pull/54?src=pr&el=footer). Last update [9e68f0c...87b9b29](https://codecov.io/gh/reata/sqllineage/pull/54?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/setup.py b/setup.py
index 2426bc9..7e0720c 100644
--- a/setup.py
+++ b/setup.py
@@ -43,6 +43,6 @@ setup(
     ],
     python_requires=">=3.5",
     install_requires=["sqlparse>=0.3.0,<0.4"],
-    entry_points={"console_scripts": ["sqllineage = sqllineage.core:main"]},
+    entry_points={"console_scripts": ["sqllineage = sqllineage.runner:main"]},
     extras_require={"ci": ci_requires},
 )
diff --git a/sqllineage/combiners.py b/sqllineage/combiners.py
new file mode 100644
index 0000000..a652340
--- /dev/null
+++ b/sqllineage/combiners.py
@@ -0,0 +1,39 @@
+from functools import reduce
+from operator import add
+
+from sqllineage.core import LineageResult
+
+
+class LineageCombiner:
+    @staticmethod
+    def combine(*args: LineageResult) -> LineageResult:
+        raise NotImplementedError
+
+
+class NaiveLineageCombiner(LineageCombiner):
+    @staticmethod
+    def combine(*args: LineageResult) -> LineageResult:
+        return reduce(add, args, LineageResult())
+
+
+class DefaultLineageCombiner(LineageCombiner):
+    @staticmethod
+    def combine(*args: LineageResult) -> LineageResult:
+        combined_result = LineageResult()
+        for lineage_result in args:
+            if lineage_result.drop:
+                for st_tables in (combined_result.read, combined_result.write):
+                    st_tables -= lineage_result.drop
+            elif lineage_result.rename:
+                for (table_old, table_new) in lineage_result.rename:
+                    for st_tables in (combined_result.read, combined_result.write):
+                        if table_old in st_tables:
+                            st_tables.remove(table_old)
+                            st_tables.add(table_new)
+            else:
+                combined_result.read |= lineage_result.read
+                combined_result.write |= lineage_result.write
+        tmp_tables = combined_result.read.intersection(combined_result.write)
+        combined_result.read -= tmp_tables
+        combined_result.write -= tmp_tables
+        return combined_result
diff --git a/sqllineage/core.py b/sqllineage/core.py
index 3ebe4dc..47d3265 100644
--- a/sqllineage/core.py
+++ b/sqllineage/core.py
@@ -1,9 +1,6 @@
-import argparse
 import re
-import sys
-from typing import List, Set
+from typing import Set, TYPE_CHECKING, Tuple
 
-import sqlparse
 from sqlparse.sql import (
     Comment,
     Comparison,
@@ -23,64 +20,73 @@ SOURCE_TABLE_TOKENS = (
     # inspired by https://github.com/andialbrecht/sqlparse/blob/master/sqlparse/keywords.py
     r"((LEFT\s+|RIGHT\s+|FULL\s+)?(INNER\s+|OUTER\s+|STRAIGHT\s+)?|(CROSS\s+|NATURAL\s+)?)?JOIN",
 )
-TARGET_TABLE_TOKENS = ("INTO", "OVERWRITE", "TABLE")
+TARGET_TABLE_TOKENS = ("INTO", "OVERWRITE", "TABLE", "VIEW")
 TEMP_TABLE_TOKENS = ("WITH",)
 
 
-class LineageParser(object):
-    def __init__(self, sql: str, encoding=None):
-        self._encoding = encoding
-        self._source_tables = set()  # type: Set[Table]
-        self._target_tables = set()  # type: Set[Table]
-        self._stmt = [
-            s
-            for s in sqlparse.parse(sql.strip(), self._encoding)
-            if s.token_first(skip_cm=True)
-        ]
-        for stmt in self._stmt:
-            if stmt.get_type() == "DROP":
-                self._extract_from_DDL_DROP(stmt)
-            elif stmt.get_type() == "ALTER":
-                self._extract_from_DDL_ALTER(stmt)
-            elif (
-                stmt.get_type() == "DELETE"
-                or stmt.token_first(skip_cm=True).normalized == "TRUNCATE"
-            ):
-                pass
-            else:
-                # DML parsing logic also applies to CREATE DDL
-                self._extract_from_DML(stmt)
-        self._tmp_tables = self._source_tables.intersection(self._target_tables)
-        self._source_tables -= self._tmp_tables
-        self._target_tables -= self._tmp_tables
+class LineageResult:
+    """Statement(s) Level Lineage Result."""
+
+    __slots__ = ["read", "write", "rename", "drop"]
+    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]
+
+    def __init__(self) -> None:
+        for attr in self.__slots__:
+            setattr(self, attr, set())
+
+    def __add__(self, other):
+        lineage_result = LineageResult()
+        for attr in self.__slots__:
+            setattr(
+                lineage_result, attr, getattr(self, attr).union(getattr(other, attr))
+            )
+        return lineage_result
 
     def __str__(self):
-        return """Statements(#): {stmt_cnt}
-Source Tables:
-    {source_tables}
-Target Tables:
-    {target_tables}
-""".format(
-            stmt_cnt=len(self.statements),
-            source_tables="\n    ".join(str(t) for t in self.source_tables),
-            target_tables="\n    ".join(str(t) for t in self.target_tables),
+        return "\n".join(
+            "table {attr}: {attr_val}".format(
+                attr=attr, attr_val=getattr(self, attr) if getattr(self, attr) else "{}"
+            )
+            for attr in self.__slots__
         )
 
-    @property
-    def statements_parsed(self) -> List[Statement]:
-        return self._stmt
 
-    @property
-    def statements(self) -> List[str]:
-        return [sqlparse.format(s.value) for s in self.statements_parsed]
+class LineageAnalyzer:
+    """SQL Statement Level Lineage Analyzer."""
 
-    @property
-    def source_tables(self) -> Set[Table]:
-        return {t for t in self._source_tables}
+    def __init__(self) -> None:
+        self._lineage_result = LineageResult()
 
-    @property
-    def target_tables(self) -> Set[Table]:
-        return {t for t in self._target_tables}
+    def analyze(self, stmt: Statement) -> LineageResult:
+        if stmt.get_type() == "DROP":
+            self._extract_from_DDL_DROP(stmt)
+        elif stmt.get_type() == "ALTER":
+            self._extract_from_DDL_ALTER(stmt)
+        elif (
+            stmt.get_type() == "DELETE"
+            or stmt.token_first(skip_cm=True).normalized == "TRUNCATE"
+        ):
+            pass
+        else:
+            # DML parsing logic also applies to CREATE DDL
+            self._extract_from_DML(stmt)
+        return self._lineage_result
+
+    def _extract_from_DDL_DROP(self, stmt: Statement) -> None:
+        for table in {
+            Table.create(t) for t in stmt.tokens if isinstance(t, Identifier)
+        }:
+            self._lineage_result.drop.add(table)
+
+    def _extract_from_DDL_ALTER(self, stmt: Statement) -> None:
+        tables = [Table.create(t) for t in stmt.tokens if isinstance(t, Identifier)]
+        keywords = [t for t in stmt.tokens if t.ttype is Keyword]
+        if any(k.normalized == "RENAME" for k in keywords) and len(tables) == 2:
+            self._lineage_result.rename.add((tables[0], tables[1]))
 
     def _extract_from_DML(self, token: Token) -> None:
         source_table_token_flag = (
@@ -106,7 +112,7 @@ Target Tables:
                 and sub_token.get_alias() is not None
             ):
                 # overwrite can't be parsed as Keyword, manual walk around
-                self._target_tables.add(Table(sub_token.get_alias()))
+                self._lineage_result.write.add(Table(sub_token.get_alias()))
                 continue
             if source_table_token_flag:
                 if self.__token_negligible_before_tablename(sub_token):
@@ -120,7 +126,7 @@ Target Tables:
                         # referring https://github.com/andialbrecht/sqlparse/issues/218 for further information
                         pass
                     else:
-                        self._source_tables.add(Table.create(sub_token))
+                        self._lineage_result.read.add(Table.create(sub_token))
                     source_table_token_flag = False
             elif target_table_token_flag:
                 if self.__token_negligible_before_tablename(sub_token):
@@ -130,7 +136,7 @@ Target Tables:
                     # referring https://github.com/andialbrecht/sqlparse/issues/483 for further information
                     if not isinstance(sub_token.token_first(skip_cm=True), Identifier):
                         raise SQLLineageException("An Identifier is expected")
-                    self._target_tables.add(
+                    self._lineage_result.write.add(
                         Table.create(sub_token.token_first(skip_cm=True))
                     )
                 elif isinstance(sub_token, Comparison):
@@ -141,12 +147,12 @@ Target Tables:
                         and isinstance(sub_token.right, Identifier)
                     ):
                         raise SQLLineageException("An Identifier is expected")
-                    self._target_tables.add(Table.create(sub_token.left))
-                    self._source_tables.add(Table.create(sub_token.right))
+                    self._lineage_result.write.add(Table.create(sub_token.left))
+                    self._lineage_result.read.add(Table.create(sub_token.right))
                 else:
                     if not isinstance(sub_token, Identifier):
                         raise SQLLineageException("An Identifier is expected")
-                    self._target_tables.add(Table.create(sub_token))
+                    self._lineage_result.write.add(Table.create(sub_token))
                 target_table_token_flag = False
             elif temp_table_token_flag:
                 if self.__token_negligible_before_tablename(sub_token):
@@ -154,62 +160,11 @@ Target Tables:
                 else:
                     if not isinstance(sub_token, Identifier):
                         raise SQLLineageException("An Identifier is expected")
-                    self._source_tables.add(Table.create(sub_token))
-                    self._target_tables.add(Table.create(sub_token))
+                    self._lineage_result.read.add(Table.create(sub_token))
+                    self._lineage_result.write.add(Table.create(sub_token))
                     self._extract_from_DML(sub_token)
                     temp_table_token_flag = False
 
-    def _extract_from_DDL_DROP(self, stmt: Statement) -> None:
-        for st_tables in (self._source_tables, self._target_tables):
-            st_tables -= {
-                Table.create(t) for t in stmt.tokens if isinstance(t, Identifier)
-            }
-
-    def _extract_from_DDL_ALTER(self, stmt: Statement) -> None:
-        tables = [Table.create(t) for t in stmt.tokens if isinstance(t, Identifier)]
-        keywords = [t for t in stmt.tokens if t.ttype is Keyword]
-        if any(k.normalized == "RENAME" for k in keywords) and len(tables) == 2:
-            for st_tables in (self._source_tables, self._target_tables):
-                if tables[0] in st_tables:
-                    st_tables.remove(tables[0])
-                    st_tables.add(tables[1])
-
     @classmethod
     def __token_negligible_before_tablename(cls, token: Token) -> bool:
         return token.is_whitespace or isinstance(token, Comment)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        prog="sqllineage", description="SQL Lineage Parser."
-    )
-    parser.add_argument(
-        "-e", metavar="<quoted-query-string>", help="SQL from command line"
-    )
-    parser.add_argument("-f", metavar="<filename>", help="SQL from files")
-    args = parser.parse_args()
-    if args.e and args.f:
-        print(
-            "WARNING: Both -e and -f options are specified. -e option will be ignored",
-            file=sys.stderr,
-        )
-    if args.f:
-        try:
-            with open(args.f) as f:
-                sql = f.read()
-            print(LineageParser(sql))
-        except FileNotFoundError:
-            print("ERROR: No such file: {}".format(args.f), file=sys.stderr)
-        except PermissionError:
-            print(
-                "ERROR: Permission denied when reading file '{}'".format(args.f),
-                file=sys.stderr,
-            )
-    elif args.e:
-        print(LineageParser(args.e))
-    else:
-        parser.print_help()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/sqllineage/runner.py b/sqllineage/runner.py
new file mode 100644
index 0000000..643bfc4
--- /dev/null
+++ b/sqllineage/runner.py
@@ -0,0 +1,113 @@
+import argparse
+import sys
+from typing import List, Set, Type
+
+import sqlparse
+from sqlparse.sql import Statement
+
+from sqllineage.combiners import DefaultLineageCombiner, LineageCombiner
+from sqllineage.core import LineageAnalyzer
+from sqllineage.models import Table
+
+
+class LineageRunner(object):
+    def __init__(
+        self,
+        sql: str,
+        encoding: str = None,
+        combiner: Type[LineageCombiner] = DefaultLineageCombiner,
+        verbose: bool = False,
+    ):
+        self._encoding = encoding
+        self._stmt = [
+            s
+            for s in sqlparse.parse(sql.strip(), self._encoding)
+            if s.token_first(skip_cm=True)
+        ]
+        self._lineage_results = [LineageAnalyzer().analyze(stmt) for stmt in self._stmt]
+        self._lineage_result = combiner.combine(*self._lineage_results)
+        self._verbose = verbose
+
+    def __str__(self):
+        statements = self.statements
+        combined = """Statements(#): {stmt_cnt}
+Source Tables:
+    {source_tables}
+Target Tables:
+    {target_tables}
+""".format(
+            stmt_cnt=len(statements),
+            source_tables="\n    ".join(str(t) for t in self.source_tables),
+            target_tables="\n    ".join(str(t) for t in self.target_tables),
+        )
+        if self._verbose:
+            result = ""
+            for i, lineage_result in enumerate(self._lineage_results):
+                stmt_short = statements[i].replace("\n", "")
+                if len(stmt_short) > 50:
+                    stmt_short = stmt_short[:50] + "..."
+                result += """Statement #{ord}: {stmt}
+{content}
+""".format(
+                    ord=i + 1, content=str(lineage_result), stmt=stmt_short
+                )
+            combined = result + "==========\nSummary:\n" + combined
+        return combined
+
+    @property
+    def statements_parsed(self) -> List[Statement]:
+        return self._stmt
+
+    @property
+    def statements(self) -> List[str]:
+        return [sqlparse.format(s.value) for s in self.statements_parsed]
+
+    @property
+    def source_tables(self) -> Set[Table]:
+        return {t for t in self._lineage_result.read}
+
+    @property
+    def target_tables(self) -> Set[Table]:
+        return {t for t in self._lineage_result.write}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="sqllineage", description="SQL Lineage Parser."
+    )
+    parser.add_argument(
+        "-e", metavar="<quoted-query-string>", help="SQL from command line"
+    )
+    parser.add_argument("-f", metavar="<filename>", help="SQL from files")
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        help="increase output verbosity, show statement level lineage result",
+        action="store_true",
+    )
+    args = parser.parse_args()
+    if args.e and args.f:
+        print(
+            "WARNING: Both -e and -f options are specified. -e option will be ignored",
+            file=sys.stderr,
+        )
+    if args.f:
+        try:
+            with open(args.f) as f:
+                sql = f.read()
+            print(LineageRunner(sql, verbose=args.verbose))
+        except FileNotFoundError:
+            print("ERROR: No such file: {}".format(args.f), file=sys.stderr)
+        except PermissionError:
+            print(
+                "ERROR: Permission denied when reading file '{}'".format(args.f),
+                file=sys.stderr,
+            )
+    elif args.e:
+        print(LineageRunner(args.e, verbose=args.verbose))
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()


Test patch:
diff --git a/tests/helpers.py b/tests/helpers.py
index 8ea1990..dc70c13 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -1,9 +1,9 @@
-from sqllineage.core import LineageParser
 from sqllineage.models import Table
+from sqllineage.runner import LineageRunner
 
 
 def helper(sql, source_tables=None, target_tables=None):
-    lp = LineageParser(sql)
+    lp = LineageRunner(sql)
     assert lp.source_tables == (
         set() if source_tables is None else {Table(t) for t in source_tables}
     )
diff --git a/tests/test_combiner.py b/tests/test_combiner.py
new file mode 100644
index 0000000..64acd2a
--- /dev/null
+++ b/tests/test_combiner.py
@@ -0,0 +1,9 @@
+import pytest
+
+from sqllineage.combiners import LineageCombiner, NaiveLineageCombiner
+
+
+def test_dummy():
+    assert NaiveLineageCombiner.combine()
+    with pytest.raises(NotImplementedError):
+        LineageCombiner.combine()
diff --git a/tests/test_core.py b/tests/test_core.py
new file mode 100644
index 0000000..a91e485
--- /dev/null
+++ b/tests/test_core.py
@@ -0,0 +1,6 @@
+from sqllineage.core import LineageResult
+
+
+def test_dummy():
+    assert str(LineageResult())
+    assert LineageResult() + LineageResult() is not None
diff --git a/tests/test_exception.py b/tests/test_exception.py
index 248716b..ce56451 100644
--- a/tests/test_exception.py
+++ b/tests/test_exception.py
@@ -1,21 +1,21 @@
 import pytest
 
-from sqllineage.core import LineageParser
 from sqllineage.exceptions import SQLLineageException
+from sqllineage.runner import LineageRunner
 
 
 def test_select_without_table():
     with pytest.raises(SQLLineageException):
-        LineageParser("select * from where foo='bar'")
+        LineageRunner("select * from where foo='bar'")
 
 
 def test_insert_without_table():
     with pytest.raises(SQLLineageException):
-        LineageParser("insert into select * from foo")
+        LineageRunner("insert into select * from foo")
 
 
 def test_with_cte_without_table():
     with pytest.raises(SQLLineageException):
-        LineageParser(
+        LineageRunner(
             "with as select * from foo insert into table bar select * from foo"
         )
diff --git a/tests/test_others.py b/tests/test_others.py
index b929099..946bf4c 100644
--- a/tests/test_others.py
+++ b/tests/test_others.py
@@ -1,4 +1,4 @@
-from sqllineage.core import LineageParser
+from sqllineage.runner import LineageRunner
 from .helpers import helper
 
 
@@ -95,30 +95,30 @@ def test_delete_from_table():
 
 def test_split_statements():
     sql = "SELECT * FROM tab1; SELECT * FROM tab2;"
-    assert len(LineageParser(sql).statements) == 2
+    assert len(LineageRunner(sql).statements) == 2
 
 
 def test_split_statements_with_heading_and_ending_new_line():
     sql = "\nSELECT * FROM tab1;\nSELECT * FROM tab2;\n"
-    assert len(LineageParser(sql).statements) == 2
+    assert len(LineageRunner(sql).statements) == 2
 
 
 def test_split_statements_with_comment():
     sql = """SELECT 1;
 
 -- SELECT 2;"""
-    assert len(LineageParser(sql).statements) == 1
+    assert len(LineageRunner(sql).statements) == 1
 
 
 def test_split_statements_with_show_create_table():
     sql = """SELECT 1;
 
 SHOW CREATE TABLE tab1;"""
-    assert len(LineageParser(sql).statements) == 2
+    assert len(LineageRunner(sql).statements) == 2
 
 
 def test_split_statements_with_desc():
     sql = """SELECT 1;
 
 DESC tab1;"""
-    assert len(LineageParser(sql).statements) == 2
+    assert len(LineageRunner(sql).statements) == 2
diff --git a/tests/test_runner.py b/tests/test_runner.py
new file mode 100644
index 0000000..ecae7d0
--- /dev/null
+++ b/tests/test_runner.py
@@ -0,0 +1,14 @@
+import pytest
+
+from sqllineage.runner import LineageRunner, main
+
+
+def test_dummy():
+    assert str(
+        LineageRunner(
+            "select col1, col2, col3, col4, col5, col6 from tab1;", verbose=True
+        )
+    )
+    with pytest.raises(SystemExit) as e:
+        main()
+    assert e.value.code == 2