Coverage for /usr/local/lib/python3.12/site-packages/prefect/automations.py: 0%

1from typing import TYPE_CHECKING, Optional, overload 

2from uuid import UUID 

3 

4from pydantic import Field 

5from typing_extensions import Self 

6 

7from prefect._internal.compatibility.async_dispatch import async_dispatch 

8from prefect.client.orchestration import get_client 

9from prefect.events.actions import ( 

10 CallWebhook, 

11 CancelFlowRun, 

12 ChangeFlowRunState, 

13 DeclareIncident, 

14 DoNothing, 

15 PauseAutomation, 

16 PauseDeployment, 

17 PauseWorkPool, 

18 PauseWorkQueue, 

19 ResumeAutomation, 

20 ResumeDeployment, 

21 ResumeWorkPool, 

22 ResumeWorkQueue, 

23 RunDeployment, 

24 SendNotification, 

25 SuspendFlowRun, 

26) 

27from prefect.events.schemas.automations import ( 

28 AutomationCore, 

29 CompositeTrigger, 

30 CompoundTrigger, 

31 EventTrigger, 

32 MetricTrigger, 

33 MetricTriggerOperator, 

34 MetricTriggerQuery, 

35 Posture, 

36 PrefectMetric, 

37 ResourceSpecification, 

38 ResourceTrigger, 

39 SequenceTrigger, 

40 Trigger, 

41) 

42from prefect.exceptions import PrefectHTTPStatusError 

43 

44__all__ = [ 

45 "AutomationCore", 

46 "EventTrigger", 

47 "ResourceTrigger", 

48 "Posture", 

49 "Trigger", 

50 "ResourceSpecification", 

51 "MetricTriggerOperator", 

52 "MetricTrigger", 

53 "PrefectMetric", 

54 "CompositeTrigger", 

55 "SequenceTrigger", 

56 "CompoundTrigger", 

57 "MetricTriggerQuery", 

58 # action types 

59 "DoNothing", 

60 "RunDeployment", 

61 "PauseDeployment", 

62 "ResumeDeployment", 

63 "CancelFlowRun", 

64 "ChangeFlowRunState", 

65 "PauseWorkQueue", 

66 "ResumeWorkQueue", 

67 "SendNotification", 

68 "CallWebhook", 

69 "PauseAutomation", 

70 "ResumeAutomation", 

71 "SuspendFlowRun", 

72 "PauseWorkPool", 

73 "ResumeWorkPool", 

74 "DeclareIncident", 

75] 

76 

77 

78class Automation(AutomationCore): 

79 id: Optional[UUID] = Field(default=None, description="The ID of this automation") 

80 

81 async def acreate(self: Self) -> Self: 

82 """ 

83 Asynchronously create a new automation. 

84 

85 Examples: 

86 

87 ```python 

88 auto_to_create = Automation( 

89 name="woodchonk", 

90 trigger=EventTrigger( 

91 expect={"animal.walked"}, 

92 match={ 

93 "genus": "Marmota", 

94 "species": "monax", 

95 }, 

96 posture="Reactive", 

97 threshold=3, 

98 within=timedelta(seconds=10), 

99 ), 

100 actions=[CancelFlowRun()] 

101 ) 

102 created_automation = await auto_to_create.acreate() 

103 ``` 

104 """ 

105 async with get_client() as client: 

106 automation = AutomationCore(**self.model_dump(exclude={"id"})) 

107 self.id = await client.create_automation(automation=automation) 

108 return self 

109 

110 @async_dispatch(acreate) 

111 def create(self: Self) -> Self: 

112 """ 

113 Create a new automation. 

114 

115 Examples: 

116 

117 ```python 

118 auto_to_create = Automation( 

119 name="woodchonk", 

120 trigger=EventTrigger( 

121 expect={"animal.walked"}, 

122 match={ 

123 "genus": "Marmota", 

124 "species": "monax", 

125 }, 

126 posture="Reactive", 

127 threshold=3, 

128 within=timedelta(seconds=10), 

129 ), 

130 actions=[CancelFlowRun()] 

131 ) 

132 created_automation = auto_to_create.create() 

133 ``` 

134 """ 

135 with get_client(sync_client=True) as client: 

136 automation = AutomationCore(**self.model_dump(exclude={"id"})) 

137 self.id = client.create_automation(automation=automation) 

138 return self 

139 

140 async def aupdate(self: Self) -> None: 

141 """ 

142 Updates an existing automation. 

143 

144 Examples: 

145 

146 ```python 

147 auto = Automation.read(id=123) 

148 auto.name = "new name" 

149 auto.update() 

150 ``` 

151 """ 

152 assert self.id is not None 

153 async with get_client() as client: 

154 automation = AutomationCore( 

155 **self.model_dump(exclude={"id", "owner_resource"}) 

156 ) 

157 await client.update_automation(automation_id=self.id, automation=automation) 

158 

159 @async_dispatch(aupdate) 

160 def update(self: Self): 

161 """ 

162 Updates an existing automation. 

163 

164 Examples: 

165 

166 

167 ```python 

168 auto = Automation.read(id=123) 

169 auto.name = "new name" 

170 auto.update() 

171 ``` 

172 """ 

173 assert self.id is not None 

174 with get_client(sync_client=True) as client: 

175 automation = AutomationCore( 

176 **self.model_dump(exclude={"id", "owner_resource"}) 

177 ) 

178 client.update_automation(automation_id=self.id, automation=automation) 

179 

180 @overload 

181 @classmethod 

182 async def aread(cls, id: UUID, name: Optional[str] = ...) -> Self: ... 

183 

184 @overload 

185 @classmethod 

186 async def aread(cls, id: None = None, name: str = ...) -> Self: ... 

187 

188 @classmethod 

189 async def aread(cls, id: Optional[UUID] = None, name: Optional[str] = None) -> Self: 

190 """ 

191 Asynchronously read an automation by ID or name. 

192 

193 Examples: 

194 

195 ```python 

196 automation = await Automation.aread(name="woodchonk") 

197 ``` 

198 

199 ```python 

200 automation = await Automation.aread(id=UUID("b3514963-02b1-47a5-93d1-6eeb131041cb")) 

201 ``` 

202 """ 

203 if id and name: 

204 raise ValueError("Only one of id or name can be provided") 

205 if not id and not name: 

206 raise ValueError("One of id or name must be provided") 

207 async with get_client() as client: 

208 if id: 

209 try: 

210 automation = await client.read_automation(automation_id=id) 

211 except PrefectHTTPStatusError as exc: 

212 if exc.response.status_code == 404: 

213 raise ValueError(f"Automation with ID {id!r} not found") 

214 raise 

215 if automation is None: 

216 raise ValueError(f"Automation with ID {id!r} not found") 

217 return cls(**automation.model_dump()) 

218 else: 

219 if TYPE_CHECKING: 

220 assert name is not None 

221 automation = await client.read_automations_by_name(name=name) 

222 if len(automation) > 0: 

223 return cls(**automation[0].model_dump()) 

224 raise ValueError(f"Automation with name {name!r} not found") 

225 

226 @overload 

227 @classmethod 

228 async def read(cls, id: UUID, name: Optional[str] = ...) -> Self: ... 

229 

230 @overload 

231 @classmethod 

232 async def read(cls, id: None = None, name: str = ...) -> Self: ... 

233 

234 @classmethod 

235 @async_dispatch(aread) 

236 def read(cls, id: Optional[UUID] = None, name: Optional[str] = None) -> Self: 

237 """ 

238 Read an automation by ID or name. 

239 

240 Examples: 

241 

242 ```python 

243 automation = Automation.read(name="woodchonk") 

244 ``` 

245 

246 ```python 

247 automation = Automation.read(id=UUID("b3514963-02b1-47a5-93d1-6eeb131041cb")) 

248 ``` 

249 """ 

250 if id and name: 

251 raise ValueError("Only one of id or name can be provided") 

252 if not id and not name: 

253 raise ValueError("One of id or name must be provided") 

254 with get_client(sync_client=True) as client: 

255 if id: 

256 try: 

257 automation = client.read_automation(automation_id=id) 

258 except PrefectHTTPStatusError as exc: 

259 if exc.response.status_code == 404: 

260 raise ValueError(f"Automation with ID {id!r} not found") 

261 raise 

262 if automation is None: 

263 raise ValueError(f"Automation with ID {id!r} not found") 

264 return cls(**automation.model_dump()) 

265 else: 

266 if TYPE_CHECKING: 

267 assert name is not None 

268 automation = client.read_automations_by_name(name=name) 

269 if len(automation) > 0: 

270 return cls(**automation[0].model_dump()) 

271 raise ValueError(f"Automation with name {name!r} not found") 

272 

273 async def adelete(self: Self) -> bool: 

274 """ 

275 Asynchronously delete an automation. 

276 

277 Examples: 

278 

279 ```python 

280 auto = Automation.read(id = 123) 

281 await auto.adelete() 

282 ``` 

283 """ 

284 if self.id is None: 

285 raise ValueError("Can't delete an automation without an id") 

286 

287 async with get_client() as client: 

288 try: 

289 await client.delete_automation(self.id) 

290 return True 

291 except PrefectHTTPStatusError as exc: 

292 if exc.response.status_code == 404: 

293 return False 

294 raise 

295 

296 @async_dispatch(adelete) 

297 def delete(self: Self) -> bool: 

298 """ 

299 Delete an automation. 

300 

301 Examples: 

302 

303 ```python 

304 auto = Automation.read(id = 123) 

305 auto.delete() 

306 ``` 

307 """ 

308 if self.id is None: 

309 raise ValueError("Can't delete an automation without an id") 

310 

311 with get_client(sync_client=True) as client: 

312 try: 

313 client.delete_automation(self.id) 

314 return True 

315 except PrefectHTTPStatusError as exc: 

316 if exc.response.status_code == 404: 

317 return False 

318 raise 

319 

320 async def adisable(self: Self) -> bool: 

321 """ 

322 Asynchronously disable an automation. 

323 

324 Raises: 

325 ValueError: If the automation does not have an id 

326 PrefectHTTPStatusError: If the automation cannot be disabled 

327 

328 Example: 

329 ```python 

330 auto = await Automation.aread(id = 123) 

331 await auto.adisable() 

332 ``` 

333 """ 

334 if self.id is None: 

335 raise ValueError("Can't disable an automation without an id") 

336 

337 async with get_client() as client: 

338 try: 

339 await client.pause_automation(self.id) 

340 return True 

341 except PrefectHTTPStatusError as exc: 

342 if exc.response.status_code == 404: 

343 return False 

344 raise 

345 

346 @async_dispatch(adisable) 

347 def disable(self: Self) -> bool: 

348 """ 

349 Disable an automation. 

350 

351 

352 Raises: 

353 ValueError: If the automation does not have an id 

354 PrefectHTTPStatusError: If the automation cannot be disabled 

355 

356 Example: 

357 ```python 

358 auto = Automation.read(id = 123) 

359 auto.disable() 

360 ``` 

361 """ 

362 if self.id is None: 

363 raise ValueError("Can't disable an automation without an id") 

364 

365 with get_client(sync_client=True) as client: 

366 try: 

367 client.pause_automation(self.id) 

368 return True 

369 except PrefectHTTPStatusError as exc: 

370 if exc.response.status_code == 404: 

371 return False 

372 raise 

373 

374 async def aenable(self: Self) -> bool: 

375 """ 

376 Asynchronously enable an automation. 

377 

378 Raises: 

379 ValueError: If the automation does not have an id 

380 PrefectHTTPStatusError: If the automation cannot be enabled 

381 

382 Example: 

383 ```python 

384 auto = await Automation.aread(id = 123) 

385 await auto.aenable() 

386 ``` 

387 """ 

388 if self.id is None: 

389 raise ValueError("Can't enable an automation without an id") 

390 

391 async with get_client() as client: 

392 try: 

393 await client.resume_automation(self.id) 

394 return True 

395 except PrefectHTTPStatusError as exc: 

396 if exc.response.status_code == 404: 

397 return False 

398 raise 

399 

400 @async_dispatch(aenable) 

401 def enable(self: Self) -> bool: 

402 """ 

403 Enable an automation. 

404 

405 Raises: 

406 ValueError: If the automation does not have an id 

407 PrefectHTTPStatusError: If the automation cannot be enabled 

408 

409 Example: 

410 ```python 

411 auto = Automation.read(id = 123) 

412 auto.enable() 

413 ``` 

414 """ 

415 if self.id is None: 

416 raise ValueError("Can't enable an automation without an id") 

417 

418 with get_client(sync_client=True) as client: 

419 try: 

420 client.resume_automation(self.id) 

421 return True 

422 except PrefectHTTPStatusError as exc: 

423 if exc.response.status_code == 404: 

424 return False 

425 raise