Coverage for open_webui/utils/recurrence.py: 24%

1"""Recurrence calculations isolated from application/DB imports for worker processes.""" 

2 

3import logging 

4from datetime import datetime, timedelta 

5from typing import Optional 

6from zoneinfo import ZoneInfo 

7 

8from anyio import fail_after, to_process 

9from dateutil.rrule import HOURLY, MINUTELY, SECONDLY, rruleset, rrulestr 

10from open_webui.constants import ERROR_MESSAGES 

11 

12log = logging.getLogger(__name__) 

13RRULE_TIMEOUT_SECONDS = 2 

14 

15 

16class RecurrenceEvaluationTimeout(ValueError): 

17 """The evaluation budget expired; the schedule may still have occurrences.""" 

18 

19 

20def _resolve_tz(tz: str = None) -> Optional[ZoneInfo]: 

21 """Safely resolve a timezone string to ZoneInfo. 

22 

23 Returns None (→ server-local fallback) when *tz* is empty, None, 

24 or an unrecognised IANA key. Logs a warning on bad keys so 

25 misconfiguration is visible in the server logs. 

26 """ 

27 if not tz: 27 ↛ 29line 27 didn't jump to line 29 because the condition on line 27 was always true

28 return None 

29 try: 

30 return ZoneInfo(tz) 

31 except (KeyError, Exception): 

32 log.warning('Unknown timezone %r — falling back to server time', tz) 

33 return None 

34 

35 

36def _parse_rule(s: str, now: Optional[datetime] = None): 

37 """Parse RRULE with clock-aligned DTSTART for sub-daily frequencies. 

38 

39 SECONDLY/MINUTELY/HOURLY rules use a fixed epoch DTSTART (2000-01-01 00:00) 

40 so intervals snap to clock boundaries (e.g. every 5min = :00, :05, :10). 

41 """ 

42 upper = s.upper() 

43 if 'EXRULE' in upper: 

44 raise ValueError('EXRULE is not supported in recurrence rules') 

45 

46 parsed = rrulestr(s, ignoretz=True) 

47 rules = parsed._rrule if isinstance(parsed, rruleset) else [parsed] 

48 if len(rules) > 1: 

49 raise ValueError('only one RRULE is supported per recurrence rule') 

50 

51 rule = rules[0] 

52 start = rule._dtstart.replace(tzinfo=None) 

53 anchor = now or datetime.now() 

54 parts = s.split() 

55 stripped = '\n'.join(part for part in parts if not part.upper().startswith('DTSTART')) or s 

56 has_dtstart = any(part.upper().startswith('DTSTART') for part in parts) 

57 step = { 

58 SECONDLY: timedelta(seconds=rule._interval), 

59 MINUTELY: timedelta(minutes=rule._interval), 

60 HOURLY: timedelta(hours=rule._interval), 

61 }.get(rule._freq) 

62 

63 if step is None: 

64 if not rule._dtstart.tzinfo: 

65 return parsed 

66 return rrulestr(stripped, dtstart=start, ignoretz=True) 

67 

68 if rule._interval < 1: 

69 raise ValueError('RRULE INTERVAL must be a positive integer') 

70 dtstart = None 

71 if has_dtstart: 

72 emitted = ((anchor - start) // step) if anchor > start else 0 

73 emitted *= len(rule._byminute or (0,)) * len(rule._bysecond or (0,)) 

74 if emitted <= 100_000: 

75 if rule._dtstart.tzinfo: 

76 dtstart = start 

77 else: 

78 return parsed 

79 if not has_dtstart or dtstart is None: 

80 epoch = datetime(2000, 1, 1) 

81 dtstart = epoch + ((anchor - epoch) // step) * step 

82 

83 return rrulestr(stripped, dtstart=dtstart, ignoretz=True) 

84 

85 

86def _next_occurrences(s: str, now: datetime, n: int) -> list[datetime]: 

87 rule = _parse_rule(s, now) 

88 occurrences = [] 

89 for _ in range(n): 

90 now = rule.after(now) 

91 if now is None: 

92 break 

93 occurrences.append(now) 

94 return occurrences 

95 

96 

97async def _get_next_occurrences(s: str, now: datetime, n: int) -> list[datetime]: 

98 # A result-count or date limit cannot bound work before the first match. 

99 try: 

100 with fail_after(RRULE_TIMEOUT_SECONDS): 

101 return await to_process.run_sync(_next_occurrences, s, now, n, cancellable=True) 

102 except TimeoutError as e: 

103 raise RecurrenceEvaluationTimeout('Schedule took too long to evaluate; simplify its recurrence rule.') from e 

104 

105 

106async def validate_rrule(s: str, tz: str = None) -> None: 

107 """Raise ValueError if the RRULE is malformed or exhausted. 

108 

109 When *tz* is provided the "now" reference uses the user's local 

110 clock so that near-future schedules are not incorrectly rejected 

111 on servers whose system clock is ahead (e.g. UTC vs US timezones). 

112 """ 

113 upper = s.upper() 

114 if 'COUNT=' in upper and 'DTSTART' not in upper: 114 ↛ 115line 114 didn't jump to line 115 because the condition on line 114 was never true

115 raise ValueError(ERROR_MESSAGES.AUTOMATION_COUNT_REQUIRES_DTSTART) 

116 zi = _resolve_tz(tz) 

117 now = datetime.now(zi).replace(tzinfo=None) if zi else datetime.now() 

118 try: 

119 occurrences = await _get_next_occurrences(s, now, 1) 

120 except RecurrenceEvaluationTimeout: 

121 raise 

122 except Exception as e: 

123 raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e)) 

124 if not occurrences: 

125 raise ValueError(ERROR_MESSAGES.AUTOMATION_NO_FUTURE_RUNS) 

126 

127 

128async def next_run_ns(s: str, tz: str = None) -> Optional[int]: 

129 """Next occurrence as epoch nanoseconds, respecting user timezone.""" 

130 zi = _resolve_tz(tz) 

131 now = datetime.now(zi) if zi else datetime.now() 

132 now_naive = now.replace(tzinfo=None) 

133 occurrences = await _get_next_occurrences(s, now_naive, 1) 

134 if not occurrences: 

135 return None 

136 dt = occurrences[0] 

137 if zi: 

138 dt = dt.replace(tzinfo=zi) 

139 return int(dt.timestamp() * 1_000_000_000) 

140 

141 

142async def next_n_runs_ns(s: str, n: int = 5, tz: str = None) -> list[int]: 

143 """Compute next N occurrences for UI preview. 

144 

145 Uses the user's timezone for the starting "now" so that the 

146 preview matches the user's local clock (same as next_run_ns). 

147 """ 

148 zi = _resolve_tz(tz) 

149 result = [] 

150 now = datetime.now(zi).replace(tzinfo=None) if zi else datetime.now() 

151 for dt in await _get_next_occurrences(s, now, n): 

152 if zi: 

153 dt_tz = dt.replace(tzinfo=zi) 

154 result.append(int(dt_tz.timestamp() * 1_000_000_000)) 

155 else: 

156 result.append(int(dt.timestamp() * 1_000_000_000)) 

157 return result 

158 

159 

160async def rrule_interval_seconds(s: str) -> Optional[int]: 

161 """Approximate interval between recurrences in seconds. 

162 

163 Returns None for one-shot (COUNT=1) schedules or rules 

164 with fewer than two future occurrences. 

165 """ 

166 s = '\n'.join(part for part in s.split() if not part.upper().startswith('DTSTART')) or s 

167 now = datetime.now() 

168 occurrences = await _get_next_occurrences(s, now, 2) 

169 if len(occurrences) < 2: 

170 return None 

171 return int((occurrences[1] - occurrences[0]).total_seconds())