Coverage for /Users/davegaeddert/Development/dropseed/plain/plain-worker/plain/worker/scheduling.py: 80%

114 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-16 22:03 -0500

1import datetime 

2import subprocess 

3 

4from plain.utils import timezone 

5 

6from .jobs import Job, load_job 

7 

8_MONTH_NAMES = { 

9 "JAN": 1, 

10 "FEB": 2, 

11 "MAR": 3, 

12 "APR": 4, 

13 "MAY": 5, 

14 "JUN": 6, 

15 "JUL": 7, 

16 "AUG": 8, 

17 "SEP": 9, 

18 "OCT": 10, 

19 "NOV": 11, 

20 "DEC": 12, 

21} 

22_DAY_NAMES = { 

23 "MON": 0, 

24 "TUE": 1, 

25 "WED": 2, 

26 "THU": 3, 

27 "FRI": 4, 

28 "SAT": 5, 

29 "SUN": 6, 

30} 

31 

32 

33class _ScheduleComponent: 

34 def __init__(self, values, raw=""): 

35 self.values = sorted(values) 

36 self._raw = raw 

37 

38 def __str__(self): 

39 if self._raw: 

40 return self._raw 

41 return ",".join(str(v) for v in self.values) 

42 

43 def __eq__(self, other): 

44 return self.values == other.values 

45 

46 @classmethod 

47 def parse(cls, value, min_allowed, max_allowed, str_conversions=None): 

48 if str_conversions is None: 

49 str_conversions = {} 

50 

51 if isinstance(value, int): 

52 if value < min_allowed or value > max_allowed: 

53 raise ValueError( 

54 f"Schedule component should be between {min_allowed} and {max_allowed}" 

55 ) 

56 return cls([value], raw=value) 

57 

58 if not isinstance(value, str): 

59 raise ValueError("Schedule component should be an int or str") 

60 

61 # First split any subcomponents and re-parse them 

62 if "," in value: 

63 return cls( 

64 sum( 

65 ( 

66 cls.parse( 

67 sub_value, min_allowed, max_allowed, str_conversions 

68 ).values 

69 for sub_value in value.split(",") 

70 ), 

71 [], 

72 ), 

73 raw=value, 

74 ) 

75 

76 if value == "*": 

77 return cls(list(range(min_allowed, max_allowed + 1)), raw=value) 

78 

79 def _convert(value): 

80 result = str_conversions.get(value.upper(), value) 

81 return int(result) 

82 

83 if "/" in value: 

84 values, step = value.split("/") 

85 values = cls.parse(values, min_allowed, max_allowed, str_conversions) 

86 return cls([v for v in values.values if v % int(step) == 0], raw=value) 

87 

88 if "-" in value: 

89 start, end = value.split("-") 

90 return cls(list(range(_convert(start), _convert(end) + 1)), raw=value) 

91 

92 return cls([_convert(value)], raw=value) 

93 

94 

95class Schedule: 

96 def __init__( 

97 self, 

98 *, 

99 minute="*", 

100 hour="*", 

101 day_of_month="*", 

102 month="*", 

103 day_of_week="*", 

104 raw="", 

105 ): 

106 self.minute = _ScheduleComponent.parse(minute, min_allowed=0, max_allowed=59) 

107 self.hour = _ScheduleComponent.parse(hour, min_allowed=0, max_allowed=23) 

108 self.day_of_month = _ScheduleComponent.parse( 

109 day_of_month, min_allowed=1, max_allowed=31 

110 ) 

111 self.month = _ScheduleComponent.parse( 

112 month, 

113 min_allowed=1, 

114 max_allowed=12, 

115 str_conversions=_MONTH_NAMES, 

116 ) 

117 self.day_of_week = _ScheduleComponent.parse( 

118 day_of_week, 

119 min_allowed=0, 

120 max_allowed=6, 

121 str_conversions=_DAY_NAMES, 

122 ) 

123 self._raw = raw 

124 

125 def __str__(self): 

126 if self._raw: 

127 return self._raw 

128 return f"{self.minute} {self.hour} {self.day_of_month} {self.month} {self.day_of_week}" 

129 

130 def __repr__(self) -> str: 

131 return f"<Schedule {self}>" 

132 

133 @classmethod 

134 def from_cron(cls, cron): 

135 raw = cron 

136 

137 if cron == "@yearly" or cron == "@annually": 

138 cron = "0 0 1 1 *" 

139 elif cron == "@monthly": 

140 cron = "0 0 1 * *" 

141 elif cron == "@weekly": 

142 cron = "0 0 * * 0" 

143 elif cron == "@daily" or cron == "@midnight": 

144 cron = "0 0 * * *" 

145 elif cron == "@hourly": 

146 cron = "0 * * * *" 

147 

148 minute, hour, day_of_month, month, day_of_week = cron.split() 

149 

150 return cls( 

151 minute=minute, 

152 hour=hour, 

153 day_of_month=day_of_month, 

154 month=month, 

155 day_of_week=day_of_week, 

156 raw=raw, 

157 ) 

158 

159 def next(self, now=None): 

160 """ 

161 Find the next datetime that matches the schedule after the given datetime. 

162 """ 

163 dt = now or timezone.localtime() # Use the defined plain timezone by default 

164 

165 # We only care about minutes, so immediately jump to the next minute 

166 dt += datetime.timedelta(minutes=1) 

167 dt = dt.replace(second=0, microsecond=0) 

168 

169 def _go_to_next_day(v): 

170 v = v + datetime.timedelta(days=1) 

171 return v.replace( 

172 hour=self.hour.values[0], 

173 minute=self.minute.values[0], 

174 ) 

175 

176 # If we don't find a value in the next 500 days, 

177 # then the schedule is probably never going to match (i.e. Feb 31) 

178 max_future = dt + datetime.timedelta(days=500) 

179 

180 while True: 

181 is_valid_day = ( 

182 dt.month in self.month.values 

183 and dt.day in self.day_of_month.values 

184 and dt.weekday() in self.day_of_week.values 

185 ) 

186 if is_valid_day: 

187 # We're on a valid day, now find the next valid hour and minute 

188 for hour in self.hour.values: 

189 if hour < dt.hour: 

190 continue 

191 for minute in self.minute.values: 

192 if hour == dt.hour and minute < dt.minute: 

193 continue 

194 candidate_datetime = dt.replace(hour=hour, minute=minute) 

195 if candidate_datetime >= dt: 

196 return candidate_datetime 

197 # If no valid time is found today, reset to the first valid minute and hour of the next day 

198 dt = _go_to_next_day(dt) 

199 else: 

200 # Increment the day until a valid month/day/weekday combination is found 

201 dt = _go_to_next_day(dt) 

202 

203 if dt > max_future: 

204 raise ValueError("No valid schedule match found in the next 500 days") 

205 

206 

207class ScheduledCommand(Job): 

208 def __init__(self, command): 

209 self.command = command 

210 

211 def __repr__(self) -> str: 

212 return f"<ScheduledCommand: {self.command}>" 

213 

214 def run(self): 

215 subprocess.run(self.command, shell=True, check=True) 

216 

217 def get_unique_key(self) -> str: 

218 # The ScheduledCommand can be used for different commands, 

219 # so we need the unique_key to separate them in the scheduling uniqueness logic 

220 return self.command 

221 

222 

223def load_schedule(schedules): 

224 jobs_schedule = [] 

225 

226 for job, schedule in schedules: 

227 if isinstance(job, str): 

228 if job.startswith("cmd:"): 

229 job = ScheduledCommand(job[4:]) 

230 else: 

231 job = load_job(job, {"args": [], "kwargs": {}}) 

232 

233 if isinstance(schedule, str): 

234 schedule = Schedule.from_cron(schedule) 

235 

236 jobs_schedule.append((job, schedule)) 

237 

238 return jobs_schedule