Coverage for mindsdb / integrations / utilities / handlers / query_utilities / update_query_utilities.py: 0%

1from mindsdb_sql_parser import ast 

2from typing import Text, List, Optional 

3 

4from .exceptions import UnsupportedColumnException 

5 

6from mindsdb.integrations.utilities.handlers.query_utilities.base_query_utilities import BaseQueryParser 

7from mindsdb.integrations.utilities.handlers.query_utilities.base_query_utilities import BaseQueryExecutor 

8 

9 

10class UPDATEQueryParser(BaseQueryParser): 

11 """ 

12 Parses an UPDATE query into its component parts. 

13 

14 Parameters 

15 ---------- 

16 query : ast.Update 

17 Given SQL UPDATE query. 

18 supported_columns : List[Text], Optional 

19 List of columns supported by the table for updating. 

20 """ 

21 def __init__(self, query: ast.Update, supported_columns: Optional[List[Text]] = None): 

22 super().__init__(query) 

23 self.supported_columns = supported_columns 

24 

25 def parse_query(self): 

26 """ 

27 Parses a SQL UPDATE statement into its components: the columns and values to update as a dictionary, and the WHERE conditions. 

28 """ 

29 values_to_update = self.parse_set_clause() 

30 where_conditions = self.parse_where_clause() 

31 

32 return values_to_update, where_conditions 

33 

34 def parse_set_clause(self): 

35 """ 

36 Parses the SET clause of the query and returns a dictionary of columns and values to update. 

37 """ 

38 values = list(self.query.update_columns.items()) 

39 

40 values_to_update = {} 

41 for value in values: 

42 if self.supported_columns: 

43 if value[0] not in self.supported_columns: 

44 raise UnsupportedColumnException(f"Unsupported column: {value[0]}") 

45 

46 values_to_update[value[0]] = value[1].value 

47 

48 return values_to_update 

49 

50 

51class UPDATEQueryExecutor(BaseQueryExecutor): 

52 """ 

53 Executes an UPDATE query. 

54 

55 Parameters 

56 ---------- 

57 df : pd.DataFrame 

58 Given table. 

59 where_conditions : List[List[Text]] 

60 WHERE conditions of the query. 

61 

62 NOTE: This class DOES NOT update the relevant records of the entity for you, it will simply return the records that need to be updated based on the WHERE conditions. 

63 

64 This class expects all of the entities to be passed in as a DataFrane and filters out the relevant records based on the WHERE conditions. 

65 Because all of the records need to be extracted to be passed in as a DataFrame, this class is not very computationally efficient. 

66 Therefore, DO NOT use this class if the API/SDK that you are using supports updating records in bulk. 

67 """