Repository: reata/sqllineage (2019-08-12 ~ 2020-04-08)

problem_statement: comment in line raise AssertionError
```sql
select
*
from  -- comment
tab_a
```

This will raise AssertionError
date: 2019-08-12
patch: diff --git a/sqllineage/core.py b/sqllineage/core.py
index 231aa26..5b010e2 100644
--- a/sqllineage/core.py
+++ b/sqllineage/core.py
@@ -3,13 +3,13 @@ import sys
 from typing import List, Set
 
 import sqlparse
-from sqlparse.sql import Function, Identifier, Parenthesis, Statement, TokenList
+from sqlparse.sql import Comment, Function, Identifier, Parenthesis, Statement, TokenList
 from sqlparse.tokens import Keyword, Token
 
 SOURCE_TABLE_TOKENS = ('FROM', 'JOIN', 'INNER JOIN', 'LEFT JOIN', 'RIGHT JOIN', 'LEFT OUTER JOIN', 'RIGHT OUTER JOIN',
                        'FULL OUTER JOIN', 'CROSS JOIN')
 TARGET_TABLE_TOKENS = ('INTO', 'OVERWRITE', 'TABLE')
-TEMP_TABLE_TOKENS = ('WITH', )
+TEMP_TABLE_TOKENS = ('WITH',)
 
 
 class LineageParser(object):
@@ -53,7 +53,7 @@ Target Tables:
     def target_tables(self) -> Set[str]:
         return self._target_tables
 
-    def _extract_from_token(self, token: Token):
+    def _extract_from_token(self, token: Token) -> None:
         source_table_token_flag = target_table_token_flag = temp_table_token_flag = False
         for sub_token in token.tokens:
             if isinstance(sub_token, TokenList):
@@ -72,7 +72,7 @@ Target Tables:
                 self._target_tables.add(sub_token.get_alias())
                 continue
             if source_table_token_flag:
-                if sub_token.is_whitespace:
+                if self.__token_negligible_before_tablename(sub_token):
                     continue
                 else:
                     assert isinstance(sub_token, Identifier)
@@ -85,7 +85,7 @@ Target Tables:
                         self._source_tables.add(sub_token.get_real_name())
                     source_table_token_flag = False
             elif target_table_token_flag:
-                if sub_token.is_whitespace:
+                if self.__token_negligible_before_tablename(sub_token):
                     continue
                 elif isinstance(sub_token, Function):
                     # insert into tab (col1, col2), tab (col1, col2) will be parsed as Function
@@ -98,7 +98,7 @@ Target Tables:
                     self._target_tables.add(sub_token.get_real_name())
                     target_table_token_flag = False
             elif temp_table_token_flag:
-                if sub_token.is_whitespace:
+                if self.__token_negligible_before_tablename(sub_token):
                     continue
                 else:
                     assert isinstance(sub_token, Identifier)
@@ -107,6 +107,10 @@ Target Tables:
                     self._extract_from_token(sub_token)
                     temp_table_token_flag = False
 
+    @classmethod
+    def __token_negligible_before_tablename(cls, token: Token) -> bool:
+        return token.is_whitespace or isinstance(token, Comment)
+
 
 def main():
     parser = argparse.ArgumentParser(prog='sqllineage', description='SQL Lineage Parser.')

test_patch: diff --git a/tests/test_select.py b/tests/test_select.py
index 8837091..69b6133 100644
--- a/tests/test_select.py
+++ b/tests/test_select.py
@@ -30,6 +30,14 @@ def test_select_with_comment():
     helper("SELECT -- comment1\n col1 FROM tab1", {"tab1"})
 
 
+def test_select_with_comment_after_from():
+    helper("SELECT col1\nFROM  -- comment\ntab1", {"tab1"})
+
+
+def test_select_with_comment_after_join():
+    helper("select * from tab1 join --comment\ntab2 on tab1.x = tab2.x", {"tab1", "tab2"})
+
+
 def test_select_keyword_as_column_alias():
     # here `as` is the column alias
     helper("SELECT 1 `as` FROM tab1", {"tab1"})


----

problem_statement: white space in left join
```sql
select * from tab_a
left  join tab_b
on tab_a.x = tab_b.x
```
an extra whitespace in "left  join" make tab_b undetectable.
date: 2020-03-23
patch: diff --git a/sqllineage/core.py b/sqllineage/core.py
index 5b010e2..0d7c505 100644
--- a/sqllineage/core.py
+++ b/sqllineage/core.py
@@ -1,4 +1,5 @@
 import argparse
+import re
 import sys
 from typing import List, Set
 
@@ -6,8 +7,9 @@ import sqlparse
 from sqlparse.sql import Comment, Function, Identifier, Parenthesis, Statement, TokenList
 from sqlparse.tokens import Keyword, Token
 
-SOURCE_TABLE_TOKENS = ('FROM', 'JOIN', 'INNER JOIN', 'LEFT JOIN', 'RIGHT JOIN', 'LEFT OUTER JOIN', 'RIGHT OUTER JOIN',
-                       'FULL OUTER JOIN', 'CROSS JOIN')
+SOURCE_TABLE_TOKENS = (r'FROM',
+                       # 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')
 TEMP_TABLE_TOKENS = ('WITH',)
 
@@ -59,7 +61,7 @@ Target Tables:
             if isinstance(sub_token, TokenList):
                 self._extract_from_token(sub_token)
             if sub_token.ttype in Keyword:
-                if sub_token.normalized in SOURCE_TABLE_TOKENS:
+                if any(re.match(regex, sub_token.normalized) for regex in SOURCE_TABLE_TOKENS):
                     source_table_token_flag = True
                 elif sub_token.normalized in TARGET_TABLE_TOKENS:
                     target_table_token_flag = True

test_patch: diff --git a/tests/test_select.py b/tests/test_select.py
index 69b6133..8f31412 100644
--- a/tests/test_select.py
+++ b/tests/test_select.py
@@ -69,6 +69,10 @@ def test_select_left_join():
     helper("SELECT * FROM tab1 LEFT JOIN tab2", {"tab1", "tab2"})
 
 
+def test_select_left_join_with_extra_space_in_middle():
+    helper("SELECT * FROM tab1 LEFT  JOIN tab2", {"tab1", "tab2"})
+
+
 def test_select_left_semi_join():
     helper("SELECT * FROM tab1 LEFT SEMI JOIN tab2", {"tab1", "tab2"})
 


----

problem_statement: temp table checking
```sql
create table tab_a as select * from tab_b;
insert overwrite table tab_c select * from tab_a;
drop table tab_a;
```

Here tab_a is clearly a temp table, and yet it's marked as source table.
Possibly related to #29 #23 
date: 2020-04-08
patch: diff --git a/.travis.yml b/.travis.yml
index eb1e3f6..59abe34 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -4,6 +4,7 @@ python:
   - "3.5"
   - "3.6"
   - "3.7"
+  - "3.8"
 # command to install dependencies
 install:
   - pip install tox-travis codecov
diff --git a/sqllineage/core.py b/sqllineage/core.py
index 0d7c505..a46e388 100644
--- a/sqllineage/core.py
+++ b/sqllineage/core.py
@@ -23,7 +23,8 @@ class LineageParser(object):
         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._target_tables -= {t.get_real_name() for t in stmt.tokens if isinstance(t, Identifier)}
+                for tables in (self._source_tables, self._target_tables):
+                    tables -= {t.get_real_name() for t in stmt.tokens if isinstance(t, Identifier)}
             else:
                 self._extract_from_token(stmt)
         self._tmp_tables = self._source_tables.intersection(self._target_tables)
diff --git a/tox.ini b/tox.ini
index 0fa8e74..9546285 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py35,py36,py37
+envlist = py35,py36,py37,py38
 
 [testenv]
 deps =

test_patch: diff --git a/tests/test_others.py b/tests/test_others.py
index ddc0c30..eabb6d4 100644
--- a/tests/test_others.py
+++ b/tests/test_others.py
@@ -31,6 +31,13 @@ def test_drop_after_create():
     helper("CREATE TABLE IF NOT EXISTS tab1 (col1 STRING);DROP TABLE IF EXISTS tab1", None, None)
 
 
+def test_drop_tmp_tab_after_create():
+    sql = """create table tab_a as select * from tab_b;
+insert overwrite table tab_c select * from tab_a;
+drop table tab_a;"""
+    helper(sql, {"tab_b"}, {"tab_c"})
+
+
 def test_create_select():
     helper("CREATE TABLE tab1 SELECT * FROM tab2", {"tab2"}, {"tab1"})