Coverage for packages/sql-smith/src/sql_smith/query/union_query.py: 40%

1from sql_smith.capability import CanUnionMixin, HasOrderByMixin 

2from sql_smith.functions import express 

3from sql_smith.query import AbstractQuery 

4 

5 

6class UnionQuery(CanUnionMixin, HasOrderByMixin, AbstractQuery): 

7 """Implements a union query.""" 

8 

9 def __init__( 

10 self, 

11 engine: "EngineInterface", 

12 left: "StatementInterface", 

13 right: "StatementInterface", 

14 ): 

15 super().__init__(engine) 

16 self._order_by = [] 

17 self._all = False 

18 self._left = left 

19 self._right = right 

20 

21 def all(self, state: bool = True) -> "UnionQuery": 

22 """Sets ALL.""" 

23 self._all = state 

24 return self 

25 

26 def as_expression(self) -> "ExpressionInterface": 

27 query = self.start_expression() 

28 query = self.__apply_all(query) 

29 query = self.__apply_right(query) 

30 query = self._apply_order_by(query) 

31 return query 

32 

33 def start_expression(self) -> "ExpressionInterface": 

34 return express("{} UNION", self._left) 

35 

36 def __apply_all(self, query: "ExpressionInterface") -> "ExpressionInterface": 

37 return query.append("ALL") if self._all else query 

38 

39 def __apply_right(self, query: "ExpressionInterface") -> "ExpressionInterface": 

40 return query.append("{}", self._right)