Repository: reata/sqllineage

Issue:
comment in line raise AssertionError
```sql
select
*
from  -- comment
tab_a
```

This will raise AssertionError

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/reata/sqllineage/pull/38?src=pr&el=h1) Report
> Merging [#38](https://codecov.io/gh/reata/sqllineage/pull/38?src=pr&el=desc) into [master](https://codecov.io/gh/reata/sqllineage/commit/d79057f068fa8d2e3c8c121355eb9af2a026d899&el=desc) will **decrease** coverage by `1.36%`.
> The diff coverage is `100.00%`.

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

```diff
@@            Coverage Diff             @@
##           master      #38      +/-   ##
==========================================
- Coverage   79.79%   78.43%   -1.37%     
==========================================
  Files           1        1              
  Lines          99      102       +3     
==========================================
+ Hits           79       80       +1     
- Misses         20       22       +2     
```


| [Impacted Files](https://codecov.io/gh/reata/sqllineage/pull/38?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [sqllineage/core.py](https://codecov.io/gh/reata/sqllineage/pull/38/diff?src=pr&el=tree#diff-c3FsbGluZWFnZS9jb3JlLnB5) | `78.43% <100.00%> (-1.37%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/reata/sqllineage/pull/38?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/38?src=pr&el=footer). Last update [d79057f...2920258](https://codecov.io/gh/reata/sqllineage/pull/38?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 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"})