Repository: macbre/sql-metadata

Issue:
missing tables returned by sql_metadata.get_query_tables
Hello,

I have 80 long queries, and for 27 of them, results are not correct for "sql_metadata.get_query_tables(ddl)"

Here is an example:

ddl_for_sql_metadata.sql
```
SELECT
  "attr1"
FROM
  (
   SELECT
     "attr2"
   FROM
     (database1.table1 "aliasTable"
   LEFT JOIN (
      SELECT
        "attr3"
      FROM
        ((
         SELECT
           "attr4"
         FROM
           database2.table2
         GROUP BY 1, 2, 3, 4
      )  "aliasTable2"
      LEFT JOIN (
         SELECT
           "attr5"
         FROM
           database3.table3
         GROUP BY 1, 2
      )  "X" ON ("table"."attr" = "table"."attr"))
   )  "Y" ON ("table"."attr" = "table"."attr"))
)  "Z"
WHERE (myCondition)
GROUP BY 1, 2, 3, 4, 5
```

Quick way to test:
```
import sql_metadata

def read_from_file(file_path):
    return open(file_path, "r").read()

ddl = read_from_file("ddl_for_sql_metadata.sql")
tables = sql_metadata.get_query_tables(ddl)
print(tables)
```

It returns only ['database2.table2', 'database3.table3']
and it should return database1.table1



Discussion:

Patch:
diff --git a/sql_metadata.py b/sql_metadata.py
index 46aa30b..4a754ed 100644
--- a/sql_metadata.py
+++ b/sql_metadata.py
@@ -171,7 +171,7 @@ def get_query_columns(query: str) -> List[str]:
     return unique(columns)
 
 
-def _get_token_normalized_value(token: str) -> str:
+def _get_token_normalized_value(token: sqlparse.sql.Token) -> str:
     return token.value.translate(str.maketrans("", "", " \n\t\r")).upper()
 
 
@@ -208,7 +208,7 @@ def _update_table_names(
             "UPDATE",
             "TABLE",
         ]
-        and last_token not in ["AS"]
+        and last_token not in ["AS", "WITH"]
         and token.value not in ["AS", "SELECT"]
     ):
         if last_token == "." and next_token != ".":
@@ -289,16 +289,18 @@ def get_query_tables(query: str) -> List[str]:
     tokens = get_query_tokens(query)
 
     for index, token in enumerate(tokens):
-        # print([token, token.ttype, last_token, last_keyword])
-
         # remove whitespaces from token value and uppercase
         token_val_norm = _get_token_normalized_value(token)
+
+        # print([token, token_val_norm, token.ttype, last_keyword])
+
         if token.is_keyword and token_val_norm in table_syntax_keywords:
             # keep the name of the last keyword, the next one can be a table name
             last_keyword = token_val_norm
             # print('keyword', last_keyword)
-        elif str(token) == "(":
+        elif str(token) == "(" and last_keyword in ["INTO", "VALUES"]:
             # reset the last_keyword for INSERT `foo` VALUES(id, bar) ...
+            # reset the last_keyword for INSERT `foo` (col1, col2) VALUES(id, bar) ...
             last_keyword = None
         elif token.is_keyword and token_val_norm in ["FORCE", "ORDER", "GROUPBY"]:
             # reset the last_keyword for queries like:


Test patch:
diff --git a/test/test_hive.py b/test/test_hive.py
index 7833a98..0306852 100644
--- a/test/test_hive.py
+++ b/test/test_hive.py
@@ -1,6 +1,8 @@
 """
 Set of unit tests for handling of Apache Hive queries
 """
+import pytest
+
 from sql_metadata import get_query_columns, get_query_tables
 
 
@@ -17,6 +19,8 @@ def test_insert_overwrite_table():
 
 
 def test_complex_hive_query():
+    pytest.skip("Improve HIVE syntax handling with a new parser (#98)")
+
     # https://cwiki.apache.org/confluence/display/Hive/LanguageManual+DML#LanguageManualDML-InsertingdataintoHiveTablesfromqueries
     dag = """
 INSERT OVERWRITE TABLE foo_report
diff --git a/test/test_query.py b/test/test_query.py
index 5ba52eb..980bbb3 100644
--- a/test/test_query.py
+++ b/test/test_query.py
@@ -1,3 +1,5 @@
+import pytest
+
 from sql_metadata import (
     preprocess_query,
     get_query_tokens,
@@ -577,3 +579,41 @@ def test_unions():
     assert ["d", "g"] == get_query_tables(
         "SELECT a,b,c FROM d UNION ALL SELECT e,f FROM g"
     )
+
+
+def test_with_brackets():
+    assert ["database1.table1", "database2.table2"] == get_query_tables(
+        """
+        SELECT
+          "xxxxx"
+        FROM
+          (database1.table1 alias
+        LEFT JOIN database2.table2 ON ("tt"."ttt"."fff" = "xx"."xxx"))
+        """
+    )
+
+    assert ["inner_table"] == get_query_tables(
+        """
+        SELECT
+            t.foo
+        FROM
+            (SELECT foo FROM inner_table
+            WHERE bar = '1') t
+        """
+    )
+
+
+def test_with_with():
+    pytest.skip("Improve WITH syntax handling with a new parser (#98)")
+
+    assert ["table3", "database2.table2"] == get_query_tables(
+        """
+        WITH
+            database1.tableFromWith AS SELECT * FROM table3
+        SELECT
+          "xxxxx"
+        FROM
+          database1.tableFromWith alias
+        LEFT JOIN database2.table2 ON ("tt"."ttt"."fff" = "xx"."xxx")
+        """
+    )