Object of type `DagRunProtocol` has no attribute `partition_key`

[m1racoli]

Under which category would you file this issue?

Airflow Core

Apache Airflow version

3.2.0

What happened and how to reproduce it?

When type checking the following code:

from airflow.sdk import CronPartitionTimetable, dag, get_current_context, task


@task
def example():
    ctx = get_current_context()
    print(type(ctx["dag_run"]))
    print(ctx["dag_run"].partition_key)


@dag(
    schedule=CronPartitionTimetable("0 23 * * *", timezone="UTC"),
)
def asset_partitioning():

    example()


asset_partitioning()

we get the following error for ctx["dag_run"].partition_key:

Object of type `DagRunProtocol` has no attribute `partition_key`

but when running the task everything works fine:

[2026-04-15 19:21:43] INFO - DAG bundles loaded: dags-folder
[2026-04-15 19:21:43] INFO - Filling up the DagBag from /usr/local/airflow/dags/asset_partitioning.py
[2026-04-15 19:21:43] INFO - TaskInstance Details:  dag_id=asset_partitioning task_id=example dagrun_id=scheduled__2026-04-15T23:00:00+00:00__2026-04-15T23:00:00__BC4C5IYH map_index=-1 run_start_date=2026-04-16T01:21:43.052300Z try_number=2 op_classpath=["airflow.providers.standard.decorators.python._PythonDecoratedOperator","airflow.sdk.bases.decorator.DecoratedOperator","airflow.providers.standard.operators.python.PythonOperator","airflow.sdk.bases.operator.BaseAsyncOperator"]
[2026-04-15 19:21:43] INFO - Done. Returned value was: None
[2026-04-15 19:21:43] INFO - <class 'airflow.sdk.api.datamodels._generated.DagRun'>
[2026-04-15 19:21:43] INFO - 2026-04-15T23:00:0

The question wether Context.dag_run is correctly typed with airflow.sdk.types.DagRunProtocol instead of airflow.sdk.api.datamodels._generated.DagRun (which is the actual type during exectuion) or if DagRunProtocol is missing partition_key in its definition.

What you think should happen instead?

No response

Operating System

No response

Deployment

None

Apache Airflow Provider(s)

No response

Versions of Apache Airflow Providers

No response

Official Helm Chart version

Not Applicable

Kubernetes Version

No response

Helm Chart configuration

No response

Docker Image customizations

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct