Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/database/query_components.py: 59%

1import datetime 

2from abc import ABC, abstractmethod 

3from collections import defaultdict 

4from collections.abc import Hashable, Iterable, Sequence 

5from functools import cached_property 

6from typing import ( 

7 Any, 

8 ClassVar, 

9 Literal, 

10 NamedTuple, 

11 Optional, 

12 Union, 

13 cast, 

14) 

15from uuid import UUID 

16 

17import sqlalchemy as sa 

18from cachetools import Cache, TTLCache 

19from jinja2 import Environment, PackageLoader, select_autoescape 

20from sqlalchemy import orm 

21from sqlalchemy.dialects import postgresql, sqlite 

22from sqlalchemy.exc import NoResultFound 

23from sqlalchemy.ext.asyncio import AsyncSession 

24from sqlalchemy.sql.type_api import TypeEngine 

25from typing_extensions import TypeVar 

26 

27from prefect.server import models, schemas 

28from prefect.server.database import orm_models 

29from prefect.server.database.dependencies import db_injector 

30from prefect.server.database.interface import PrefectDBInterface 

31from prefect.server.exceptions import FlowRunGraphTooLarge, ObjectNotFoundError 

32from prefect.server.schemas.graph import Edge, Graph, GraphArtifact, GraphState, Node 

33from prefect.server.schemas.states import StateType 

34from prefect.server.utilities.database import UUID as UUIDTypeDecorator 

35from prefect.server.utilities.database import Timestamp, bindparams_from_clause 

36from prefect.types._datetime import DateTime 

37 

38T = TypeVar("T", infer_variance=True) 

39 

40 

41class FlowRunGraphV2Node(NamedTuple): 

42 kind: Literal["flow-run", "task-run"] 

43 id: UUID 

44 label: str 

45 state_type: StateType 

46 start_time: DateTime 

47 end_time: Optional[DateTime] 

48 parent_ids: Optional[list[UUID]] 

49 child_ids: Optional[list[UUID]] 

50 encapsulating_ids: Optional[list[UUID]] 

51 

52 

53ONE_HOUR = 60 * 60 

54 

55 

56jinja_env: Environment = Environment( 

57 loader=PackageLoader("prefect.server.database", package_path="sql"), 

58 autoescape=select_autoescape(), 

59 trim_blocks=True, 

60) 

61 

62 

63class BaseQueryComponents(ABC): 

64 """ 

65 Abstract base class used to inject dialect-specific SQL operations into Prefect. 

66 """ 

67 

68 _configuration_cache: ClassVar[Cache[str, dict[str, Any]]] = TTLCache( 

69 maxsize=100, ttl=ONE_HOUR 

70 ) 

71 

72 def unique_key(self) -> tuple[Hashable, ...]: 

73 """ 

74 Returns a key used to determine whether to instantiate a new DB interface. 

75 """ 

76 return (self.__class__,) 

77 

78 # --- dialect-specific SqlAlchemy bindings 

79 

80 @abstractmethod 

81 def insert( 

82 self, obj: type[orm_models.Base] 

83 ) -> Union[postgresql.Insert, sqlite.Insert]: 

84 """dialect-specific insert statement""" 

85 

86 # --- dialect-specific JSON handling 

87 

88 @property 

89 @abstractmethod 

90 def uses_json_strings(self) -> bool: 

91 """specifies whether the configured dialect returns JSON as strings""" 

92 

93 @abstractmethod 

94 def cast_to_json(self, json_obj: sa.ColumnElement[T]) -> sa.ColumnElement[T]: 

95 """casts to JSON object if necessary""" 

96 

97 @abstractmethod 

98 def build_json_object( 

99 self, *args: Union[str, sa.ColumnElement[Any]] 

100 ) -> sa.ColumnElement[Any]: 

101 """builds a JSON object from sequential key-value pairs""" 

102 

103 @abstractmethod 

104 def json_arr_agg(self, json_array: sa.ColumnElement[Any]) -> sa.ColumnElement[Any]: 

105 """aggregates a JSON array""" 

106 

107 # --- dialect-optimized subqueries 

108 

109 @abstractmethod 

110 def make_timestamp_intervals( 110 ↛ exitline 110 didn't return from function 'make_timestamp_intervals' because

111 self, 

112 start_time: datetime.datetime, 

113 end_time: datetime.datetime, 

114 interval: datetime.timedelta, 

115 ) -> sa.Select[tuple[datetime.datetime, datetime.datetime]]: ... 

116 

117 @abstractmethod 

118 def set_state_id_on_inserted_flow_runs_statement( 118 ↛ exitline 118 didn't return from function 'set_state_id_on_inserted_flow_runs_statement' because

119 self, 

120 inserted_flow_run_ids: Sequence[UUID], 

121 insert_flow_run_states: Iterable[dict[str, Any]], 

122 ) -> sa.Update: ... 

123 

124 @db_injector 

125 def get_scheduled_flow_runs_from_work_queues( 

126 self, 

127 db: PrefectDBInterface, 

128 limit_per_queue: Optional[int] = None, 

129 work_queue_ids: Optional[list[UUID]] = None, 

130 scheduled_before: Optional[DateTime] = None, 

131 ) -> sa.Select[tuple[orm_models.FlowRun, UUID]]: 

132 """ 

133 Returns all scheduled runs in work queues, subject to provided parameters. 

134 

135 This query returns a `(orm_models.FlowRun, orm_models.WorkQueue.id)` pair; calling 

136 `result.all()` will return both; calling `result.scalars().unique().all()` 

137 will return only the flow run because it grabs the first result. 

138 """ 

139 

140 FlowRun, WorkQueue = db.FlowRun, db.WorkQueue 

141 

142 # get any work queues that have a concurrency limit, and compute available 

143 # slots as their limit less the number of running flows 

144 concurrency_queues = ( 

145 sa.select( 

146 WorkQueue.id, 

147 sa.func.greatest( 

148 0, 

149 WorkQueue.concurrency_limit - sa.func.count(FlowRun.id), 

150 ).label("available_slots"), 

151 ) 

152 .select_from(WorkQueue) 

153 .join( 

154 FlowRun, 

155 sa.and_( 

156 FlowRun.work_queue_name == WorkQueue.name, 

157 FlowRun.state_type.in_( 

158 (StateType.RUNNING, StateType.PENDING, StateType.CANCELLING) 

159 ), 

160 ), 

161 isouter=True, 

162 ) 

163 .where(WorkQueue.concurrency_limit.is_not(None)) 

164 .group_by(WorkQueue.id) 

165 .cte("concurrency_queues") 

166 ) 

167 

168 # use the available slots information to generate a join 

169 # for all scheduled runs 

170 scheduled_flow_runs, join_criteria = self._get_scheduled_flow_runs_join( 

171 work_queue_query=concurrency_queues, 

172 limit_per_queue=limit_per_queue, 

173 scheduled_before=scheduled_before, 

174 ) 

175 

176 # starting with the work queue table, join the limited queues to get the 

177 # concurrency information and the scheduled flow runs to load all applicable 

178 # runs. this will return all the scheduled runs allowed by the parameters 

179 query = ( 

180 # return a flow run and work queue id 

181 sa.select( 

182 orm.aliased(FlowRun, scheduled_flow_runs), WorkQueue.id.label("wq_id") 

183 ) 

184 .select_from(WorkQueue) 

185 .join( 

186 concurrency_queues, 

187 WorkQueue.id == concurrency_queues.c.id, 

188 isouter=True, 

189 ) 

190 .join(scheduled_flow_runs, join_criteria) 

191 .where( 

192 WorkQueue.is_paused.is_(False), 

193 WorkQueue.id.in_(work_queue_ids) if work_queue_ids else sa.true(), 

194 ) 

195 .order_by( 

196 scheduled_flow_runs.c.next_scheduled_start_time, 

197 scheduled_flow_runs.c.id, 

198 ) 

199 ) 

200 

201 return query 

202 

203 @db_injector 

204 def _get_scheduled_flow_runs_join( 

205 self, 

206 db: PrefectDBInterface, 

207 work_queue_query: sa.CTE, 

208 limit_per_queue: Optional[int], 

209 scheduled_before: Optional[DateTime], 

210 ) -> tuple[sa.FromClause, sa.ColumnExpressionArgument[bool]]: 

211 """Used by self.get_scheduled_flow_runs_from_work_queue, allowing just 

212 this function to be changed on a per-dialect basis""" 

213 

214 FlowRun = db.FlowRun 

215 

216 # precompute for readability 

217 scheduled_before_clause = ( 

218 FlowRun.next_scheduled_start_time <= scheduled_before 

219 if scheduled_before is not None 

220 else sa.true() 

221 ) 

222 

223 # get scheduled flow runs with lateral join where the limit is the 

224 # available slots per queue 

225 scheduled_flow_runs = ( 

226 sa.select(FlowRun) 

227 .where( 

228 FlowRun.work_queue_name == db.WorkQueue.name, 

229 FlowRun.state_type == StateType.SCHEDULED, 

230 scheduled_before_clause, 

231 ) 

232 .with_for_update(skip_locked=True) 

233 # priority given to runs with earlier next_scheduled_start_time 

234 .order_by(FlowRun.next_scheduled_start_time) 

235 # if null, no limit will be applied 

236 .limit(sa.func.least(limit_per_queue, work_queue_query.c.available_slots)) 

237 .lateral("scheduled_flow_runs") 

238 ) 

239 

240 # Perform a cross-join 

241 join_criteria = sa.true() 

242 

243 return scheduled_flow_runs, join_criteria 

244 

245 # ------------------------------------------------------- 

246 # Workers 

247 # ------------------------------------------------------- 

248 

249 @property 

250 @abstractmethod 

251 def _get_scheduled_flow_runs_from_work_pool_template_path(self) -> str: 

252 """ 

253 Template for the query to get scheduled flow runs from a work pool 

254 """ 

255 

256 @db_injector 

257 async def get_scheduled_flow_runs_from_work_pool( 

258 self, 

259 db: PrefectDBInterface, 

260 session: AsyncSession, 

261 limit: Optional[int] = None, 

262 worker_limit: Optional[int] = None, 

263 queue_limit: Optional[int] = None, 

264 work_pool_ids: Optional[list[UUID]] = None, 

265 work_queue_ids: Optional[list[UUID]] = None, 

266 scheduled_before: Optional[DateTime] = None, 

267 scheduled_after: Optional[DateTime] = None, 

268 respect_queue_priorities: bool = False, 

269 ) -> list[schemas.responses.WorkerFlowRunResponse]: 

270 template = jinja_env.get_template( 

271 self._get_scheduled_flow_runs_from_work_pool_template_path 

272 ) 

273 

274 raw_query = sa.text( 

275 template.render( 

276 work_pool_ids=work_pool_ids, 

277 work_queue_ids=work_queue_ids, 

278 respect_queue_priorities=respect_queue_priorities, 

279 scheduled_before=scheduled_before, 

280 scheduled_after=scheduled_after, 

281 ) 

282 ) 

283 

284 bindparams: list[sa.BindParameter[Any]] = [] 

285 

286 if scheduled_before: 

287 bindparams.append( 

288 sa.bindparam("scheduled_before", scheduled_before, type_=Timestamp) 

289 ) 

290 

291 if scheduled_after: 

292 bindparams.append( 

293 sa.bindparam("scheduled_after", scheduled_after, type_=Timestamp) 

294 ) 

295 

296 # if work pool IDs were provided, bind them 

297 if work_pool_ids: 297 ↛ 309line 297 didn't jump to line 309 because the condition on line 297 was always true

298 assert all(isinstance(i, UUID) for i in work_pool_ids) 

299 bindparams.append( 

300 sa.bindparam( 

301 "work_pool_ids", 

302 work_pool_ids, 

303 expanding=True, 

304 type_=UUIDTypeDecorator, 

305 ) 

306 ) 

307 

308 # if work queue IDs were provided, bind them 

309 if work_queue_ids: 309 ↛ 310line 309 didn't jump to line 310 because the condition on line 309 was never true

310 assert all(isinstance(i, UUID) for i in work_queue_ids) 

311 bindparams.append( 

312 sa.bindparam( 

313 "work_queue_ids", 

314 work_queue_ids, 

315 expanding=True, 

316 type_=UUIDTypeDecorator, 

317 ) 

318 ) 

319 

320 query = raw_query.bindparams( 

321 *bindparams, 

322 limit=1000 if limit is None else limit, 

323 worker_limit=1000 if worker_limit is None else worker_limit, 

324 queue_limit=1000 if queue_limit is None else queue_limit, 

325 ) 

326 

327 FlowRun = db.FlowRun 

328 orm_query = ( 

329 sa.select( 

330 sa.column("run_work_pool_id", UUIDTypeDecorator), 

331 sa.column("run_work_queue_id", UUIDTypeDecorator), 

332 FlowRun, 

333 ) 

334 .from_statement(query) 

335 # indicate that the state relationship isn't being loaded 

336 .options(orm.noload(FlowRun.state)) 

337 ) 

338 

339 result: sa.Result[ 

340 tuple[UUID, UUID, orm_models.FlowRun] 

341 ] = await session.execute(orm_query) 

342 

343 return [ 

344 schemas.responses.WorkerFlowRunResponse( 

345 work_pool_id=run_work_pool_id, 

346 work_queue_id=run_work_queue_id, 

347 flow_run=schemas.core.FlowRun.model_validate( 

348 flow_run, from_attributes=True 

349 ), 

350 ) 

351 for (run_work_pool_id, run_work_queue_id, flow_run) in result.t 

352 ] 

353 

354 @db_injector 

355 async def read_configuration_value( 

356 self, db: PrefectDBInterface, session: AsyncSession, key: str 

357 ) -> Optional[dict[str, Any]]: 

358 """ 

359 Read a configuration value by key. 

360 

361 Configuration values should not be changed at run time, so retrieved 

362 values are cached in memory. 

363 

364 The main use of configurations is encrypting blocks, this speeds up nested 

365 block document queries. 

366 """ 

367 Configuration = db.Configuration 

368 value = None 

369 try: 

370 value = self._configuration_cache[key] 

371 except KeyError: 

372 query = sa.select(Configuration).where(Configuration.key == key) 

373 if (configuration := await session.scalar(query)) is not None: 373 ↛ anywhereline 373 didn't jump anywhere: it always raised an exception.

374 value = self._configuration_cache[key] = configuration.value 

375 return value 

376 

377 def clear_configuration_value_cache_for_key(self, key: str) -> None: 

378 """Removes a configuration key from the cache.""" 

379 self._configuration_cache.pop(key, None) 

380 

381 @cached_property 

382 def _flow_run_graph_v2_query(self): 

383 query = self._build_flow_run_graph_v2_query() 

384 param_names = set(bindparams_from_clause(query)) 

385 required = {"flow_run_id", "max_nodes", "since"} 

386 assert param_names >= required, ( 

387 "_build_flow_run_graph_v2_query result is missing required bind params: " 

388 f"{sorted(required - param_names)}" 

389 ) 

390 return query 

391 

392 @abstractmethod 

393 def _build_flow_run_graph_v2_query(self) -> sa.Select[FlowRunGraphV2Node]: 

394 """The flow run graph query, per database flavour 

395 

396 The query must accept the following bind parameters: 

397 

398 flow_run_id: UUID 

399 since: DateTime 

400 max_nodes: int 

401 

402 """ 

403 

404 @db_injector 

405 async def flow_run_graph_v2( 

406 self, 

407 db: PrefectDBInterface, 

408 session: AsyncSession, 

409 flow_run_id: UUID, 

410 since: DateTime, 

411 max_nodes: int, 

412 max_artifacts: int, 

413 ) -> Graph: 

414 """Returns the query that selects all of the nodes and edges for a flow run graph (version 2).""" 

415 FlowRun = db.FlowRun 

416 result = await session.execute( 

417 sa.select( 

418 sa.func.coalesce( 

419 FlowRun.start_time, FlowRun.expected_start_time, type_=Timestamp 

420 ), 

421 FlowRun.end_time, 

422 ).where(FlowRun.id == flow_run_id) 

423 ) 

424 try: 

425 start_time, end_time = result.t.one() 

426 except NoResultFound: 

427 raise ObjectNotFoundError(f"Flow run {flow_run_id} not found") 

428 

429 query = self._flow_run_graph_v2_query 

430 results = await session.execute( 

431 query, 

432 params=dict(flow_run_id=flow_run_id, since=since, max_nodes=max_nodes + 1), 

433 ) 

434 

435 graph_artifacts = await self._get_flow_run_graph_artifacts( 

436 db, session, flow_run_id, max_artifacts 

437 ) 

438 graph_states = await self._get_flow_run_graph_states(session, flow_run_id) 

439 

440 nodes: list[tuple[UUID, Node]] = [] 

441 root_node_ids: list[UUID] = [] 

442 

443 for row in results.t: 

444 if not row.parent_ids: 

445 root_node_ids.append(row.id) 

446 

447 nodes.append( 

448 ( 

449 row.id, 

450 Node( 

451 kind=row.kind, 

452 id=row.id, 

453 label=row.label, 

454 state_type=row.state_type, 

455 start_time=row.start_time, 

456 end_time=row.end_time, 

457 parents=[Edge(id=id) for id in row.parent_ids or []], 

458 children=[Edge(id=id) for id in row.child_ids or []], 

459 encapsulating=[ 

460 Edge(id=id) 

461 # ensure encapsulating_ids is deduplicated 

462 for id in dict.fromkeys(row.encapsulating_ids or ()) 

463 ], 

464 artifacts=graph_artifacts.get(row.id, []), 

465 ), 

466 ) 

467 ) 

468 

469 if len(nodes) > max_nodes: 

470 raise FlowRunGraphTooLarge( 

471 f"The graph of flow run {flow_run_id} has more than " 

472 f"{max_nodes} nodes." 

473 ) 

474 

475 return Graph( 

476 start_time=start_time, 

477 end_time=end_time, 

478 root_node_ids=root_node_ids, 

479 nodes=nodes, 

480 artifacts=graph_artifacts.get(None, []), 

481 states=graph_states, 

482 ) 

483 

484 async def _get_flow_run_graph_artifacts( 

485 self, 

486 db: PrefectDBInterface, 

487 session: AsyncSession, 

488 flow_run_id: UUID, 

489 max_artifacts: int, 

490 ) -> dict[Optional[UUID], list[GraphArtifact]]: 

491 """Get the artifacts for a flow run grouped by task run id. 

492 

493 Does not recurse into subflows. 

494 Artifacts for the flow run without a task run id are grouped under None. 

495 """ 

496 Artifact, ArtifactCollection = db.Artifact, db.ArtifactCollection 

497 

498 query = ( 

499 sa.select(Artifact, ArtifactCollection.id.label("latest_in_collection_id")) 

500 .where(Artifact.flow_run_id == flow_run_id, Artifact.type != "result") 

501 .join( 

502 ArtifactCollection, 

503 onclause=sa.and_( 

504 ArtifactCollection.key == Artifact.key, 

505 ArtifactCollection.latest_id == Artifact.id, 

506 ), 

507 isouter=True, 

508 ) 

509 .order_by(Artifact.created.asc()) 

510 .limit(max_artifacts) 

511 ) 

512 

513 results = await session.execute(query) 

514 

515 artifacts_by_task: dict[Optional[UUID], list[GraphArtifact]] = defaultdict(list) 

516 for artifact, latest_in_collection_id in results.t: 

517 artifacts_by_task[artifact.task_run_id].append( 

518 GraphArtifact( 

519 id=artifact.id, 

520 created=artifact.created, 

521 key=artifact.key, 

522 type=artifact.type, 

523 # We're only using the data field for progress artifacts for now 

524 data=artifact.data if artifact.type == "progress" else None, 

525 is_latest=artifact.key is None 

526 or latest_in_collection_id is not None, # pyright: ignore[reportUnnecessaryComparison] 

527 ) 

528 ) 

529 

530 return dict(artifacts_by_task) 

531 

532 async def _get_flow_run_graph_states( 

533 self, session: AsyncSession, flow_run_id: UUID 

534 ) -> list[GraphState]: 

535 """Get the flow run states for a flow run graph.""" 

536 states = await models.flow_run_states.read_flow_run_states(session, flow_run_id) 

537 return [ 

538 GraphState.model_validate(state, from_attributes=True) for state in states 

539 ] 

540 

541 

542class AsyncPostgresQueryComponents(BaseQueryComponents): 

543 # --- Postgres-specific SqlAlchemy bindings 

544 

545 def insert(self, obj: type[orm_models.Base]) -> postgresql.Insert: 

546 return postgresql.insert(obj) 

547 

548 # --- Postgres-specific JSON handling 

549 

550 @property 

551 def uses_json_strings(self) -> bool: 

552 return False 

553 

554 def cast_to_json(self, json_obj: sa.ColumnElement[T]) -> sa.ColumnElement[T]: 

555 return json_obj 

556 

557 def build_json_object( 

558 self, *args: Union[str, sa.ColumnElement[Any]] 

559 ) -> sa.ColumnElement[Any]: 

560 return sa.func.jsonb_build_object(*args) 

561 

562 def json_arr_agg(self, json_array: sa.ColumnElement[Any]) -> sa.ColumnElement[Any]: 

563 return sa.func.jsonb_agg(json_array) 

564 

565 # --- Postgres-optimized subqueries 

566 

567 def make_timestamp_intervals( 

568 self, 

569 start_time: datetime.datetime, 

570 end_time: datetime.datetime, 

571 interval: datetime.timedelta, 

572 ) -> sa.Select[tuple[datetime.datetime, datetime.datetime]]: 

573 dt = sa.func.generate_series( 

574 start_time, end_time, interval, type_=Timestamp() 

575 ).column_valued("dt") 

576 return ( 

577 sa.select( 

578 dt.label("interval_start"), 

579 sa.type_coerce( 

580 dt + sa.bindparam("interval", interval, type_=sa.Interval()), 

581 type_=Timestamp(), 

582 ).label("interval_end"), 

583 ) 

584 .where(dt < end_time) 

585 .limit(500) # grab at most 500 intervals 

586 ) 

587 

588 @db_injector 

589 def set_state_id_on_inserted_flow_runs_statement( 

590 self, 

591 db: PrefectDBInterface, 

592 inserted_flow_run_ids: Sequence[UUID], 

593 insert_flow_run_states: Iterable[dict[str, Any]], 

594 ) -> sa.Update: 

595 """Given a list of flow run ids and associated states, set the state_id 

596 to the appropriate state for all flow runs""" 

597 # postgres supports `UPDATE ... FROM` syntax 

598 FlowRun, FlowRunState = db.FlowRun, db.FlowRunState 

599 stmt = ( 

600 sa.update(FlowRun) 

601 .where( 

602 FlowRun.id.in_(inserted_flow_run_ids), 

603 FlowRunState.flow_run_id == FlowRun.id, 

604 FlowRunState.id.in_([r["id"] for r in insert_flow_run_states]), 

605 ) 

606 .values(state_id=FlowRunState.id) 

607 # no need to synchronize as these flow runs are entirely new 

608 .execution_options(synchronize_session=False) 

609 ) 

610 return stmt 

611 

612 @property 

613 def _get_scheduled_flow_runs_from_work_pool_template_path(self) -> str: 

614 """ 

615 Template for the query to get scheduled flow runs from a work pool 

616 """ 

617 return "postgres/get-runs-from-worker-queues.sql.jinja" 

618 

619 @db_injector 

620 def _build_flow_run_graph_v2_query( 

621 self, db: PrefectDBInterface 

622 ) -> sa.Select[FlowRunGraphV2Node]: 

623 """Postgresql version of the V2 FlowRun graph data query 

624 

625 This SQLA query is built just once and then cached per DB interface 

626 

627 """ 

628 # the parameters this query takes as inputs 

629 param_flow_run_id = sa.bindparam("flow_run_id", type_=UUIDTypeDecorator) 

630 param_since = sa.bindparam("since", type_=Timestamp) 

631 param_max_nodes = sa.bindparam("max_nodes", type_=sa.Integer) 

632 

633 Flow, FlowRun, TaskRun = db.Flow, db.FlowRun, db.TaskRun 

634 input = sa.func.jsonb_each(TaskRun.task_inputs).table_valued( 

635 "key", "value", name="input" 

636 ) 

637 argument = ( 

638 sa.func.jsonb_array_elements(input.c.value, type_=postgresql.JSONB()) 

639 .table_valued(sa.column("value", postgresql.JSONB())) 

640 .render_derived(name="argument") 

641 ) 

642 edges = ( 

643 sa.select( 

644 sa.case((FlowRun.id.is_not(None), "flow-run"), else_="task-run").label( 

645 "kind" 

646 ), 

647 sa.func.coalesce(FlowRun.id, TaskRun.id).label("id"), 

648 sa.func.coalesce(Flow.name + " / " + FlowRun.name, TaskRun.name).label( 

649 "label" 

650 ), 

651 sa.func.coalesce(FlowRun.state_type, TaskRun.state_type).label( 

652 "state_type" 

653 ), 

654 sa.func.coalesce( 

655 FlowRun.start_time, 

656 FlowRun.expected_start_time, 

657 TaskRun.start_time, 

658 TaskRun.expected_start_time, 

659 ).label("start_time"), 

660 sa.func.coalesce( 

661 FlowRun.end_time, 

662 TaskRun.end_time, 

663 sa.case( 

664 ( 

665 TaskRun.state_type == StateType.COMPLETED, 

666 TaskRun.expected_start_time, 

667 ), 

668 else_=sa.null(), 

669 ), 

670 ).label("end_time"), 

671 sa.cast(argument.c.value["id"].astext, type_=UUIDTypeDecorator).label( 

672 "parent" 

673 ), 

674 (input.c.key == "__parents__").label("has_encapsulating_task"), 

675 ) 

676 .join_from(TaskRun, input, onclause=sa.true(), isouter=True) 

677 .join(argument, onclause=sa.true(), isouter=True) 

678 .join( 

679 FlowRun, 

680 isouter=True, 

681 onclause=FlowRun.parent_task_run_id == TaskRun.id, 

682 ) 

683 .join(Flow, isouter=True, onclause=Flow.id == FlowRun.flow_id) 

684 .where( 

685 TaskRun.flow_run_id == param_flow_run_id, 

686 sa.func.coalesce( 

687 FlowRun.start_time, 

688 FlowRun.expected_start_time, 

689 TaskRun.start_time, 

690 TaskRun.expected_start_time, 

691 ).is_not(None), 

692 ) 

693 # -- the order here is important to speed up building the two sets of 

694 # -- edges in the with_parents and with_children CTEs below 

695 .order_by(sa.func.coalesce(FlowRun.id, TaskRun.id)) 

696 ).cte("edges") 

697 children, parents = edges.alias("children"), edges.alias("parents") 

698 with_encapsulating = ( 

699 sa.select( 

700 children.c.id, 

701 sa.func.array_agg( 

702 postgresql.aggregate_order_by(parents.c.id, parents.c.start_time) 

703 ).label("encapsulating_ids"), 

704 ) 

705 .join(parents, onclause=parents.c.id == children.c.parent) 

706 .where(children.c.has_encapsulating_task.is_(True)) 

707 .group_by(children.c.id) 

708 ).cte("with_encapsulating") 

709 with_parents = ( 

710 sa.select( 

711 children.c.id, 

712 sa.func.array_agg( 

713 postgresql.aggregate_order_by(parents.c.id, parents.c.start_time) 

714 ).label("parent_ids"), 

715 ) 

716 .join(parents, onclause=parents.c.id == children.c.parent) 

717 .where(children.c.has_encapsulating_task.is_distinct_from(True)) 

718 .group_by(children.c.id) 

719 .cte("with_parents") 

720 ) 

721 with_children = ( 

722 sa.select( 

723 parents.c.id, 

724 sa.func.array_agg( 

725 postgresql.aggregate_order_by(children.c.id, children.c.start_time) 

726 ).label("child_ids"), 

727 ) 

728 .join(children, onclause=children.c.parent == parents.c.id) 

729 .where(children.c.has_encapsulating_task.is_distinct_from(True)) 

730 .group_by(parents.c.id) 

731 .cte("with_children") 

732 ) 

733 

734 graph = ( 

735 sa.select( 

736 edges.c.kind, 

737 edges.c.id, 

738 edges.c.label, 

739 edges.c.state_type, 

740 edges.c.start_time, 

741 edges.c.end_time, 

742 with_parents.c.parent_ids, 

743 with_children.c.child_ids, 

744 with_encapsulating.c.encapsulating_ids, 

745 ) 

746 .distinct(edges.c.id) 

747 .join(with_parents, isouter=True, onclause=with_parents.c.id == edges.c.id) 

748 .join( 

749 with_children, isouter=True, onclause=with_children.c.id == edges.c.id 

750 ) 

751 .join( 

752 with_encapsulating, 

753 isouter=True, 

754 onclause=with_encapsulating.c.id == edges.c.id, 

755 ) 

756 .cte("nodes") 

757 ) 

758 query = ( 

759 sa.select( 

760 graph.c.kind, 

761 graph.c.id, 

762 graph.c.label, 

763 graph.c.state_type, 

764 graph.c.start_time, 

765 graph.c.end_time, 

766 graph.c.parent_ids, 

767 graph.c.child_ids, 

768 graph.c.encapsulating_ids, 

769 ) 

770 .where(sa.or_(graph.c.end_time.is_(None), graph.c.end_time >= param_since)) 

771 .order_by(graph.c.start_time, graph.c.end_time) 

772 .limit(param_max_nodes) 

773 ) 

774 return cast(sa.Select[FlowRunGraphV2Node], query) 

775 

776 

777class UUIDList(sa.TypeDecorator[list[UUID]]): 

778 """Map a JSON list of strings back to a list of UUIDs at the result loading stage""" 

779 

780 impl: Union[TypeEngine[Any], type[TypeEngine[Any]]] = sa.JSON() 

781 cache_ok: Optional[bool] = True 

782 

783 def process_result_value( 

784 self, value: Optional[list[Union[str, UUID]]], dialect: sa.Dialect 

785 ) -> Optional[list[UUID]]: 

786 if value is None: 

787 return value 

788 return [v if isinstance(v, UUID) else UUID(v) for v in value] 

789 

790 

791class AioSqliteQueryComponents(BaseQueryComponents): 

792 # --- Sqlite-specific SqlAlchemy bindings 

793 

794 def insert(self, obj: type[orm_models.Base]) -> sqlite.Insert: