fix: left join with extra spaces in the middle confuses parser (#40)

sqllineage/core.py

@@ -1,13 +1,15 @@
 import argparse
+import re
 import sys
 from typing import List, Set
 
 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 @@ def _extract_from_token(self, token: Token) -> None:
             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

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