Skip to content

prefect.server.models.task_runs

Functions for interacting with task run ORM objects. Intended for internal use by the Prefect REST API.

count_task_runs async

Count task runs.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
flow_filter schemas.filters.FlowFilter

only count task runs whose flows match these filters

None
flow_run_filter schemas.filters.FlowRunFilter

only count task runs whose flow runs match these filters

None
task_run_filter schemas.filters.TaskRunFilter

only count task runs that match these filters

None
deployment_filter schemas.filters.DeploymentFilter

only count task runs whose deployments match these filters

None

Returns:

Name Type Description
int int

count of task runs

Source code in prefect/server/models/task_runs.py
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
@inject_db
async def count_task_runs(
    session: sa.orm.Session,
    db: PrefectDBInterface,
    flow_filter: schemas.filters.FlowFilter = None,
    flow_run_filter: schemas.filters.FlowRunFilter = None,
    task_run_filter: schemas.filters.TaskRunFilter = None,
    deployment_filter: schemas.filters.DeploymentFilter = None,
) -> int:
    """
    Count task runs.

    Args:
        session: a database session
        flow_filter: only count task runs whose flows match these filters
        flow_run_filter: only count task runs whose flow runs match these filters
        task_run_filter: only count task runs that match these filters
        deployment_filter: only count task runs whose deployments match these filters
    Returns:
        int: count of task runs
    """

    query = select(sa.func.count(sa.text("*"))).select_from(db.TaskRun)

    query = await _apply_task_run_filters(
        query,
        flow_filter=flow_filter,
        flow_run_filter=flow_run_filter,
        task_run_filter=task_run_filter,
        deployment_filter=deployment_filter,
        db=db,
    )

    result = await session.execute(query)
    return result.scalar()

create_task_run async

Creates a new task run.

If a task run with the same flow_run_id, task_key, and dynamic_key already exists, the existing task run will be returned. If the provided task run has a state attached, it will also be created.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run schemas.core.TaskRun

a task run model

required

Returns:

Type Description

db.TaskRun: the newly-created or existing task run

Source code in prefect/server/models/task_runs.py
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@inject_db
async def create_task_run(
    session: sa.orm.Session,
    task_run: schemas.core.TaskRun,
    db: PrefectDBInterface,
    orchestration_parameters: dict = None,
):
    """
    Creates a new task run.

    If a task run with the same flow_run_id, task_key, and dynamic_key already exists,
    the existing task run will be returned. If the provided task run has a state
    attached, it will also be created.

    Args:
        session: a database session
        task_run: a task run model

    Returns:
        db.TaskRun: the newly-created or existing task run
    """

    now = pendulum.now("UTC")

    # if a dynamic key exists, we need to guard against conflicts
    insert_stmt = (
        (await db.insert(db.TaskRun))
        .values(
            created=now,
            **task_run.dict(
                shallow=True, exclude={"state", "created"}, exclude_unset=True
            ),
        )
        .on_conflict_do_nothing(
            index_elements=db.task_run_unique_upsert_columns,
        )
    )
    await session.execute(insert_stmt)

    query = (
        sa.select(db.TaskRun)
        .where(
            sa.and_(
                db.TaskRun.flow_run_id == task_run.flow_run_id,
                db.TaskRun.task_key == task_run.task_key,
                db.TaskRun.dynamic_key == task_run.dynamic_key,
            )
        )
        .limit(1)
        .execution_options(populate_existing=True)
    )
    result = await session.execute(query)
    model = result.scalar()

    if model.created == now and task_run.state:
        await models.task_runs.set_task_run_state(
            session=session,
            task_run_id=model.id,
            state=task_run.state,
            force=True,
            orchestration_parameters=orchestration_parameters,
        )
    return model

delete_task_run async

Delete a task run by id.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run_id UUID

the task run id to delete

required

Returns:

Name Type Description
bool bool

whether or not the task run was deleted

Source code in prefect/server/models/task_runs.py
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
@inject_db
async def delete_task_run(
    session: sa.orm.Session, task_run_id: UUID, db: PrefectDBInterface
) -> bool:
    """
    Delete a task run by id.

    Args:
        session: a database session
        task_run_id: the task run id to delete

    Returns:
        bool: whether or not the task run was deleted
    """

    result = await session.execute(
        delete(db.TaskRun).where(db.TaskRun.id == task_run_id)
    )
    return result.rowcount > 0

read_task_run async

Read a task run by id.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run_id UUID

the task run id

required

Returns:

Type Description

db.TaskRun: the task run

Source code in prefect/server/models/task_runs.py
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
@inject_db
async def read_task_run(
    session: sa.orm.Session, task_run_id: UUID, db: PrefectDBInterface
):
    """
    Read a task run by id.

    Args:
        session: a database session
        task_run_id: the task run id

    Returns:
        db.TaskRun: the task run
    """

    model = await session.get(db.TaskRun, task_run_id)
    return model

read_task_runs async

Read task runs.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
flow_filter schemas.filters.FlowFilter

only select task runs whose flows match these filters

None
flow_run_filter schemas.filters.FlowRunFilter

only select task runs whose flow runs match these filters

None
task_run_filter schemas.filters.TaskRunFilter

only select task runs that match these filters

None
deployment_filter schemas.filters.DeploymentFilter

only select task runs whose deployments match these filters

None
offset int

Query offset

None
limit int

Query limit

None
sort schemas.sorting.TaskRunSort

Query sort

schemas.sorting.TaskRunSort.ID_DESC

Returns:

Type Description

List[db.TaskRun]: the task runs

Source code in prefect/server/models/task_runs.py
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
@inject_db
async def read_task_runs(
    session: sa.orm.Session,
    db: PrefectDBInterface,
    flow_filter: schemas.filters.FlowFilter = None,
    flow_run_filter: schemas.filters.FlowRunFilter = None,
    task_run_filter: schemas.filters.TaskRunFilter = None,
    deployment_filter: schemas.filters.DeploymentFilter = None,
    offset: int = None,
    limit: int = None,
    sort: schemas.sorting.TaskRunSort = schemas.sorting.TaskRunSort.ID_DESC,
):
    """
    Read task runs.

    Args:
        session: a database session
        flow_filter: only select task runs whose flows match these filters
        flow_run_filter: only select task runs whose flow runs match these filters
        task_run_filter: only select task runs that match these filters
        deployment_filter: only select task runs whose deployments match these filters
        offset: Query offset
        limit: Query limit
        sort: Query sort

    Returns:
        List[db.TaskRun]: the task runs
    """

    query = select(db.TaskRun).order_by(sort.as_sql_sort(db))

    query = await _apply_task_run_filters(
        query,
        flow_filter=flow_filter,
        flow_run_filter=flow_run_filter,
        task_run_filter=task_run_filter,
        deployment_filter=deployment_filter,
        db=db,
    )

    if offset is not None:
        query = query.offset(offset)

    if limit is not None:
        query = query.limit(limit)

    result = await session.execute(query)
    return result.scalars().unique().all()

set_task_run_state async

Creates a new orchestrated task run state.

Setting a new state on a run is the one of the principal actions that is governed by Prefect's orchestration logic. Setting a new run state will not guarantee creation, but instead trigger orchestration rules to govern the proposed state input. If the state is considered valid, it will be written to the database. Otherwise, a it's possible a different state, or no state, will be created. A force flag is supplied to bypass a subset of orchestration logic.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run_id UUID

the task run id

required
state schemas.states.State

a task run state model

required
force bool

if False, orchestration rules will be applied that may alter or prevent the state transition. If True, orchestration rules are not applied.

False

Returns:

Type Description
OrchestrationResult

OrchestrationResult object

Source code in prefect/server/models/task_runs.py
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
async def set_task_run_state(
    session: sa.orm.Session,
    task_run_id: UUID,
    state: schemas.states.State,
    force: bool = False,
    task_policy: BaseOrchestrationPolicy = None,
    orchestration_parameters: dict = None,
) -> OrchestrationResult:
    """
    Creates a new orchestrated task run state.

    Setting a new state on a run is the one of the principal actions that is governed by
    Prefect's orchestration logic. Setting a new run state will not guarantee creation,
    but instead trigger orchestration rules to govern the proposed `state` input. If
    the state is considered valid, it will be written to the database. Otherwise, a
    it's possible a different state, or no state, will be created. A `force` flag is
    supplied to bypass a subset of orchestration logic.

    Args:
        session: a database session
        task_run_id: the task run id
        state: a task run state model
        force: if False, orchestration rules will be applied that may alter or prevent
            the state transition. If True, orchestration rules are not applied.

    Returns:
        OrchestrationResult object
    """

    # load the task run
    run = await models.task_runs.read_task_run(session=session, task_run_id=task_run_id)

    if not run:
        raise ObjectNotFoundError(f"Task run with id {task_run_id} not found")

    initial_state = run.state.as_state() if run.state else None
    initial_state_type = initial_state.type if initial_state else None
    proposed_state_type = state.type if state else None
    intended_transition = (initial_state_type, proposed_state_type)

    if force or task_policy is None:
        task_policy = MinimalTaskPolicy

    orchestration_rules = task_policy.compile_transition_rules(*intended_transition)
    global_rules = GlobalTaskPolicy.compile_transition_rules(*intended_transition)

    context = TaskOrchestrationContext(
        session=session,
        run=run,
        initial_state=initial_state,
        proposed_state=state,
    )

    if orchestration_parameters is not None:
        context.parameters = orchestration_parameters

    # apply orchestration rules and create the new task run state
    async with contextlib.AsyncExitStack() as stack:
        for rule in orchestration_rules:
            context = await stack.enter_async_context(
                rule(context, *intended_transition)
            )

        for rule in global_rules:
            context = await stack.enter_async_context(
                rule(context, *intended_transition)
            )

        await context.validate_proposed_state()

    if context.orchestration_error is not None:
        raise context.orchestration_error

    result = OrchestrationResult(
        state=context.validated_state,
        status=context.response_status,
        details=context.response_details,
    )

    return result

update_task_run async

Updates a task run.

Parameters:

Name Type Description Default
session AsyncSession

a database session

required
task_run_id UUID

the task run id to update

required
task_run schemas.actions.TaskRunUpdate

a task run model

required

Returns:

Name Type Description
bool bool

whether or not matching rows were found to update

Source code in prefect/server/models/task_runs.py
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@inject_db
async def update_task_run(
    session: AsyncSession,
    task_run_id: UUID,
    task_run: schemas.actions.TaskRunUpdate,
    db: PrefectDBInterface,
) -> bool:
    """
    Updates a task run.

    Args:
        session: a database session
        task_run_id: the task run id to update
        task_run: a task run model

    Returns:
        bool: whether or not matching rows were found to update
    """
    update_stmt = (
        sa.update(db.TaskRun).where(db.TaskRun.id == task_run_id)
        # exclude_unset=True allows us to only update values provided by
        # the user, ignoring any defaults on the model
        .values(**task_run.dict(shallow=True, exclude_unset=True))
    )
    result = await session.execute(update_stmt)
    return result.rowcount > 0