Coverage for packages/sql-smith/src/sql_smith/query_factory.py: 88%

1from typing import Any, Union 

2 

3from sql_smith.engine import BasicEngine 

4 

5 

6class QueryFactory: 

7 """Factory class for creating queries""" 

8 

9 def __init__(self, engine: "EngineInterface" = None): 

10 if engine is None: 

11 self._engine = BasicEngine() 

12 else: 

13 self._engine = engine 

14 

15 def select(self, *columns: Union[str, "StatementInterface"]) -> "SelectQuery": 

16 """Create a SELECT query""" 

17 query = self._engine.make_select() 

18 if len(columns) == 0: 

19 return query 

20 

21 return query.columns(*columns) 

22 

23 def select_distinct( 

24 self, *columns: Union[str, "StatementInterface"] 

25 ) -> "SelectQuery": 

26 """Create a SELECT DISTINCT query""" 

27 return self.select(*columns).distinct() 

28 

29 def insert(self, table: str, column_values: dict[str, Any] = None) -> "InsertQuery": 

30 """Create an INSERT query""" 

31 query = self._engine.make_insert().into(table) 

32 if column_values: 

33 query = query.map(column_values) 

34 

35 return query 

36 

37 def delete(self, table: str) -> "DeleteQuery": 

38 """Create a DELETE query""" 

39 return self._engine.make_delete().from_(table) 

40 

41 def update(self, table: str, value_set=None) -> "UpdateQuery": 

42 """Create an UPDATE query""" 

43 query = self._engine.make_update().table(table) 

44 if value_set: 

45 query.set(value_set) 

46 

47 return query