Merge pull request #51 from mindsdb/identifier-tostring-fix

Keep escape symbols in identifier parts when dump to string

Author: StpMax

mindsdb_sql_parser/ast/select/identifier.py

@@ -1,6 +1,6 @@
 import re
 from copy import copy, deepcopy
-from typing import List
+from typing import List, Optional
 
 from mindsdb_sql_parser.ast.base import ASTNode
 from mindsdb_sql_parser.utils import indent
@@ -27,19 +27,28 @@ def path_str_to_parts(path_str: str):
 }
 
 
-def get_reserved_words():
-    from mindsdb_sql_parser.lexer import MindsDBLexer
+_reserved_keywords: set[str] = None
 
-    reserved = RESERVED_KEYWORDS
-    for word in MindsDBLexer.tokens:
-        if '_' not in word:
-            # exclude combinations
-            reserved.add(word)
-    return reserved
+
+def get_reserved_words() -> set[str]:
+    global _reserved_keywords
+
+    if _reserved_keywords is None:
+        from mindsdb_sql_parser.lexer import MindsDBLexer
+
+        _reserved_keywords = RESERVED_KEYWORDS
+        for word in MindsDBLexer.tokens:
+            if '_' not in word:
+                # exclude combinations
+                _reserved_keywords.add(word)
+    return _reserved_keywords
 
 
 class Identifier(ASTNode):
-    def __init__(self, path_str=None, parts=None, is_outer=False, with_rollup=False, *args, **kwargs):
+    def __init__(
+            self, path_str=None, parts=None, is_outer=False, with_rollup=False,
+            is_quoted: Optional[List[bool]] = None, *args, **kwargs
+    ):
         super().__init__(*args, **kwargs)
         assert path_str or parts, "Either path_str or parts must be provided for an Identifier"
         assert not (path_str and parts), "Provide either path_str or parts, but not both"
@@ -48,7 +57,7 @@ def __init__(self, path_str=None, parts=None, is_outer=False, with_rollup=False,
 
         if path_str and not parts:
             parts, is_quoted = path_str_to_parts(path_str)
-        else:
+        elif is_quoted is None:
             is_quoted = [False] * len(parts)
         assert isinstance(parts, list)
         self.parts = parts
@@ -63,22 +72,26 @@ def from_path_str(self, value, *args, **kwargs):
         parts, _ = path_str_to_parts(value)
         return Identifier(parts=parts, *args, **kwargs)
 
-    def parts_to_str(self):
-        out_parts = []
+    def append(self, other: "Identifier") -> None:
+        self.parts += other.parts
+        self.is_quoted += other.is_quoted
+
+    def iter_parts_str(self):
         reserved_words = get_reserved_words()
-        for part in self.parts:
+        for part, is_quoted in zip(self.parts, self.is_quoted):
             if isinstance(part, Star):
                 part = str(part)
             else:
                 if (
-                    not no_wrap_identifier_regex.fullmatch(part)
-                    or
-                    part.upper() in reserved_words
+                    is_quoted
+                    or not no_wrap_identifier_regex.fullmatch(part)
+                    or part.upper() in reserved_words
                 ):
                     part = f'`{part}`'
+            yield part
 
-            out_parts.append(part)
-        return '.'.join(out_parts)
+    def parts_to_str(self):
+        return '.'.join(self.iter_parts_str())
 
     def to_tree(self, *args, level=0, **kwargs):
         alias_str = f', alias={self.alias.to_tree()}' if self.alias else ''

mindsdb_sql_parser/ast/show.py

@@ -53,8 +53,8 @@ def get_string(self, *args, **kwargs):
         from_str = ''
         if self.from_table:
             ar = [
-                f'FROM {i}'
-                for i in self.from_table.parts
+                f'FROM {part}'
+                for part in self.from_table.iter_parts_str()
             ]
             ar.reverse()
             from_str = ' ' + ' '.join(ar)

mindsdb_sql_parser/parser.py

@@ -538,7 +538,7 @@ def show(self, p):
         value0 = command.from_table
         value1 = p.identifier
         if value0 is not None:
-            value1.parts = value1.parts + value0.parts
+            value1.append(value0)
 
         command.from_table = value1
         return command
@@ -549,7 +549,7 @@ def show(self, p):
         value0 = command.in_table
         value1 = p.identifier
         if value0 is not None:
-            value1.parts = value1.parts + value0.parts
+            value1.append(value0)
 
         command.in_table = value1
         return command
@@ -1824,17 +1824,21 @@ def json_value(self, p):
        'identifier DOT star')
     def identifier(self, p):
         node = p[0]
-        is_quoted = False
         if isinstance(p[2], Star):
             node.parts.append(p[2])
+            node.is_quoted.append(False)
         elif isinstance(p[2], int):
             node.parts.append(str(p[2]))
+            node.is_quoted.append(False)
         elif isinstance(p[2], str):
             node.parts.append(p[2])
+            node.is_quoted.append(False)
+        elif isinstance(p[2], Identifier):
+            node.append(p[2])
         else:
-            node.parts += p[2].parts
-            is_quoted = p[2].is_quoted[0]
-        node.is_quoted.append(is_quoted)
+            # fallback, shouldn't happen
+            node.parts.append(str(p[2]))
+            node.is_quoted.append(False)
         return node
 
     @_('quote_string',

tests/test_base_sql/test_ast.py

@@ -31,3 +31,21 @@ def test_identifier_deepcopy_is_quoted(self):
         ident = Identifier('`a`')
         ident2 = deepcopy(ident)
         assert ident2.is_quoted == [True]
+
+    def test_identifier_to_string(self):
+        test_cases = [
+            'test',
+            'Test',
+            'TEST',
+            '`test`',
+            '`Test`',
+            '`TEST`'
+        ]
+
+        for test_case in test_cases:
+            assert Identifier(test_case).to_string() == test_case
+
+        for i in range(len(test_cases)):
+            for test_case in test_cases:
+                test_str = f'{test_case}.{test_cases[i]}'
+                assert Identifier(test_str).to_string() == test_str

tests/test_base_sql/test_select_structure.py

@@ -726,9 +726,10 @@ def test_backticks(self):
         sql = "SELECT `name`, `status` FROM `mindsdb`.`wow stuff predictors`.`even-dashes-work`.`nice`"
         ast = parse_sql(sql)
 
-        expected_ast = Select(targets=[Identifier(parts=['name']), Identifier(parts=['status'])],
-                              from_table=Identifier(parts=['mindsdb', 'wow stuff predictors', 'even-dashes-work', 'nice']),
-                              )
+        expected_ast = Select(
+            targets=[Identifier('`name`'), Identifier('`status`')],
+            from_table=Identifier('`mindsdb`.`wow stuff predictors`.`even-dashes-work`.`nice`'),
+        )
 
         assert ast.to_tree() == expected_ast.to_tree()
         assert str(ast) == str(expected_ast)
@@ -1224,10 +1225,10 @@ def test_double_quote_render_skip(self):
         sql = 'select `KEY_ID`, `a`.* from `Table1` where `id`=2'
 
         expected_ast = Select(
-            targets=[Identifier('KEY_ID'), Identifier(parts=['a', Star()])],
-            from_table=Identifier(parts=['Table1']),
+            targets=[Identifier('`KEY_ID`'), Identifier(parts=['a', Star()], is_quoted=[True, False])],
+            from_table=Identifier(parts=['Table1'], is_quoted=[True]),
             where=BinaryOperation(op='=', args=[
-                Identifier('id'), Constant(2)
+                Identifier('`id`'), Constant(2)
             ])
         )
 

tests/test_base_sql/test_show.py

@@ -88,12 +88,12 @@ def test_from_where(self):
         assert ast.to_tree() == expected_ast.to_tree()
 
     def test_full_columns(self):
-        sql = "SHOW FULL COLUMNS FROM `concrete` FROM `files`"
+        sql = "SHOW FULL COLUMNS FROM `ccc` FROM `fff`"
         ast = parse_sql(sql)
         expected_ast = Show(
             category='COLUMNS',
             modes=['FULL'],
-            from_table=Identifier('files.concrete')
+            from_table=Identifier('`fff`.`ccc`')
         )
 
         assert str(ast) == str(expected_ast)

tests/test_mindsdb/test_create_predictor.py

@@ -124,7 +124,7 @@ def test_create_predictor_quotes(self):
         ast = parse_sql(sql)
         expected_ast = CreatePredictor(
             name=Identifier('xxx'),
-            integration_name=Identifier('yyy'),
+            integration_name=Identifier('`yyy`'),
             query_str="SELECT * FROM zzz",
             targets=[Identifier('sss')],
         )

tests/test_mysql/test_mysql_parser.py

@@ -42,23 +42,23 @@ def test_select_varialbe_complex(self):
         assert str(ast) == str(expected_ast)
 
     def test_show_index(self):
-        sql = "SHOW INDEX FROM predictors"
+        sql = "SHOW INDEX FROM `predictors`"
         ast = parse_sql(sql)
         expected_ast = Show(
             category='INDEX',
-            from_table=Identifier('predictors')
+            from_table=Identifier('`predictors`')
         )
 
         assert str(ast).lower() == sql.lower()
         assert str(ast) == str(expected_ast)
         assert ast.to_tree() == expected_ast.to_tree()
 
     def test_show_index_from_db(self):
-        sql = "SHOW INDEX FROM predictors FROM db"
+        sql = "SHOW INDEX FROM `predictors` FROM db"
         ast = parse_sql(sql)
         expected_ast = Show(
             category='INDEX',
-            from_table=Identifier('db.predictors'),
+            from_table=Identifier('db.`predictors`'),
         )
 
         # assert str(ast).lower() == sql.lower()

tests/test_standard_render.py

@@ -82,7 +82,4 @@ def test_standard_render():
 
         # inject function
         module.parse_sql = parse_sql2
-
         check_module(module)
-
-