Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/models/flow_runs.py: 63%
221 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1"""
2Functions for interacting with flow run ORM objects.
3Intended for internal use by the Prefect REST API.
4"""
6import contextlib
7import datetime
8from itertools import chain
9from typing import (
10 TYPE_CHECKING,
11 Any,
12 Dict,
13 List,
14 Optional,
15 Sequence,
16 Tuple,
17 Type,
18 TypeVar,
19 Union,
20 cast,
21)
22from uuid import UUID
24import sqlalchemy as sa
25from sqlalchemy import delete, select
26from sqlalchemy.ext.asyncio import AsyncSession
27from sqlalchemy.orm import load_only, selectinload
28from sqlalchemy.sql import Select
30import prefect.server.models as models
31import prefect.server.schemas as schemas
32from prefect.logging.loggers import get_logger
33from prefect.server.database import PrefectDBInterface, db_injector, orm_models
34from prefect.server.exceptions import ObjectNotFoundError
35from prefect.server.orchestration.core_policy import MinimalFlowPolicy
36from prefect.server.orchestration.global_policy import GlobalFlowPolicy
37from prefect.server.orchestration.policies import (
38 FlowRunOrchestrationPolicy,
39)
40from prefect.server.orchestration.rules import FlowOrchestrationContext
41from prefect.server.schemas.core import TaskRunResult
42from prefect.server.schemas.graph import Graph
43from prefect.server.schemas.responses import OrchestrationResult
44from prefect.server.schemas.states import State
45from prefect.server.utilities.schemas import PrefectBaseModel
46from prefect.settings import (
47 PREFECT_API_MAX_FLOW_RUN_GRAPH_ARTIFACTS,
48 PREFECT_API_MAX_FLOW_RUN_GRAPH_NODES,
49)
50from prefect.types import KeyValueLabels
51from prefect.types._datetime import DateTime, earliest_possible_datetime, now
53if TYPE_CHECKING: 53 ↛ 54line 53 didn't jump to line 54 because the condition on line 53 was never true
54 import logging
56logger: "logging.Logger" = get_logger("flow_runs")
59T = TypeVar("T", bound=tuple[Any, ...])
62@db_injector
63async def create_flow_run(
64 db: PrefectDBInterface,
65 session: AsyncSession,
66 flow_run: schemas.core.FlowRun,
67 orchestration_parameters: Optional[dict[str, Any]] = None,
68) -> orm_models.FlowRun:
69 """Creates a new flow run.
71 If the provided flow run has a state attached, it will also be created.
73 Args:
74 session: a database session
75 flow_run: a flow run model
77 Returns:
78 orm_models.FlowRun: the newly-created flow run
79 """
80 right_now = now("UTC")
81 # model: Union[orm_models.FlowRun, None] = None
83 flow_run.labels = await with_system_labels_for_flow_run(
84 session=session, flow_run=flow_run
85 )
87 flow_run_dict = dict(
88 **flow_run.model_dump_for_orm(
89 exclude={
90 "created",
91 "state",
92 "estimated_run_time",
93 "estimated_start_time_delta",
94 },
95 exclude_unset=True,
96 ),
97 created=right_now,
98 )
100 # if no idempotency key was provided, create the run directly
101 if not flow_run.idempotency_key:
102 model = db.FlowRun(**flow_run_dict)
103 session.add(model)
104 await session.flush()
106 # otherwise let the database take care of enforcing idempotency
107 else:
108 insert_stmt = (
109 db.queries.insert(db.FlowRun)
110 .values(**flow_run_dict)
111 .on_conflict_do_nothing(
112 index_elements=db.orm.flow_run_unique_upsert_columns,
113 )
114 )
115 await session.execute(insert_stmt)
117 # read the run to see if idempotency was applied or not
118 query = (
119 sa.select(db.FlowRun)
120 .where(
121 sa.and_(
122 db.FlowRun.flow_id == flow_run.flow_id,
123 db.FlowRun.idempotency_key == flow_run.idempotency_key,
124 )
125 )
126 .limit(1)
127 .execution_options(populate_existing=True)
128 .options(
129 selectinload(db.FlowRun.work_queue).selectinload(db.WorkQueue.work_pool)
130 )
131 )
132 result = await session.execute(query)
133 model = result.scalar_one()
135 # if the flow run was created in this function call then we need to set the
136 # state. If it was created idempotently, the created time won't match.
137 if model.created == right_now and flow_run.state:
138 await models.flow_runs.set_flow_run_state(
139 session=session,
140 flow_run_id=model.id,
141 state=flow_run.state,
142 force=True,
143 orchestration_parameters=orchestration_parameters,
144 )
145 return model
148@db_injector
149async def update_flow_run(
150 db: PrefectDBInterface,
151 session: AsyncSession,
152 flow_run_id: UUID,
153 flow_run: schemas.actions.FlowRunUpdate,
154) -> bool:
155 """
156 Updates a flow run.
158 Args:
159 session: a database session
160 flow_run_id: the flow run id to update
161 flow_run: a flow run model
163 Returns:
164 bool: whether or not matching rows were found to update
165 """
166 update_stmt = (
167 sa.update(db.FlowRun)
168 .where(db.FlowRun.id == flow_run_id)
169 # exclude_unset=True allows us to only update values provided by
170 # the user, ignoring any defaults on the model
171 .values(**flow_run.model_dump_for_orm(exclude_unset=True))
172 )
173 result = await session.execute(update_stmt)
174 return result.rowcount > 0
177@db_injector
178async def read_flow_run(
179 db: PrefectDBInterface,
180 session: AsyncSession,
181 flow_run_id: UUID,
182 for_update: bool = False,
183) -> Optional[orm_models.FlowRun]:
184 """
185 Reads a flow run by id.
187 Args:
188 session: A database session
189 flow_run_id: a flow run id
191 Returns:
192 orm_models.FlowRun: the flow run
193 """
194 select = (
195 sa.select(db.FlowRun)
196 .where(db.FlowRun.id == flow_run_id)
197 .options(
198 selectinload(db.FlowRun.work_queue).selectinload(db.WorkQueue.work_pool)
199 )
200 )
202 if for_update:
203 select = select.with_for_update()
205 result = await session.execute(select)
206 return result.scalar()
209async def _apply_flow_run_filters(
210 db: PrefectDBInterface,
211 query: Select[T],
212 flow_filter: Optional[schemas.filters.FlowFilter] = None,
213 flow_run_filter: Optional[schemas.filters.FlowRunFilter] = None,
214 task_run_filter: Optional[schemas.filters.TaskRunFilter] = None,
215 deployment_filter: Optional[schemas.filters.DeploymentFilter] = None,
216 work_pool_filter: Optional[schemas.filters.WorkPoolFilter] = None,
217 work_queue_filter: Optional[schemas.filters.WorkQueueFilter] = None,
218) -> Select[T]:
219 """
220 Applies filters to a flow run query as a combination of EXISTS subqueries.
221 """
223 if flow_run_filter:
224 query = query.where(flow_run_filter.as_sql_filter())
226 if deployment_filter:
227 deployment_exists_clause = select(db.Deployment).where(
228 db.Deployment.id == db.FlowRun.deployment_id,
229 deployment_filter.as_sql_filter(),
230 )
231 query = query.where(deployment_exists_clause.exists())
233 if work_pool_filter:
234 work_pool_exists_clause = select(db.WorkPool).where(
235 db.WorkQueue.id == db.FlowRun.work_queue_id,
236 db.WorkPool.id == db.WorkQueue.work_pool_id,
237 work_pool_filter.as_sql_filter(),
238 )
240 query = query.where(work_pool_exists_clause.exists())
242 if work_queue_filter:
243 work_queue_exists_clause = select(db.WorkQueue).where(
244 db.WorkQueue.id == db.FlowRun.work_queue_id,
245 work_queue_filter.as_sql_filter(),
246 )
247 query = query.where(work_queue_exists_clause.exists())
249 if flow_filter or task_run_filter:
250 flow_or_task_run_exists_clause: Union[
251 Select[Tuple[db.Flow]],
252 Select[Tuple[db.TaskRun]],
253 ]
255 if flow_filter:
256 flow_or_task_run_exists_clause = select(db.Flow).where(
257 db.Flow.id == db.FlowRun.flow_id,
258 flow_filter.as_sql_filter(),
259 )
261 if task_run_filter:
262 if not flow_filter:
263 flow_or_task_run_exists_clause = select(db.TaskRun).where(
264 db.TaskRun.flow_run_id == db.FlowRun.id
265 )
266 else:
267 flow_or_task_run_exists_clause = flow_or_task_run_exists_clause.join(
268 db.TaskRun,
269 db.TaskRun.flow_run_id == db.FlowRun.id,
270 )
271 flow_or_task_run_exists_clause = flow_or_task_run_exists_clause.where(
272 db.FlowRun.id == db.TaskRun.flow_run_id,
273 task_run_filter.as_sql_filter(),
274 )
276 query = query.where(flow_or_task_run_exists_clause.exists())
278 return query
281@db_injector
282async def read_flow_runs(
283 db: PrefectDBInterface,
284 session: AsyncSession,
285 columns: Optional[list[str]] = None,
286 flow_filter: Optional[schemas.filters.FlowFilter] = None,
287 flow_run_filter: Optional[schemas.filters.FlowRunFilter] = None,
288 task_run_filter: Optional[schemas.filters.TaskRunFilter] = None,
289 deployment_filter: Optional[schemas.filters.DeploymentFilter] = None,
290 work_pool_filter: Optional[schemas.filters.WorkPoolFilter] = None,
291 work_queue_filter: Optional[schemas.filters.WorkQueueFilter] = None,
292 offset: Optional[int] = None,
293 limit: Optional[int] = None,
294 sort: schemas.sorting.FlowRunSort = schemas.sorting.FlowRunSort.ID_DESC,
295) -> Sequence[orm_models.FlowRun]:
296 """
297 Read flow runs.
299 Args:
300 session: a database session
301 columns: a list of the flow run ORM columns to load, for performance
302 flow_filter: only select flow runs whose flows match these filters
303 flow_run_filter: only select flow runs match these filters
304 task_run_filter: only select flow runs whose task runs match these filters
305 deployment_filter: only select flow runs whose deployments match these filters
306 offset: Query offset
307 limit: Query limit
308 sort: Query sort
310 Returns:
311 List[orm_models.FlowRun]: flow runs
312 """
313 query = (
314 select(db.FlowRun)
315 .order_by(*sort.as_sql_sort())
316 .options(
317 selectinload(db.FlowRun.work_queue).selectinload(db.WorkQueue.work_pool)
318 )
319 )
321 if columns:
322 query = query.options(load_only(*columns))
324 query = await _apply_flow_run_filters(
325 db,
326 query,
327 flow_filter=flow_filter,
328 flow_run_filter=flow_run_filter,
329 task_run_filter=task_run_filter,
330 deployment_filter=deployment_filter,
331 work_pool_filter=work_pool_filter,
332 work_queue_filter=work_queue_filter,
333 )
335 if offset is not None:
336 query = query.offset(offset)
338 if limit is not None: 338 ↛ 341line 338 didn't jump to line 341 because the condition on line 338 was always true
339 query = query.limit(limit)
341 result = await session.execute(query)
342 return result.scalars().unique().all()
345async def cleanup_flow_run_concurrency_slots(
346 session: AsyncSession,
347 flow_run: orm_models.FlowRun,
348) -> None:
349 """
350 Cleanup flow run related resources, such as releasing concurrency slots.
351 All operations should be idempotent and safe to call multiple times.
352 IMPORTANT: This run may no longer exist in the database when this operation occurs.
353 """
355 if ( 355 ↛ 365line 355 didn't jump to line 365 because the condition on line 355 was never true
356 flow_run.deployment_id
357 and flow_run.state
358 and flow_run.state.type
359 in (
360 schemas.states.StateType.PENDING,
361 schemas.states.StateType.RUNNING,
362 schemas.states.StateType.CANCELLING,
363 )
364 ):
365 deployment = await models.deployments.read_deployment(
366 session, flow_run.deployment_id
367 )
368 if deployment and deployment.concurrency_limit_id:
369 await models.concurrency_limits_v2.bulk_decrement_active_slots(
370 session, [deployment.concurrency_limit_id], 1
371 )
374class DependencyResult(PrefectBaseModel):
375 id: UUID
376 name: str
377 upstream_dependencies: List[TaskRunResult]
378 state: Optional[State]
379 expected_start_time: Optional[datetime.datetime]
380 start_time: Optional[datetime.datetime]
381 end_time: Optional[datetime.datetime]
382 total_run_time: Optional[datetime.timedelta]
383 estimated_run_time: Optional[datetime.timedelta]
384 untrackable_result: bool
387async def read_task_run_dependencies(
388 session: AsyncSession,
389 flow_run_id: UUID,
390) -> List[DependencyResult]:
391 """
392 Get a task run dependency map for a given flow run.
393 """
394 flow_run = await models.flow_runs.read_flow_run(
395 session=session, flow_run_id=flow_run_id
396 )
397 if not flow_run:
398 raise ObjectNotFoundError(f"Flow run with id {flow_run_id} not found")
400 task_runs = await models.task_runs.read_task_runs(
401 session=session,
402 flow_run_filter=schemas.filters.FlowRunFilter(
403 id=schemas.filters.FlowRunFilterId(any_=[flow_run_id])
404 ),
405 )
407 dependency_graph = []
409 for task_run in task_runs:
410 inputs = list(set(chain(*task_run.task_inputs.values())))
411 untrackable_result_status = (
412 False
413 if task_run.state is None
414 else task_run.state.state_details.untrackable_result
415 )
416 dependency_graph.append(
417 DependencyResult(
418 id=task_run.id,
419 upstream_dependencies=inputs,
420 state=task_run.state,
421 expected_start_time=task_run.expected_start_time,
422 name=task_run.name,
423 start_time=task_run.start_time,
424 end_time=task_run.end_time,
425 total_run_time=task_run.total_run_time,
426 estimated_run_time=task_run.estimated_run_time,
427 untrackable_result=untrackable_result_status,
428 )
429 )
431 return dependency_graph
434@db_injector
435async def count_flow_runs(
436 db: PrefectDBInterface,
437 session: AsyncSession,
438 flow_filter: Optional[schemas.filters.FlowFilter] = None,
439 flow_run_filter: Optional[schemas.filters.FlowRunFilter] = None,
440 task_run_filter: Optional[schemas.filters.TaskRunFilter] = None,
441 deployment_filter: Optional[schemas.filters.DeploymentFilter] = None,
442 work_pool_filter: Optional[schemas.filters.WorkPoolFilter] = None,
443 work_queue_filter: Optional[schemas.filters.WorkQueueFilter] = None,
444) -> int:
445 """
446 Count flow runs.
448 Args:
449 session: a database session
450 flow_filter: only count flow runs whose flows match these filters
451 flow_run_filter: only count flow runs that match these filters
452 task_run_filter: only count flow runs whose task runs match these filters
453 deployment_filter: only count flow runs whose deployments match these filters
454 work_pool_filter: only count flow runs whose work pool matches these filters
455 work_queue_filter: only count flow runs whose work queue matches these filters
457 Returns:
458 int: count of flow runs
459 """
461 query = select(sa.func.count(None)).select_from(db.FlowRun)
463 # Fast path: JOIN instead of correlated EXISTS for the four FK filter dimensions.
464 # Safe because deployment_id, work_queue_id, and flow_id are single-valued scalar
465 # FK columns on FlowRun — each JOIN adds at most one row per flow_run, so
466 # COUNT(*) is not inflated. NULL FK rows are excluded by the INNER JOIN, which
467 # matches EXISTS semantics. task_run_filter is excluded: TaskRun.flow_run_id is
468 # the 1-to-N side, so a TaskRun JOIN would produce multiple rows per flow_run
469 # and corrupt the count. Any new filter added here must be N:1 to FlowRun.
470 if (
471 deployment_filter or work_pool_filter or work_queue_filter or flow_filter
472 ) and not task_run_filter:
473 if flow_run_filter:
474 query = query.where(flow_run_filter.as_sql_filter())
476 if deployment_filter:
477 query = query.join(
478 db.Deployment, db.Deployment.id == db.FlowRun.deployment_id
479 )
480 query = query.where(deployment_filter.as_sql_filter())
482 if work_queue_filter or work_pool_filter:
483 # WorkPool is not directly joined to FlowRun; go through WorkQueue
484 # as the bridge: FlowRun → WorkQueue → WorkPool.
485 query = query.join(
486 db.WorkQueue, db.WorkQueue.id == db.FlowRun.work_queue_id
487 )
488 if work_queue_filter:
489 query = query.where(work_queue_filter.as_sql_filter())
490 if work_pool_filter:
491 query = query.join(
492 db.WorkPool, db.WorkPool.id == db.WorkQueue.work_pool_id
493 )
494 query = query.where(work_pool_filter.as_sql_filter())
496 if flow_filter:
497 query = query.join(db.Flow, db.Flow.id == db.FlowRun.flow_id)
498 query = query.where(flow_filter.as_sql_filter())
500 else:
501 query = await _apply_flow_run_filters(
502 db,
503 query,
504 flow_filter=flow_filter,
505 flow_run_filter=flow_run_filter,
506 task_run_filter=task_run_filter,
507 deployment_filter=deployment_filter,
508 work_pool_filter=work_pool_filter,
509 work_queue_filter=work_queue_filter,
510 )
512 result = await session.execute(query)
513 return result.scalar_one()
516@db_injector
517async def delete_flow_run(
518 db: PrefectDBInterface, session: AsyncSession, flow_run_id: UUID
519) -> bool:
520 """
521 Delete a flow run by flow_run_id, handling concurrency limits if applicable.
523 Args:
524 session: A database session
525 flow_run_id: a flow run id
527 Returns:
528 bool: whether or not the flow run was deleted
529 """
530 flow_run = await read_flow_run(session, flow_run_id)
531 if not flow_run:
532 return False
534 deployment_id = flow_run.deployment_id
536 if deployment_id:
537 await cleanup_flow_run_concurrency_slots(session=session, flow_run=flow_run)
539 # Delete the flow run
540 result = await session.execute(
541 delete(db.FlowRun).where(db.FlowRun.id == flow_run_id)
542 )
544 return result.rowcount > 0
547@db_injector
548async def delete_flow_runs(
549 db: PrefectDBInterface,
550 session: AsyncSession,
551 flow_run_ids: List[UUID],
552) -> List[UUID]:
553 """
554 Delete multiple flow runs by their IDs, handling concurrency limits.
556 Args:
557 session: A database session
558 flow_run_ids: a list of flow run ids to delete
560 Returns:
561 List[UUID]: the IDs of the flow runs that were deleted
562 """
563 if not flow_run_ids:
564 return []
566 # Read all flow runs to handle concurrency cleanup
567 flow_runs = await session.execute(
568 select(db.FlowRun).where(db.FlowRun.id.in_(flow_run_ids))
569 )
570 flow_run_list = flow_runs.scalars().all()
572 if not flow_run_list:
573 return []
575 # Cleanup concurrency slots for each flow run that has a deployment
576 for flow_run in flow_run_list:
577 if flow_run.deployment_id:
578 await cleanup_flow_run_concurrency_slots(session=session, flow_run=flow_run)
580 # Get the IDs of flow runs that exist
581 existing_ids = [fr.id for fr in flow_run_list]
583 # Delete all flow runs in one query
584 await session.execute(delete(db.FlowRun).where(db.FlowRun.id.in_(existing_ids)))
586 return existing_ids
589async def set_flow_run_state(
590 session: AsyncSession,
591 flow_run_id: UUID,
592 state: schemas.states.State,
593 force: bool = False,
594 flow_policy: Optional[Type[FlowRunOrchestrationPolicy]] = None,
595 orchestration_parameters: Optional[Dict[str, Any]] = None,
596 client_version: Optional[str] = None,
597) -> OrchestrationResult:
598 """
599 Creates a new orchestrated flow run state.
601 Setting a new state on a run is the one of the principal actions that is governed by
602 Prefect's orchestration logic. Setting a new run state will not guarantee creation,
603 but instead trigger orchestration rules to govern the proposed `state` input. If
604 the state is considered valid, it will be written to the database. Otherwise, a
605 it's possible a different state, or no state, will be created. A `force` flag is
606 supplied to bypass a subset of orchestration logic.
608 Args:
609 session: a database session
610 flow_run_id: the flow run id
611 state: a flow run state model
612 force: if False, orchestration rules will be applied that may alter or prevent
613 the state transition. If True, orchestration rules are not applied.
615 Returns:
616 OrchestrationResult object
617 """
619 # load the flow run
620 run = await models.flow_runs.read_flow_run(
621 session=session,
622 flow_run_id=flow_run_id,
623 # Lock the row to prevent orchestration race conditions
624 for_update=True,
625 )
627 if not run:
628 raise ObjectNotFoundError(f"Flow run with id {flow_run_id} not found")
630 initial_state = run.state.as_state() if run.state else None
631 initial_state_type = initial_state.type if initial_state else None
632 proposed_state_type = state.type if state else None
633 intended_transition = (initial_state_type, proposed_state_type)
635 if force or flow_policy is None:
636 flow_policy = MinimalFlowPolicy
638 orchestration_rules = flow_policy.compile_transition_rules(*intended_transition) # type: ignore
639 global_rules = GlobalFlowPolicy.compile_transition_rules(*intended_transition)
641 context = FlowOrchestrationContext(
642 session=session,
643 run=run,
644 initial_state=initial_state,
645 proposed_state=state,
646 client_version=client_version,
647 )
649 if orchestration_parameters is not None:
650 context.parameters = orchestration_parameters
652 # apply orchestration rules and create the new flow run state
653 async with contextlib.AsyncExitStack() as stack:
654 for rule in orchestration_rules: 654 ↛ 659line 654 didn't jump to line 659 because the loop on line 654 didn't complete
655 context = await stack.enter_async_context(
656 rule(context, *intended_transition)
657 )
659 for rule in global_rules:
660 context = await stack.enter_async_context(
661 rule(context, *intended_transition)
662 )
664 await context.validate_proposed_state()
666 if context.orchestration_error is not None: 666 ↛ 669line 666 didn't jump to line 669 because the condition on line 666 was always true
667 raise context.orchestration_error
669 result = OrchestrationResult(
670 state=context.validated_state,
671 status=context.response_status,
672 details=context.response_details,
673 )
675 return result
678@db_injector
679async def read_flow_run_graph(
680 db: PrefectDBInterface,
681 session: AsyncSession,
682 flow_run_id: UUID,
683 since: datetime.datetime = earliest_possible_datetime(),
684) -> Graph:
685 """Given a flow run, return the graph of it's task and subflow runs. If a `since`
686 datetime is provided, only return items that may have changed since that time."""
687 if isinstance(since, str): 687 ↛ 688line 687 didn't jump to line 688 because the condition on line 687 was never true
688 since = DateTime.fromisoformat(since)
690 return await db.queries.flow_run_graph_v2(
691 session=session,
692 flow_run_id=flow_run_id,
693 since=since,
694 max_nodes=PREFECT_API_MAX_FLOW_RUN_GRAPH_NODES.value(),
695 max_artifacts=PREFECT_API_MAX_FLOW_RUN_GRAPH_ARTIFACTS.value(),
696 )
699async def with_system_labels_for_flow_run(
700 session: AsyncSession,
701 flow_run: Union[schemas.core.FlowRun, schemas.actions.FlowRunCreate],
702) -> schemas.core.KeyValueLabels:
703 """Augment user supplied labels with system default labels for a flow
704 run."""
706 user_supplied_labels = flow_run.labels or {}
708 # `deployment_id` is deprecated on `schemas.actions.FlowRunCreate`. Only
709 # check `deployment_id` if given an instance of a `schemas.core.FlowRun`.
710 if isinstance(flow_run, schemas.core.FlowRun) and flow_run.deployment_id:
711 deployment = await models.deployments.read_deployment(
712 session, deployment_id=flow_run.deployment_id
713 )
714 if deployment:
715 # Use the deployment flow run utility for consistent label generation
716 return await models.deployments.with_system_labels_for_deployment_flow_run(
717 session=session,
718 deployment=deployment,
719 user_supplied_labels=user_supplied_labels,
720 )
722 # If the flow run is not part of a deployment, generate basic flow labels
723 default_labels = cast(
724 schemas.core.KeyValueLabels,
725 {
726 "prefect.flow.id": str(flow_run.flow_id),
727 },
728 )
730 parent_labels = await models.flows.read_flow_labels(session, flow_run.flow_id) or {}
732 return parent_labels | default_labels | user_supplied_labels
735@db_injector
736async def update_flow_run_labels(
737 db: PrefectDBInterface,
738 session: AsyncSession,
739 flow_run_id: UUID,
740 labels: KeyValueLabels,
741) -> bool:
742 """
743 Update flow run labels by patching existing labels with new values.
744 Args:
745 session: A database session
746 flow_run_id: the flow run id to update
747 labels: the new labels to patch into existing labels
748 Returns:
749 bool: whether the update was successful
750 """
751 # First read the existing flow run to get current labels
752 flow_run: Optional[orm_models.FlowRun] = await read_flow_run(session, flow_run_id)
753 if not flow_run:
754 raise ObjectNotFoundError(f"Flow run with id {flow_run_id} not found")
756 # Merge existing labels with new labels
757 current_labels = flow_run.labels or {}
758 updated_labels = {**current_labels, **labels}
760 try:
761 # Update the flow run with merged labels
762 result = await session.execute(
763 sa.update(db.FlowRun)
764 .where(db.FlowRun.id == flow_run_id)
765 .values(labels=updated_labels)
766 )
767 success = result.rowcount > 0
768 if success:
769 await session.commit() # Explicitly commit
770 return success
771 except Exception:
772 raise