Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/models/flows.py: 73%

1""" 

2Functions for interacting with flow ORM objects. 

3Intended for internal use by the Prefect REST API. 

4""" 

5 

6from typing import List, Optional, Sequence, TypeVar, Union 

7from uuid import UUID 

8 

9import sqlalchemy as sa 

10from sqlalchemy import delete, select 

11from sqlalchemy.ext.asyncio import AsyncSession 

12from sqlalchemy.sql import Select 

13 

14import prefect.server.schemas as schemas 

15from prefect.server.database import PrefectDBInterface, db_injector, orm_models 

16from prefect.server.events import clients 

17from prefect.server.events.schemas import lifecycle 

18from prefect.types._datetime import now 

19 

20T = TypeVar("T", bound=tuple) 

21 

22 

23async def emit_flow_created_event(flow: orm_models.Flow) -> None: 

24 """Emit an event when a flow is created.""" 

25 async with clients.PrefectServerEventsClient() as events_client: 

26 await events_client.emit(lifecycle.flow_created_event(flow, now("UTC"))) 

27 

28 

29async def emit_flow_updated_event(flow: orm_models.Flow) -> None: 

30 """Emit an event when a flow is updated.""" 

31 async with clients.PrefectServerEventsClient() as events_client: 

32 await events_client.emit(lifecycle.flow_updated_event(flow, now("UTC"))) 

33 

34 

35async def emit_flow_deleted_event(flow: orm_models.Flow) -> None: 

36 """Emit an event when a flow is deleted.""" 

37 async with clients.PrefectServerEventsClient() as events_client: 

38 await events_client.emit(lifecycle.flow_deleted_event(flow, now("UTC"))) 

39 

40 

41@db_injector 

42async def create_flow( 

43 db: PrefectDBInterface, session: AsyncSession, flow: schemas.core.Flow 

44) -> orm_models.Flow: 

45 """ 

46 Creates a new flow. 

47 

48 If a flow with the same name already exists, the existing flow is returned. 

49 

50 Args: 

51 session: a database session 

52 flow: a flow model 

53 

54 Returns: 

55 orm_models.Flow: the newly-created or existing flow 

56 """ 

57 

58 insert_stmt = ( 

59 db.queries.insert(db.Flow) 

60 .values(**flow.model_dump_for_orm(exclude_unset=True)) 

61 .on_conflict_do_nothing( 

62 index_elements=db.orm.flow_unique_upsert_columns, 

63 ) 

64 ) 

65 insert_result = await session.execute(insert_stmt) 

66 inserted = insert_result.rowcount == 1 

67 

68 query = ( 

69 sa.select(db.Flow) 

70 .where(db.Flow.name == flow.name) 

71 .limit(1) 

72 .execution_options(populate_existing=True) 

73 ) 

74 result = await session.execute(query) 

75 model = result.scalar_one() 

76 

77 if inserted: 

78 await emit_flow_created_event(model) 

79 

80 return model 

81 

82 

83@db_injector 

84async def update_flow( 

85 db: PrefectDBInterface, 

86 session: AsyncSession, 

87 flow_id: UUID, 

88 flow: schemas.actions.FlowUpdate, 

89) -> bool: 

90 """ 

91 Updates a flow. 

92 

93 Args: 

94 session: a database session 

95 flow_id: the flow id to update 

96 flow: a flow update model 

97 

98 Returns: 

99 bool: whether or not matching rows were found to update 

100 """ 

101 existing = await read_flow(session, flow_id) 

102 if existing is None: 

103 return False 

104 

105 update_stmt = ( 

106 sa.update(db.Flow) 

107 .where(db.Flow.id == flow_id) 

108 # exclude_unset=True allows us to only update values provided by 

109 # the user, ignoring any defaults on the model 

110 .values(**flow.model_dump_for_orm(exclude_unset=True)) 

111 ) 

112 await session.execute(update_stmt) 

113 

114 await session.refresh(existing) 

115 await emit_flow_updated_event(existing) 

116 return True 

117 

118 

119@db_injector 

120async def read_flow( 

121 db: PrefectDBInterface, session: AsyncSession, flow_id: UUID 

122) -> Optional[orm_models.Flow]: 

123 """ 

124 Reads a flow by id. 

125 

126 Args: 

127 session: A database session 

128 flow_id: a flow id 

129 

130 Returns: 

131 orm_models.Flow: the flow 

132 """ 

133 return await session.get(db.Flow, flow_id) 

134 

135 

136@db_injector 

137async def read_flow_by_name( 

138 db: PrefectDBInterface, session: AsyncSession, name: str 

139) -> Optional[orm_models.Flow]: 

140 """ 

141 Reads a flow by name. 

142 

143 Args: 

144 session: A database session 

145 name: a flow name 

146 

147 Returns: 

148 orm_models.Flow: the flow 

149 """ 

150 

151 result = await session.execute(select(db.Flow).filter_by(name=name)) 

152 return result.scalar() 

153 

154 

155async def _apply_flow_filters( 

156 db: PrefectDBInterface, 

157 query: Select[T], 

158 flow_filter: Union[schemas.filters.FlowFilter, None] = None, 

159 flow_run_filter: Union[schemas.filters.FlowRunFilter, None] = None, 

160 task_run_filter: Union[schemas.filters.TaskRunFilter, None] = None, 

161 deployment_filter: Union[schemas.filters.DeploymentFilter, None] = None, 

162 work_pool_filter: Union[schemas.filters.WorkPoolFilter, None] = None, 

163) -> Select[T]: 

164 """ 

165 Applies filters to a flow query as a combination of EXISTS subqueries. 

166 """ 

167 

168 if flow_filter: 

169 query = query.where(flow_filter.as_sql_filter()) 

170 

171 if deployment_filter or work_pool_filter: 

172 deployment_exists_clause = select(db.Deployment).where( 

173 db.Deployment.flow_id == db.Flow.id 

174 ) 

175 

176 if deployment_filter: 

177 deployment_exists_clause = deployment_exists_clause.where( 

178 deployment_filter.as_sql_filter(), 

179 ) 

180 

181 if work_pool_filter: 

182 deployment_exists_clause = deployment_exists_clause.join( 

183 db.WorkQueue, 

184 db.WorkQueue.id == db.Deployment.work_queue_id, 

185 ) 

186 deployment_exists_clause = deployment_exists_clause.join( 

187 db.WorkPool, 

188 db.WorkPool.id == db.WorkQueue.work_pool_id, 

189 ).where(work_pool_filter.as_sql_filter()) 

190 

191 query = query.where(deployment_exists_clause.exists()) 

192 

193 if flow_run_filter or task_run_filter: 

194 flow_run_exists_clause = select(db.FlowRun).where( 

195 db.FlowRun.flow_id == db.Flow.id 

196 ) 

197 

198 if flow_run_filter: 

199 flow_run_exists_clause = flow_run_exists_clause.where( 

200 flow_run_filter.as_sql_filter() 

201 ) 

202 

203 if task_run_filter: 

204 flow_run_exists_clause = flow_run_exists_clause.join( 

205 db.TaskRun, 

206 db.TaskRun.flow_run_id == db.FlowRun.id, 

207 ).where(task_run_filter.as_sql_filter()) 

208 

209 query = query.where(flow_run_exists_clause.exists()) 

210 

211 return query 

212 

213 

214@db_injector 

215async def read_flows( 

216 db: PrefectDBInterface, 

217 session: AsyncSession, 

218 flow_filter: Union[schemas.filters.FlowFilter, None] = None, 

219 flow_run_filter: Union[schemas.filters.FlowRunFilter, None] = None, 

220 task_run_filter: Union[schemas.filters.TaskRunFilter, None] = None, 

221 deployment_filter: Union[schemas.filters.DeploymentFilter, None] = None, 

222 work_pool_filter: Union[schemas.filters.WorkPoolFilter, None] = None, 

223 sort: schemas.sorting.FlowSort = schemas.sorting.FlowSort.NAME_ASC, 

224 offset: Union[int, None] = None, 

225 limit: Union[int, None] = None, 

226) -> Sequence[orm_models.Flow]: 

227 """ 

228 Read multiple flows. 

229 

230 Args: 

231 session: A database session 

232 flow_filter: only select flows that match these filters 

233 flow_run_filter: only select flows whose flow runs match these filters 

234 task_run_filter: only select flows whose task runs match these filters 

235 deployment_filter: only select flows whose deployments match these filters 

236 work_pool_filter: only select flows whose work pools match these filters 

237 offset: Query offset 

238 limit: Query limit 

239 

240 Returns: 

241 List[orm_models.Flow]: flows 

242 """ 

243 

244 query = select(db.Flow).order_by(*sort.as_sql_sort()) 

245 

246 query = await _apply_flow_filters( 

247 db, 

248 query, 

249 flow_filter=flow_filter, 

250 flow_run_filter=flow_run_filter, 

251 task_run_filter=task_run_filter, 

252 deployment_filter=deployment_filter, 

253 work_pool_filter=work_pool_filter, 

254 ) 

255 

256 if offset is not None: 

257 query = query.offset(offset) 

258 

259 if limit is not None: 259 ↛ 262line 259 didn't jump to line 262 because the condition on line 259 was always true

260 query = query.limit(limit) 

261 

262 result = await session.execute(query) 

263 return result.scalars().unique().all() 

264 

265 

266@db_injector 

267async def count_flows( 

268 db: PrefectDBInterface, 

269 session: AsyncSession, 

270 flow_filter: Union[schemas.filters.FlowFilter, None] = None, 

271 flow_run_filter: Union[schemas.filters.FlowRunFilter, None] = None, 

272 task_run_filter: Union[schemas.filters.TaskRunFilter, None] = None, 

273 deployment_filter: Union[schemas.filters.DeploymentFilter, None] = None, 

274 work_pool_filter: Union[schemas.filters.WorkPoolFilter, None] = None, 

275) -> int: 

276 """ 

277 Count flows. 

278 

279 Args: 

280 session: A database session 

281 flow_filter: only count flows that match these filters 

282 flow_run_filter: only count flows whose flow runs match these filters 

283 task_run_filter: only count flows whose task runs match these filters 

284 deployment_filter: only count flows whose deployments match these filters 

285 work_pool_filter: only count flows whose work pools match these filters 

286 

287 Returns: 

288 int: count of flows 

289 """ 

290 

291 query = select(sa.func.count(None)).select_from(db.Flow) 

292 

293 query = await _apply_flow_filters( 

294 db, 

295 query, 

296 flow_filter=flow_filter, 

297 flow_run_filter=flow_run_filter, 

298 task_run_filter=task_run_filter, 

299 deployment_filter=deployment_filter, 

300 work_pool_filter=work_pool_filter, 

301 ) 

302 

303 result = await session.execute(query) 

304 return result.scalar_one() 

305 

306 

307@db_injector 

308async def delete_flow( 

309 db: PrefectDBInterface, session: AsyncSession, flow_id: UUID 

310) -> bool: 

311 """ 

312 Delete a flow by id. 

313 

314 Args: 

315 session: A database session 

316 flow_id: a flow id 

317 

318 Returns: 

319 bool: whether or not the flow was deleted 

320 """ 

321 existing = await read_flow(session, flow_id) 

322 if existing is None: 

323 return False 

324 

325 await emit_flow_deleted_event(existing) 

326 

327 await session.execute(delete(db.Flow).where(db.Flow.id == flow_id)) 

328 return True 

329 

330 

331@db_injector 

332async def delete_flows( 

333 db: PrefectDBInterface, 

334 session: AsyncSession, 

335 flow_ids: List[UUID], 

336) -> List[UUID]: 

337 """ 

338 Delete multiple flows by their IDs. 

339 

340 This also deletes all associated deployments (hard delete). 

341 

342 Args: 

343 session: A database session 

344 flow_ids: a list of flow ids to delete 

345 

346 Returns: 

347 List[UUID]: the IDs of the flows that were deleted 

348 """ 

349 if not flow_ids: 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true

350 return [] 

351 

352 # Get existing flows 

353 result = await session.execute(select(db.Flow).where(db.Flow.id.in_(flow_ids))) 

354 existing_flows = list(result.scalars().all()) 

355 

356 if not existing_flows: 

357 return [] 

358 

359 existing_ids = [flow.id for flow in existing_flows] 

360 

361 for flow in existing_flows: 

362 await emit_flow_deleted_event(flow) 

363 

364 # Delete associated deployments (hard delete - cascade will handle related data) 

365 await session.execute( 

366 delete(db.Deployment).where(db.Deployment.flow_id.in_(existing_ids)) 

367 ) 

368 

369 # Delete the flows 

370 await session.execute(delete(db.Flow).where(db.Flow.id.in_(existing_ids))) 

371 

372 return existing_ids 

373 

374 

375@db_injector 

376async def read_flow_labels( 

377 db: PrefectDBInterface, 

378 session: AsyncSession, 

379 flow_id: UUID, 

380) -> Union[schemas.core.KeyValueLabels, None]: 

381 result = await session.execute(select(db.Flow.labels).where(db.Flow.id == flow_id)) 

382 

383 return result.scalar()