test

Author: StpMax

mindsdb_sql_parser/ast/select/union.py

@@ -9,11 +9,13 @@ def __init__(self,
                  left,
                  right,
                  unique=True,
+                 distinct_key=False,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.left = left
         self.right = right
         self.unique = unique
+        self.distinct_key = distinct_key
 
         if self.alias:
             self.parentheses = True
@@ -26,7 +28,7 @@ def to_tree(self, *args, level=0, **kwargs):
         right_str = f'\n{ind1}right=\n{self.right.to_tree(level=level + 2)},'
 
         cls_name = self.__class__.__name__
-        out_str = f'{ind}{cls_name}(unique={repr(self.unique)},' \
+        out_str = f'{ind}{cls_name}(unique={repr(self.unique)}, distinct_key={repr(self.distinct_key)}' \
                   f'{left_str}' \
                   f'{right_str}' \
                   f'\n{ind})'
@@ -38,6 +40,8 @@ def get_string(self, *args, **kwargs):
         keyword = self.operation
         if not self.unique:
             keyword += ' ALL'
+        if self.distinct_key:
+            keyword += ' DISTINCT'
         out_str = f"""{left_str}\n{keyword}\n{right_str}"""
 
         return out_str

tests/test_base_sql/test_union.py

@@ -12,21 +12,26 @@ def test_single_select_error(self):
 
     def test_union_base(self):
         for keyword, cls in {'union': Union, 'intersect': Intersect, 'except': Except}.items():
-            sql = f"""SELECT col1 FROM tab1 
-            {keyword} 
-            SELECT col1 FROM tab2"""
+            for rule in ['', 'distinct']:
+                sql = f"""SELECT col1 FROM tab1 
+                {keyword} {rule}
+                SELECT col1 FROM tab2"""
 
-            ast = parse_sql(sql)
-            expected_ast = cls(unique=True,
-                                 left=Select(targets=[Identifier('col1')],
-                                             from_table=Identifier(parts=['tab1']),
-                                             ),
-                                 right=Select(targets=[Identifier('col1')],
-                                             from_table=Identifier(parts=['tab2']),
-                                             ),
-                                 )
-            assert ast.to_tree() == expected_ast.to_tree()
-            assert str(ast) == str(expected_ast)
+                ast = parse_sql(sql)
+                expected_ast = cls(
+                    unique=True,
+                    distinct_key=rule == 'distinct',
+                    left=Select(
+                        targets=[Identifier('col1')],
+                        from_table=Identifier(parts=['tab1']),
+                    ),
+                    right=Select(
+                        targets=[Identifier('col1')],
+                        from_table=Identifier(parts=['tab2']),
+                    ),
+                )
+                assert ast.to_tree() == expected_ast.to_tree()
+                assert str(ast) == str(expected_ast)
 
     def test_union_all(self):
         for keyword, cls in {'union': Union, 'intersect': Intersect, 'except': Except}.items():