fix: table dropped after creation should be left out of source or tar…

…get tables (#41)

* chore: test against Python 3.8

* fix: table dropped after creation should be left out of source or target tables

.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

sqllineage/core.py

@@ -23,7 +23,8 @@ def __init__(self, sql: str, encoding=None):
         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)

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"})
 

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py35,py36,py37
+envlist = py35,py36,py37,py38
 
 [testenv]
 deps =