fix(parser): extract the column in between to filter_column_list (#44)

Author: Ifffff

parser/mysql_parser/lexer.py

@@ -16,7 +16,7 @@
 
 reserved = sorted(set(presto_tokens).difference(presto_nonreserved))
 
-tokens = ['INTEGER', 'DECIMAL', 'NUMBER',
+tokens = ['INTEGER', 'NUMBER', 'DOUBLE',
           'IDENTIFIER', 'DIGIT_IDENTIFIER',
           'QUOTED_IDENTIFIER', 'BACKQUOTED_IDENTIFIER',
           'STRING', 'PERIOD',
@@ -56,6 +56,15 @@
 t_ignore = ' \t'
 
 
+def t_DOUBLE(t):
+    r"""(\d+(?:\.\d*)?(?:[eE][+-]?\d+)?|\d*(?:\.\d+)(?:[eE][+-]?\d+)?)"""
+    if 'e' in t.value or 'E' in t.value or '.' in t.value:
+        t.type = 'DOUBLE'
+    else:
+        t.type = "INTEGER"
+    return t
+
+
 def t_INTEGER(t):
     r'\d+'
     t.type = "INTEGER"

parser/mysql_parser/parser.py

@@ -1089,9 +1089,9 @@ def p_quoted_identifier(p):
 
 
 def p_number(p):
-    r"""number : DECIMAL
+    r"""number : DOUBLE
                | INTEGER"""
-    if p.slice[1].type == "DECIMAL":
+    if p.slice[1].type == "DOUBLE":
         p[0] = DoubleLiteral(p.lineno(1), p.lexpos(1), p[1])
     else:
         p[0] = LongLiteral(p.lineno(1), p.lexpos(1), p[1])

parser/parser_utils.py

@@ -108,24 +108,8 @@ def visit_like_predicate(self, node, context):
                     if isinstance(pattern, StringLiteral):
                         if not pattern.value.startswith('%'):
                             can_query_range = True
-                    qualifed_name = node.value.name
                     if can_query_range:
-                        if len(qualifed_name.parts) == 2:
-                            table_or_alias_name = qualifed_name.parts[0]
-                            for _table in self.table_list:
-                                if _table['alias'] == table_or_alias_name or _table[
-                                    'table_name'] == table_or_alias_name:
-                                    filter_column_list = _table['filter_column_list']
-                                    filter_column_list.append({
-                                        'column_name': qualifed_name.parts[1],
-                                        'opt': 'like'
-                                    })
-                        else:
-                            filter_column_list = self.table_list[-1]['filter_column_list']
-                            filter_column_list.append({
-                                'column_name': qualifed_name.parts[0],
-                                'opt': 'like'
-                            })
+                        self.add_filter_column_with_qualified_name_reference(node.value, 'like')
 
                 return self.visit_expression(node, context)
 
@@ -141,22 +125,7 @@ def visit_in_predicate(self, node, context):
                         self.in_count_list.append(len(node.value_list.values))
 
                     if isinstance(value, QualifiedNameReference):
-                        if len(value.name.parts) == 2:
-                            table_or_alias_name = value.name.parts[0]
-                            for _table in self.table_list:
-                                if _table['alias'] == table_or_alias_name or _table[
-                                    'table_name'] == table_or_alias_name:
-                                    filter_column_list = _table['filter_column_list']
-                                    filter_column_list.append({
-                                        'column_name': value.name.parts[1],
-                                        'opt': 'in'
-                                    })
-                        else:
-                            filter_column_list = self.table_list[-1]['filter_column_list']
-                            filter_column_list.append({
-                                'column_name': value.name.parts[0],
-                                'opt': 'in'
-                            })
+                        self.add_filter_column_with_qualified_name_reference(value, 'in')
 
                 self.process(node.value, None)
                 self.process(node.value_list, None)
@@ -258,6 +227,30 @@ def visit_delete(self, node, context):
                     self.process(node.where, context)
                 return None
 
+            def visit_between_predicate(self, node, context):
+                if isinstance(node.value, QualifiedNameReference):
+                    self.add_filter_column_with_qualified_name_reference(node.value, 'between')
+                return None
+
+            def add_filter_column_with_qualified_name_reference(self, qualified_name_reference: QualifiedNameReference,
+                                                                opt):
+                if len(qualified_name_reference.name.parts) == 2:
+                    table_or_alias_name = qualified_name_reference.name.parts[0]
+                    for _table in self.table_list:
+                        if _table['alias'] == table_or_alias_name or _table[
+                            'table_name'] == table_or_alias_name:
+                            filter_column_list = _table['filter_column_list']
+                            filter_column_list.append({
+                                'column_name': qualified_name_reference.name.parts[1],
+                                'opt': opt
+                            })
+                else:
+                    filter_column_list = self.table_list[-1]['filter_column_list']
+                    filter_column_list.append({
+                        'column_name': qualified_name_reference.name.parts[0],
+                        'opt': opt
+                    })
+
         visitor = FormatVisitor()
         visitor.process(statement, None)
         return visitor

test/parser/test_parser_utils.py

@@ -219,6 +219,16 @@ def test_update_parameterize(self):
         statement = format_sql(statement_node, 0)
         assert statement == """UPDATE t1 SET c = ? WHERE id = ?"""
 
+    def test_between(self):
+        sql = """
+         select max(`successRate`) AS `successRate` from `table_850d` where `period` between '2022-07-11 00:00:00' and '2022-07-11 23:59:59' and `successRate` < 0.35;
+        """
+        visitor = ParserUtils.format_statement(parser.parse(sql))
+        table_list = visitor.table_list
+        assert table_list == [{'table_name': 'table_850d', 'alias': '',
+                               'filter_column_list': [{'column_name': 'period', 'opt': 'between'},
+                                                      {'column_name': 'successRate', 'opt': '<'}]}]
+
 
 if __name__ == '__main__':
     unittest.main()