diff --git a/airflow-core/src/airflow/executors/workloads/__init__.py b/airflow-core/src/airflow/executors/workloads/__init__.py index dca4c991f637b..462e38ad0aaac 100644 --- a/airflow-core/src/airflow/executors/workloads/__init__.py +++ b/airflow-core/src/airflow/executors/workloads/__init__.py @@ -24,7 +24,7 @@ from airflow.executors.workloads.base import BaseWorkload, BundleInfo from airflow.executors.workloads.callback import CallbackFetchMethod, ExecuteCallback -from airflow.executors.workloads.task import ExecuteTask +from airflow.executors.workloads.task import ExecuteTask, TaskInstanceDTO from airflow.executors.workloads.trigger import RunTrigger All = Annotated[ @@ -32,4 +32,15 @@ Field(discriminator="type"), ] -__all__ = ["All", "BaseWorkload", "BundleInfo", "CallbackFetchMethod", "ExecuteCallback", "ExecuteTask"] +TaskInstance = TaskInstanceDTO + +__all__ = [ + "All", + "BaseWorkload", + "BundleInfo", + "CallbackFetchMethod", + "ExecuteCallback", + "ExecuteTask", + "TaskInstance", + "TaskInstanceDTO", +] diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py new file mode 100644 index 0000000000000..8ca86f9704b4b --- /dev/null +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -0,0 +1,27 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from airflow.executors import workloads +from airflow.executors.workloads import TaskInstance, TaskInstanceDTO + + +def test_task_instance_alias_keeps_backwards_compat(): + assert TaskInstance is TaskInstanceDTO + assert workloads.TaskInstance is TaskInstanceDTO + assert workloads.TaskInstanceDTO is TaskInstanceDTO