Coverage for mindsdb / api / executor / sql_query / steps / delete_step.py: 32%

1import copy 

2 

3from mindsdb_sql_parser.ast import ( 

4 Identifier, 

5 Constant, 

6 Delete, 

7 Parameter, 

8 Tuple, 

9) 

10from mindsdb.integrations.utilities.query_traversal import query_traversal 

11from mindsdb.api.executor.planner.steps import DeleteStep 

12 

13from mindsdb.api.executor.sql_query.result_set import ResultSet 

14 

15from .base import BaseStepCall 

16 

17 

18class DeleteStepCall(BaseStepCall): 

19 

20 bind = DeleteStep 

21 

22 def call(self, step): 

23 if len(step.table.parts) > 1: 

24 integration_name = step.table.parts[0] 

25 table_name_parts = step.table.parts[1:] 

26 else: 

27 integration_name = self.context['database'] 

28 table_name_parts = step.table.parts 

29 

30 dn = self.session.datahub.get(integration_name) 

31 

32 # make command 

33 query = Delete( 

34 table=Identifier(parts=table_name_parts), 

35 where=copy.deepcopy(step.where), 

36 ) 

37 

38 # fill params 

39 def fill_params(node, **kwargs): 

40 if isinstance(node, Parameter): 

41 rs = self.steps_data[node.value.step_num] 

42 items = [Constant(i) for i in rs.get_column_values(col_idx=0)] 

43 return Tuple(items) 

44 

45 query_traversal(query.where, fill_params) 

46 

47 response = dn.query(query=query, session=self.session) 

48 return ResultSet(affected_rows=response.affected_rows)