Coverage for packages/sql-smith/src/sql_smith/query/insert_query.py: 97%

1from typing import Any, Dict 

2 

3from sql_smith.functions import express, identify, identify_all, listing, param_all 

4from sql_smith.query import AbstractQuery 

5 

6 

7class InsertQuery(AbstractQuery): 

8 """Implements the INSERT query.""" 

9 

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

11 AbstractQuery.__init__(self, engine) 

12 self._into = None 

13 self._columns = None 

14 self._values = [] 

15 

16 def into(self, table: str) -> "InsertQuery": 

17 """Sets the table.""" 

18 self._into = identify(table) 

19 return self 

20 

21 def map(self, column_values: Dict[str, Any]): 

22 """Maps a dictionary to columns and values.""" 

23 return self.columns(*column_values.keys()).values(*column_values.values()) 

24 

25 def columns(self, *columns) -> "InsertQuery": 

26 """Sets the columns to insert.""" 

27 self._columns = listing(identify_all(*columns)) 

28 return self 

29 

30 def values(self, *values) -> "InsertQuery": 

31 """Appends values.""" 

32 self._values.append(express("({})", listing(param_all(*values)))) 

33 return self 

34 

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

36 query = self.start_expression() 

37 query = self.__apply_into(query) 

38 query = self.__apply_columns(query) 

39 query = self.__apply_values(query) 

40 

41 return query 

42 

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

44 return express("INSERT") 

45 

46 def __apply_into(self, query): 

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

48 

49 def __apply_columns(self, query): 

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

51 

52 def __apply_values(self, query): 

53 return ( 

54 query.append("VALUES {}", listing(self._values)) if self._values else query 

55 )