Repository: reata/sqllineage

Issue:
drop table parsed as target table
```
$ sqllineage -e "DROP TABLE IF EXISTS tab1"
Statements(#): 1
Source Tables:
    
Target Tables:
    tab1
```

expect:
When a table is in target_tables, dropping table result in a removal from target_tables. Otherwise  the DROP DML doesn't affect the result.

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/reata/sqllineage/pull/24?src=pr&el=h1) Report
> Merging [#24](https://codecov.io/gh/reata/sqllineage/pull/24?src=pr&el=desc) into [master](https://codecov.io/gh/reata/sqllineage/commit/245a83548482fbf78f90a26cea97a534c61be2e5?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `100%`.

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

```diff
@@           Coverage Diff           @@
##           master      #24   +/-   ##
=======================================
  Coverage   79.79%   79.79%           
=======================================
  Files           1        1           
  Lines          99       99           
=======================================
  Hits           79       79           
  Misses         20       20
```


| [Impacted Files](https://codecov.io/gh/reata/sqllineage/pull/24?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [sqllineage/core.py](https://codecov.io/gh/reata/sqllineage/pull/24/diff?src=pr&el=tree#diff-c3FsbGluZWFnZS9jb3JlLnB5) | `79.79% <100%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/reata/sqllineage/pull/24?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/24?src=pr&el=footer). Last update [245a835...5fcd262](https://codecov.io/gh/reata/sqllineage/pull/24?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/sqllineage/core.py b/sqllineage/core.py
index 74c178d..2faa914 100644
--- a/sqllineage/core.py
+++ b/sqllineage/core.py
@@ -4,7 +4,7 @@ from typing import List, Set
 
 import sqlparse
 from sqlparse.sql import Function, Identifier, Parenthesis, Statement, TokenList
-from sqlparse.tokens import DDL, Keyword, Token
+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')
@@ -18,9 +18,9 @@ class LineageParser(object):
         self._encoding = encoding
         self._source_tables = set()
         self._target_tables = set()
-        self._stmt = sqlparse.parse(sql, self._encoding)
+        self._stmt = sqlparse.parse(sql.strip(), self._encoding)
         for stmt in self._stmt:
-            if stmt.token_first().ttype == DDL and stmt.token_first().normalized == "DROP":
+            if stmt.get_type() == "DROP":
                 self._target_tables -= {t.get_real_name() for t in stmt.tokens if isinstance(t, Identifier)}
             else:
                 self._extract_from_token(stmt)
@@ -76,7 +76,7 @@ Target Tables:
                     continue
                 else:
                     assert isinstance(sub_token, Identifier)
-                    if isinstance(sub_token.token_first(), Parenthesis):
+                    if isinstance(sub_token.token_first(skip_cm=True), Parenthesis):
                         # SELECT col1 FROM (SELECT col2 FROM tab1) dt, the subquery will be parsed as Identifier
                         # and this Identifier's get_real_name method would return alias name dt
                         # referring https://github.com/andialbrecht/sqlparse/issues/218 for further information
@@ -90,8 +90,8 @@ Target Tables:
                 elif isinstance(sub_token, Function):
                     # insert into tab (col1, col2), tab (col1, col2) will be parsed as Function
                     # referring https://github.com/andialbrecht/sqlparse/issues/483 for further information
-                    assert isinstance(sub_token.token_first(), Identifier)
-                    self._target_tables.add(sub_token.token_first().get_real_name())
+                    assert isinstance(sub_token.token_first(skip_cm=True), Identifier)
+                    self._target_tables.add(sub_token.token_first(skip_cm=True).get_real_name())
                     target_table_token_flag = False
                 else:
                     assert isinstance(sub_token, Identifier)


Test patch:
diff --git a/tests/test_others.py b/tests/test_others.py
index 6fc57c9..7e0a287 100644
--- a/tests/test_others.py
+++ b/tests/test_others.py
@@ -22,6 +22,11 @@ def test_drop():
     helper("DROP TABLE IF EXISTS tab1", None, None)
 
 
+def test_drop_with_comment():
+    helper("""--comment
+DROP TABLE IF EXISTS tab1""", None, None)
+
+
 def test_drop_after_create():
     helper("CREATE TABLE IF NOT EXISTS tab1 (col1 STRING);DROP TABLE IF EXISTS tab1", None, None)
 
@@ -33,3 +38,8 @@ def test_create_select():
 def test_split_statements():
     sql = "SELECT * FROM tab1; SELECT * FROM tab2;"
     assert len(LineageParser(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