Coverage for packages/sql-smith/src/sql_smith/query/update_query.py: 100%

1from typing import Any, Union 

2 

3from sql_smith.capability import HasWhereMixin 

4from sql_smith.functions import express, identify, listing, param 

5from sql_smith.query import AbstractQuery 

6 

7 

8class UpdateQuery(HasWhereMixin, AbstractQuery): 

9 """Implements an UPDATE query.""" 

10 

11 def __init__(self, engine: "EngineInterface"): 

12 super().__init__(engine) 

13 self._table = None 

14 self._set = None 

15 self._where = [] 

16 

17 def table(self, table: Union[str, "StatementInterface"]) -> "UpdateQuery": 

18 """Sets the table.""" 

19 self._table = identify(table) 

20 return self 

21 

22 def set(self, value_dict: dict[str, Any]): 

23 """Sets the column and values with a dictionary.""" 

24 self._set = listing( 

25 list( 

26 map( 

27 lambda k, v: express("{} = {}", identify(k), param(v)), 

28 value_dict.keys(), 

29 value_dict.values(), 

30 ) 

31 ) 

32 ) 

33 return self 

34 

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

36 query = self.start_expression() 

37 query = self.__apply_table(query) 

38 query = self.__apply_set(query) 

39 query = self._apply_where(query) 

40 

41 return query 

42 

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

44 return express("UPDATE") 

45 

46 def __apply_table(self, query: "ExpressionInterface") -> "ExpressionInterface": 

47 return query.append("{}", self._table) if self._table else query 

48 

49 def __apply_set(self, query: "ExpressionInterface") -> "ExpressionInterface": 

50 return query.append("SET {}", self._set) if self._set else query