Coverage for mindsdb / api / executor / planner / ts_utils.py: 90%

1from mindsdb_sql_parser.ast import Identifier, Operation, BinaryOperation, BetweenOperation, OrderBy 

2 

3from mindsdb.api.executor.planner.exceptions import PlanningException 

4 

5 

6def find_time_filter(op, time_column_name): 

7 if not op: 

8 return 

9 if op.op == 'and': 

10 left = find_time_filter(op.args[0], time_column_name) 

11 right = find_time_filter(op.args[1], time_column_name) 

12 if left and right: 

13 raise PlanningException('Can provide only one filter by predictor order_by column, found two') 

14 

15 return left or right 

16 elif ((isinstance(op.args[0], Identifier) and op.args[0].parts[-1].lower() == time_column_name.lower()) 

17 or (isinstance(op.args[1], Identifier) and op.args[1].parts[-1].lower() == time_column_name.lower())): 

18 return op 

19 

20 

21def replace_time_filter(op, time_filter, new_filter): 

22 if op == time_filter: 

23 return new_filter 

24 if isinstance(op, BinaryOperation): 

25 op.args[0] = replace_time_filter(op.args[0], time_filter, new_filter) 

26 op.args[1] = replace_time_filter(op.args[1], time_filter, new_filter) 

27 return op 

28 

29 

30def find_and_remove_time_filter(op, time_filter): 

31 if isinstance(op, BinaryOperation) or isinstance(op, BetweenOperation): 

32 if op == time_filter: 

33 return None 

34 elif op.op == 'and': 

35 # TODO maybe OR operation too? 

36 

37 # next level 

38 left_arg = find_and_remove_time_filter(op.args[0], time_filter) 

39 right_arg = find_and_remove_time_filter(op.args[1], time_filter) 

40 

41 # if found in one arg return other 

42 if left_arg is None: 

43 return right_arg 

44 if right_arg is None: 

45 return left_arg 

46 

47 op.args = [left_arg, right_arg] 

48 return op 

49 

50 return op 

51 

52 

53def validate_ts_where_condition(op, allowed_columns, allow_and=True): 

54 """Error if the where condition caontains invalid ops, is nested or filters on some column that's not time or partition""" 

55 if not op: 

56 return 

57 allowed_ops = ['and', '>', '>=', '=', '<', '<=', 'between', 'in'] 

58 if not allow_and: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true

59 allowed_ops.remove('and') 

60 if op.op not in allowed_ops: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true

61 raise PlanningException( 

62 f'For time series predictors only the following operations are allowed in WHERE: {str(allowed_ops)}, found instead: {str(op)}.') 

63 

64 for arg in op.args: 

65 if isinstance(arg, Identifier): 

66 if arg.parts[-1].lower() not in allowed_columns: 

67 raise PlanningException( 

68 f'For time series predictor only the following columns are allowed in WHERE: {str(allowed_columns)}, found instead: {str(arg)}.') 

69 # remove alias 

70 arg.parts = [arg.parts[-1]] 

71 

72 if isinstance(op.args[0], Operation): 

73 validate_ts_where_condition(op.args[0], allowed_columns, allow_and=True) 

74 if isinstance(op.args[1], Operation): 

75 validate_ts_where_condition(op.args[1], allowed_columns, allow_and=True) 

76 

77 

78def recursively_check_join_identifiers_for_ambiguity(item, aliased_fields=None): 

79 if item is None: 

80 return 

81 elif isinstance(item, Identifier): 

82 if len(item.parts) == 1: 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true

83 if aliased_fields is not None and item.parts[0] in aliased_fields: 

84 # is alias 

85 return 

86 raise PlanningException(f'Ambigous identifier {str(item)}, provide table name for operations on a join.') 

87 elif isinstance(item, Operation): 

88 recursively_check_join_identifiers_for_ambiguity(item.args, aliased_fields=aliased_fields) 

89 elif isinstance(item, OrderBy): 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true

90 recursively_check_join_identifiers_for_ambiguity(item.field, aliased_fields=aliased_fields) 

91 elif isinstance(item, list): 

92 for arg in item: 

93 recursively_check_join_identifiers_for_ambiguity(arg, aliased_fields=aliased_fields)