Repository: macbre/sql-metadata

Issue:
Tables with leading digits are not properly recognized
https://github.com/macbre/index-digest/runs/2567720826

```
ERROR    indexdigest.linters.linter_0006_not_used_columns_and_tables:linter_0006_not_used_columns_and_tables.py:91 Unable to extract tables and columns used from the query: SELECT test
ERROR    indexdigest.linters.linter_0006_not_used_columns_and_tables:linter_0006_not_used_columns_and_tables.py:91 Unable to extract tables and columns used from the query: SELECT t.val as value, count(*) FROM 0020_big_table as t WHERE id BETWEEN 10 AND 20 GROUP BY val
```

```python
Parser("SELECT t.val as value, count(*) FROM 0020_big_table as t WHERE id BETWEEN 10 AND 20 GROUP BY val").tables  # []
```

However:

```python
Parser("SELECT t.val as value, count(*) FROM `0020_big_table` as t WHERE id BETWEEN 10 AND 20 GROUP BY val").tables  # ['0020_big_table']
```

> **Identifiers may begin with a digit** but unless quoted may not consist solely of digits.
>
> https://dev.mysql.com/doc/refman/8.0/en/identifiers.html

Comments:

Patch:
diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py
index 887c505..10be534 100644
--- a/sql_metadata/parser.py
+++ b/sql_metadata/parser.py
@@ -5,8 +5,8 @@ import re
 from typing import Dict, List, Optional, Tuple, Union
 
 import sqlparse
-from sqlparse.sql import TokenList
-from sqlparse.tokens import Whitespace
+from sqlparse.sql import Token
+from sqlparse.tokens import Name, Number, Whitespace
 
 from sql_metadata.generalizator import Generalizator
 from sql_metadata.keywords_lists import (
@@ -18,7 +18,7 @@ from sql_metadata.keywords_lists import (
     TABLE_ADJUSTMENT_KEYWORDS,
     WITH_ENDING_KEYWORDS,
 )
-from sql_metadata.token import SQLToken, EmptyToken
+from sql_metadata.token import EmptyToken, SQLToken
 from sql_metadata.utils import UniqueList
 
 
@@ -60,6 +60,8 @@ class Parser:  # pylint: disable=R0902
         self._is_in_with_block = False
         self._with_columns_candidates = dict()
 
+        self.sqlparse_tokens = None
+
     @property
     def query(self) -> str:
         """
@@ -81,7 +83,8 @@ class Parser:  # pylint: disable=R0902
         if not parsed:
             return tokens
 
-        sqlparse_tokens = TokenList(parsed[0].tokens).flatten()
+        self.sqlparse_tokens = parsed[0].tokens
+        sqlparse_tokens = self._flatten_sqlparse()
         non_empty_tokens = [
             token for token in sqlparse_tokens if token.ttype is not Whitespace
         ]
@@ -744,13 +747,47 @@ class Parser:  # pylint: disable=R0902
         # 0. remove newlines
         query = self._raw_query.replace("\n", " ")
         # 1. remove quotes "
-        query = query.replace('"', "")
+        query = query.replace('"', "`")
 
         # 2. `database`.`table` notation -> database.table
         query = re.sub(r"`([^`]+)`\.`([^`]+)`", r"\1.\2", query)
 
         return query
 
+    def _flatten_sqlparse(self):
+        for token in self.sqlparse_tokens:
+            # sqlparse returns mysql digit starting identifiers as group
+            # check https://github.com/andialbrecht/sqlparse/issues/337
+            is_grouped_mysql_digit_name = (
+                token.is_group
+                and len(token.tokens) == 2
+                and token.tokens[0].ttype is Number.Integer
+                and (
+                    token.tokens[1].is_group and token.tokens[1].tokens[0].ttype is Name
+                )
+            )
+            if token.is_group and not is_grouped_mysql_digit_name:
+                yield from token.flatten()
+            elif is_grouped_mysql_digit_name:
+                # we have digit starting name
+                new_tok = Token(
+                    value=f"{token.tokens[0].normalized}"
+                    f"{token.tokens[1].tokens[0].normalized}",
+                    ttype=token.tokens[1].tokens[0].ttype,
+                )
+                new_tok.parent = token.parent
+                yield new_tok
+                if len(token.tokens[1].tokens) > 1:
+                    # unfortunately there might be nested groups
+                    remaining_tokens = token.tokens[1].tokens[1:]
+                    for tok in remaining_tokens:
+                        if tok.is_group:
+                            yield from tok.flatten()
+                        else:
+                            yield tok
+            else:
+                yield token
+
     @property
     def _is_create_table_query(self) -> bool:
         """


Test patch:
diff --git a/test/test_getting_columns.py b/test/test_getting_columns.py
index 26c0ddd..9ba60d4 100644
--- a/test/test_getting_columns.py
+++ b/test/test_getting_columns.py
@@ -250,3 +250,16 @@ def test_columns_and_sql_functions():
     assert Parser(
         "select count(col)+max(col2)+ min(col3)+ count(distinct  col4) + custom_func(col5) from dual"
     ).columns == ["col", "col2", "col3", "col4", "col5"]
+
+
+def test_columns_starting_with_keywords():
+    query = "SELECT `schema_name`, full_table_name, `column_name`, `catalog_name`, `table_name`, column_length, annotation FROM corporate.all_tables"
+    assert Parser(query).columns == [
+        "schema_name",
+        "full_table_name",
+        "column_name",
+        "catalog_name",
+        "table_name",
+        "column_length",
+        "annotation",
+    ]
diff --git a/test/test_getting_tables.py b/test/test_getting_tables.py
index e53744a..9657e64 100644
--- a/test/test_getting_tables.py
+++ b/test/test_getting_tables.py
@@ -432,9 +432,15 @@ def test_get_tables_with_leading_digits():
     assert ["0020_big_table"] == Parser(
         "SELECT t.val as value, count(*) FROM `0020_big_table`"
     ).tables
+    assert ["0020_big_table"] == Parser(
+        'SELECT t.val as value, count(*) FROM "0020_big_table" as t WHERE id BETWEEN 10 AND 20 GROUP BY val'
+    ).tables
     assert ["0020_big_table"] == Parser(
         "SELECT t.val as value, count(*) FROM 0020_big_table as t WHERE id BETWEEN 10 AND 20 GROUP BY val"
     ).tables
+    assert ["0020_big_table"] == Parser(
+        "SELECT t.val as value, count(*) FROM `0020_big_table` as t WHERE id BETWEEN 10 AND 20 GROUP BY val"
+    ).tables
     assert ["0020_big_table"] == Parser(
         "SELECT t.val as value, count(*) FROM 0020_big_table"
     ).tables
diff --git a/test/test_unions.py b/test/test_unions.py
new file mode 100644
index 0000000..f0506be
--- /dev/null
+++ b/test/test_unions.py
@@ -0,0 +1,81 @@
+from sql_metadata import Parser
+
+
+def test_union():
+    query = """
+    SELECT
+ACCOUNTING_ENTITY.VERSION as "accountingEntityVersion",
+ACCOUNTING_ENTITY.ACTIVE as "active",
+ACCOUNTING_ENTITY.CATEGORY as "category",
+ACCOUNTING_ENTITY.CREATION_DATE as "creationDate",
+ACCOUNTING_ENTITY.DESCRIPTION as "description",
+ACCOUNTING_ENTITY.ID as "accountingEntityId",
+ACCOUNTING_ENTITY.MINIMAL_REMAINDER as "minimalRemainder",
+ACCOUNTING_ENTITY.REMAINDER as "remainder",
+ACCOUNTING_ENTITY.SYSTEM_TYPE_ID as "aeSystemTypeId",
+ACCOUNTING_ENTITY.DATE_CREATION as "dateCreation",
+ACCOUNTING_ENTITY.DATE_LAST_MODIFICATION as "dateLastModification",
+ACCOUNTING_ENTITY.USER_CREATION as "userCreation",
+ACCOUNTING_ENTITY.USER_LAST_MODIFICATION as "userLastModification"
+FROM ACCOUNTING_ENTITY
+WHERE ACCOUNTING_ENTITY.ID IN (
+SELECT DPD.ACCOUNTING_ENTITY_ID AS "ACCOUNTINGENTITYID" FROM DEBT D
+INNER JOIN DUTY_PER_DEBT DPD ON DPD.DEBT_ID = D.ID
+INNER JOIN DECLARATION_V2 DV2 ON DV2.ID = D.DECLARATION_ID
+WHERE DV2.DECLARATION_REF = #MRNFORMOVEMENT
+UNION
+SELECT BX.ACCOUNTING_ENTITY_ID AS "ACCOUNTINGENTITYID" FROM BENELUX BX
+INNER JOIN DECLARATION_V2 DV2 ON DV2.ID = BX.DECLARATION_ID
+WHERE DV2.DECLARATION_REF = #MRNFORMOVEMENT
+UNION
+SELECT CA4D.ACCOUNTING_ENTITY_ID AS "ACCOUNTINGENTITYID" FROM RESERVATION R
+INNER JOIN CA4_RESERVATIONS_DECLARATION CA4D ON CA4D.ID = R.CA4_ID
+INNER JOIN DECLARATION_V2 DV2 ON DV2.ID = R.DECLARATION_ID
+WHERE DV2.DECLARATION_REF = #MRNFORMOVEMENT)
+    """
+
+    parser = Parser(query)
+    assert parser.tables == [
+        "ACCOUNTING_ENTITY",
+        "DEBT",
+        "DUTY_PER_DEBT",
+        "DECLARATION_V2",
+        "BENELUX",
+        "RESERVATION",
+        "CA4_RESERVATIONS_DECLARATION",
+    ]
+    assert parser.columns_dict == {
+        "join": [
+            "DUTY_PER_DEBT.DEBT_ID",
+            "DEBT.ID",
+            "DECLARATION_V2.ID",
+            "DEBT.DECLARATION_ID",
+            "BENELUX.DECLARATION_ID",
+            "CA4_RESERVATIONS_DECLARATION.ID",
+            "RESERVATION.CA4_ID",
+            "RESERVATION.DECLARATION_ID",
+        ],
+        "select": [
+            "ACCOUNTING_ENTITY.VERSION",
+            "ACCOUNTING_ENTITY.ACTIVE",
+            "ACCOUNTING_ENTITY.CATEGORY",
+            "ACCOUNTING_ENTITY.CREATION_DATE",
+            "ACCOUNTING_ENTITY.DESCRIPTION",
+            "ACCOUNTING_ENTITY.ID",
+            "ACCOUNTING_ENTITY.MINIMAL_REMAINDER",
+            "ACCOUNTING_ENTITY.REMAINDER",
+            "ACCOUNTING_ENTITY.SYSTEM_TYPE_ID",
+            "ACCOUNTING_ENTITY.DATE_CREATION",
+            "ACCOUNTING_ENTITY.DATE_LAST_MODIFICATION",
+            "ACCOUNTING_ENTITY.USER_CREATION",
+            "ACCOUNTING_ENTITY.USER_LAST_MODIFICATION",
+            "DUTY_PER_DEBT.ACCOUNTING_ENTITY_ID",
+            "BENELUX.ACCOUNTING_ENTITY_ID",
+            "CA4_RESERVATIONS_DECLARATION.ACCOUNTING_ENTITY_ID",
+        ],
+        "where": [
+            "ACCOUNTING_ENTITY.ID",
+            "DECLARATION_V2.DECLARATION_REF",
+            "#MRNFORMOVEMENT",
+        ],
+    }