Coverage for /Users/davegaeddert/Development/dropseed/plain/plain-worker/plain/worker/jobs.py: 27%

119 statements  

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

1import datetime 

2import inspect 

3import logging 

4from importlib import import_module 

5 

6from plain.models import IntegrityError, Model 

7from plain.utils import timezone 

8 

9logger = logging.getLogger(__name__) 

10 

11 

12def load_job(job_class_path, parameters): 

13 module_path, class_name = job_class_path.rsplit(".", 1) 

14 module = import_module(module_path) 

15 job_class = getattr(module, class_name) 

16 args, kwargs = JobParameters.from_json(parameters) 

17 return job_class(*args, **kwargs) 

18 

19 

20class JobParameters: 

21 @staticmethod 

22 def to_json(args, kwargs): 

23 serialized_args = [] 

24 for arg in args: 

25 if isinstance(arg, Model): 

26 serialized_args.append(ModelInstanceParameter.from_instance(arg)) 

27 else: 

28 serialized_args.append(arg) 

29 

30 serialized_kwargs = {} 

31 for key, value in kwargs.items(): 

32 if isinstance(value, Model): 

33 serialized_kwargs[key] = ModelInstanceParameter.from_instance(value) 

34 else: 

35 serialized_kwargs[key] = value 

36 

37 return {"args": serialized_args, "kwargs": serialized_kwargs} 

38 

39 @staticmethod 

40 def from_json(data): 

41 args = [] 

42 for arg in data["args"]: 

43 if ModelInstanceParameter.is_gid(arg): 

44 args.append(ModelInstanceParameter.to_instance(arg)) 

45 else: 

46 args.append(arg) 

47 

48 kwargs = {} 

49 for key, value in data["kwargs"].items(): 

50 if ModelInstanceParameter.is_gid(value): 

51 kwargs[key] = ModelInstanceParameter.to_instance(value) 

52 else: 

53 kwargs[key] = value 

54 

55 return args, kwargs 

56 

57 

58class ModelInstanceParameter: 

59 """ 

60 A string representation of a model instance, 

61 so we can convert a single parameter (model instance itself) 

62 into a string that can be serialized and stored in the database. 

63 """ 

64 

65 @staticmethod 

66 def from_instance(instance): 

67 return f"gid://{instance._meta.package_label}/{instance._meta.model_name}/{instance.pk}" 

68 

69 @staticmethod 

70 def to_instance(s): 

71 if not s.startswith("gid://"): 

72 raise ValueError("Invalid ModelInstanceParameter string") 

73 package, model, pk = s[6:].split("/") 

74 from plain.packages import packages 

75 

76 model = packages.get_model(package, model) 

77 return model.objects.get(pk=pk) 

78 

79 @staticmethod 

80 def is_gid(x): 

81 if not isinstance(x, str): 

82 return False 

83 return x.startswith("gid://") 

84 

85 

86class JobType(type): 

87 """ 

88 Metaclass allows us to capture the original args/kwargs 

89 used to instantiate the job, so we can store them in the database 

90 when we schedule the job. 

91 """ 

92 

93 def __call__(self, *args, **kwargs): 

94 instance = super().__call__(*args, **kwargs) 

95 instance._init_args = args 

96 instance._init_kwargs = kwargs 

97 return instance 

98 

99 

100class Job(metaclass=JobType): 

101 def run(self): 

102 raise NotImplementedError 

103 

104 def run_in_worker( 

105 self, 

106 *, 

107 queue: str | None = None, 

108 delay: int | datetime.timedelta | datetime.datetime | None = None, 

109 priority: int | None = None, 

110 retries: int | None = None, 

111 retry_attempt: int = 0, 

112 unique_key: str | None = None, 

113 ): 

114 from .models import JobRequest 

115 

116 try: 

117 # Try to automatically annotate the source of the job 

118 caller = inspect.stack()[1] 

119 source = f"{caller.filename}:{caller.lineno}" 

120 except (IndexError, AttributeError): 

121 source = "" 

122 

123 parameters = JobParameters.to_json(self._init_args, self._init_kwargs) 

124 

125 if queue is None: 

126 queue = self.get_queue() 

127 

128 if priority is None: 

129 priority = self.get_priority() 

130 

131 if retries is None: 

132 retries = self.get_retries() 

133 

134 if delay is None: 

135 start_at = None 

136 elif isinstance(delay, int): 

137 start_at = timezone.now() + datetime.timedelta(seconds=delay) 

138 elif isinstance(delay, datetime.timedelta): 

139 start_at = timezone.now() + delay 

140 elif isinstance(delay, datetime.datetime): 

141 start_at = delay 

142 else: 

143 raise ValueError(f"Invalid delay: {delay}") 

144 

145 if unique_key is None: 

146 unique_key = self.get_unique_key() 

147 

148 if unique_key: 

149 # Only need to look at in progress jobs 

150 # if we also have a unique key. 

151 # Otherwise it's up to the user to use _in_progress() 

152 if running := self._in_progress(unique_key): 

153 return running 

154 

155 try: 

156 job_request = JobRequest( 

157 job_class=self._job_class_str(), 

158 parameters=parameters, 

159 start_at=start_at, 

160 source=source, 

161 queue=queue, 

162 priority=priority, 

163 retries=retries, 

164 retry_attempt=retry_attempt, 

165 unique_key=unique_key, 

166 ) 

167 job_request.save( 

168 clean_and_validate=False 

169 ) # So IntegrityError is raised on unique instead of potentially confusing ValidationError... 

170 return job_request 

171 except IntegrityError as e: 

172 logger.warning("Job already in progress: %s", e) 

173 # Try to return the _in_progress list again 

174 return self._in_progress(unique_key) 

175 

176 def _job_class_str(self): 

177 return f"{self.__module__}.{self.__class__.__name__}" 

178 

179 def _in_progress(self, unique_key): 

180 """Get all JobRequests and Jobs that are currently in progress, regardless of queue.""" 

181 from .models import Job, JobRequest 

182 

183 job_class = self._job_class_str() 

184 

185 job_requests = JobRequest.objects.filter( 

186 job_class=job_class, 

187 unique_key=unique_key, 

188 ) 

189 

190 jobs = Job.objects.filter( 

191 job_class=job_class, 

192 unique_key=unique_key, 

193 ) 

194 

195 return list(job_requests) + list(jobs) 

196 

197 def get_unique_key(self) -> str: 

198 """ 

199 A unique key to prevent duplicate jobs from being queued. 

200 Enabled by returning a non-empty string. 

201 

202 Note that this is not a "once and only once" guarantee, but rather 

203 an "at least once" guarantee. Jobs should still be idempotent in case 

204 multiple instances are queued in a race condition. 

205 """ 

206 return "" 

207 

208 def get_queue(self) -> str: 

209 return "default" 

210 

211 def get_priority(self) -> int: 

212 return 0 

213 

214 def get_retries(self) -> int: 

215 return 0 

216 

217 def get_retry_delay(self, attempt: int) -> int: 

218 """ 

219 Calcluate a delay in seconds before the next retry attempt. 

220 

221 On the first retry, attempt will be 1. 

222 """ 

223 return 0