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

Discussion:

Patch:
diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py
index 887c505..2ad3e07 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
 
 
@@ -59,6 +59,9 @@ class Parser:  # pylint: disable=R0902
         self._is_in_nested_function = False
         self._is_in_with_block = False
         self._with_columns_candidates = dict()
+        self._column_aliases_max_subquery_level = dict()
+
+        self.sqlparse_tokens = None
 
     @property
     def query(self) -> str:
@@ -81,7 +84,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
         ]
@@ -145,10 +149,15 @@ class Parser:  # pylint: disable=R0902
                 if (
                     token.last_keyword_normalized in KEYWORDS_BEFORE_COLUMNS
                     and token.previous_token.normalized not in ["AS", ")"]
-                    and token.previous_token.table_prefixed_column(tables_aliases)
-                    not in columns
-                    and token.left_expanded not in self.columns_aliases_names
+                    and not token.is_alias_without_as
+                    and (
+                        token.left_expanded not in self.columns_aliases_names
+                        or token.token_is_alias_of_self_not_from_subquery(
+                            aliases_levels=self._column_aliases_max_subquery_level
+                        )
+                    )
                 ):
+
                     if (
                         token.normalized not in FUNCTIONS_IGNORED
                         and not (
@@ -331,6 +340,13 @@ class Parser:  # pylint: disable=R0902
                 ) and token.value not in with_names + subqueries_names:
                     alias = token.left_expanded
                     column_aliases_names.append(alias)
+                    current_level = self._column_aliases_max_subquery_level.setdefault(
+                        alias, 0
+                    )
+                    if token.subquery_level > current_level:
+                        self._column_aliases_max_subquery_level[
+                            alias
+                        ] = token.subquery_level
 
         self._columns_aliases_names = column_aliases_names
         return self._columns_aliases_names
@@ -726,7 +742,7 @@ class Parser:  # pylint: disable=R0902
         Returns a list of columns between tw tokens
         """
         loop_token = start_token
-        aliases = []
+        aliases = UniqueList()
         while loop_token.next_token != end_token:
             if loop_token.next_token.left_expanded in self._aliases_to_check:
                 alias_token = loop_token.next_token
@@ -744,13 +760,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:
         """
diff --git a/sql_metadata/token.py b/sql_metadata/token.py
index e59e6c0..8f8a044 100644
--- a/sql_metadata/token.py
+++ b/sql_metadata/token.py
@@ -161,6 +161,22 @@ class SQLToken:  # pylint: disable=R0902
             and not self.previous_token.is_comment
         )
 
+    @property
+    def is_alias_of_self(self) -> bool:
+        """
+        Checks if a given token is an alias but at the same time
+        is also an alias of self, so not really an alias
+        """
+
+        end_of_column = self.find_nearest_token(
+            [",", "FROM"], value_attribute="normalized", direction="right"
+        )
+        while end_of_column.is_in_nested_function:
+            end_of_column = end_of_column.find_nearest_token(
+                [",", "FROM"], value_attribute="normalized", direction="right"
+            )
+        return end_of_column.previous_token.normalized == self.normalized
+
     @property
     def is_in_with_columns(self) -> bool:
         """
@@ -171,6 +187,17 @@ class SQLToken:  # pylint: disable=R0902
             and self.find_nearest_token(")", direction="right").is_with_columns_end
         )
 
+    def token_is_alias_of_self_not_from_subquery(self, aliases_levels: Dict) -> bool:
+        """
+        Checks if token is also an alias, but is an alias of self that is not
+        coming from a subquery, that means it's a valid column
+        """
+        return (
+            self.last_keyword_normalized == "SELECT"
+            and self.is_alias_of_self
+            and self.subquery_level == aliases_levels[self.value]
+        )
+
     def table_prefixed_column(self, table_aliases: Dict) -> str:
         """
         Substitutes table alias with actual table name


Test patch:
diff --git a/test/test_getting_columns.py b/test/test_getting_columns.py
index 26c0ddd..6fd4d8b 100644
--- a/test/test_getting_columns.py
+++ b/test/test_getting_columns.py
@@ -250,3 +250,103 @@ 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, column_weight, annotation 
+    FROM corporate.all_tables
+    """
+    parser = Parser(query)
+    assert parser.columns == [
+        "schema_name",
+        "full_table_name",
+        "column_name",
+        "catalog_name",
+        "table_name",
+        "column_length",
+        "column_weight",
+        "annotation",
+    ]
+
+
+def test_columns_with_keywords_parts():
+    query = """
+    SELECT column_length, column_weight, table_random, drop_20, create_table
+    FROM sample_table
+    """
+    assert Parser(query).columns == [
+        "column_length",
+        "column_weight",
+        "table_random",
+        "drop_20",
+        "create_table",
+    ]
+
+
+def test_columns_with_complex_aliases_same_as_columns():
+    query = """
+    select targetingtype, sellerid, sguid, 'd01' as datetype, adgroupname, targeting, 
+    customersearchterm, 
+    'product_search_term' as `type`, 
+    sum(impressions) as impr, 
+    sum(clicks) as clicks, 
+    sum(seventotalunits) as sold, 
+    sum(sevenadvertisedskuunits) as advertisedskuunits, 
+    sum(sevenotherskuunits) as otherskuunits, 
+    sum(sevendaytotalsales) as totalsales, 
+    round(sum(spend), 4) as spend, if(sum(impressions) > 0, 
+    round(sum(clicks)/sum(impressions), 4), 0) as ctr, 
+    if(sum(clicks) > 0, round(sum(seventotalunits)/sum(clicks), 4), 0) as cr, 
+    if(sum(clicks) > 0, round(sum(spend)/sum(clicks), 2), 0) as cpc 
+    from amazon_pl.search_term_report_impala 
+    where reportday >= to_date('2021-05-16 00:00:00.0') 
+    and reportday <= to_date('2021-05-16 00:00:00.0') 
+    and targetingtype in ('auto','manual') 
+    and sguid is not null and sguid != '' 
+    group by targetingtype,sellerid,sguid,adgroupname,targeting,customersearchterm 
+    order by impr desc
+    """
+    parser = Parser(query)
+    assert parser.columns == [
+        "targetingtype",
+        "sellerid",
+        "sguid",
+        "adgroupname",
+        "targeting",
+        "customersearchterm",
+        "impressions",
+        "clicks",
+        "seventotalunits",
+        "sevenadvertisedskuunits",
+        "sevenotherskuunits",
+        "sevendaytotalsales",
+        "spend",
+        "reportday",
+    ]
+
+
+def test_columns_with_aliases_same_as_columns():
+    query = """
+    select 
+    round(sum(impressions),1) as impressions, 
+    sum(clicks) as clicks
+    from amazon_pl.search_term_report_impala 
+    """
+    parser = Parser(query)
+    assert parser.columns == ["impressions", "clicks"]
+    assert parser.columns_aliases == {}
+
+    query = """
+    select
+    if(sum(clicks) > 0, round(sum(seventotalunits)/sum(clicks), 4), 0) as clicks, 
+    if(sum(clicks) > 0, round(sum(spend)/sum(clicks), 2), 0) as cpc 
+    from amazon_pl.search_term_report_impala 
+    """
+    parser = Parser(query)
+    assert parser.columns == ["clicks", "seventotalunits", "spend"]
+    assert parser.columns_aliases == {
+        "clicks": ["clicks", "seventotalunits"],
+        "cpc": ["clicks", "spend"],
+    }
diff --git a/test/test_getting_tables.py b/test/test_getting_tables.py
index c9bcc16..9657e64 100644
--- a/test/test_getting_tables.py
+++ b/test/test_getting_tables.py
@@ -418,3 +418,29 @@ with ur
         "zekerheid_accijnsbeweging": "accounting_entity.remainder",
         "zekerheid_douanevervoer": "accounting_entity.remainder",
     }
+
+
+def test_get_tables_with_leading_digits():
+    # see #139
+
+    # Identifiers may begin with a digit but unless quoted may not consist solely of digits.
+    assert ["0020"] == Parser("SELECT * FROM `0020`").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
+    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",
+        ],
+    }