fix logic

Author: wjddn279

providers/edge3/src/airflow/providers/edge3/cli/worker.py

@@ -51,12 +51,12 @@
     status_file_path,
     write_pid_to_pidfile,
 )
-from airflow.providers.edge3.utils import is_callback_execute
 from airflow.providers.edge3.models.edge_worker import (
     EdgeWorkerDuplicateException,
     EdgeWorkerState,
     EdgeWorkerVersionException,
 )
+from airflow.providers.edge3.utils.types import is_callback_execute
 from airflow.utils.net import getfqdn
 from airflow.utils.state import TaskInstanceState
 

providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py

@@ -28,11 +28,11 @@
 from airflow.executors.base_executor import BaseExecutor
 from airflow.models.taskinstance import TaskInstance
 from airflow.providers.common.compat.sdk import Stats, conf, timezone
-from airflow.providers.edge3.utils import is_callback_execute
 from airflow.providers.edge3.models.db import EdgeDBManager, check_db_manager_config
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
 from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, reset_metrics
+from airflow.providers.edge3.utils.types import is_callback_execute
 from airflow.utils.db import DBLocks, create_global_lock
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.state import TaskInstanceState

providers/edge3/src/airflow/providers/edge3/utils/__init__.py

@@ -0,0 +1,16 @@
+# 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.

…ge3/src/airflow/providers/edge3/utils.py …c/airflow/providers/edge3/utils/types.pyproviders/edge3/src/airflow/providers/edge3/utils.py renamed to providers/edge3/src/airflow/providers/edge3/utils/types.py providers/edge3/src/airflow/providers/edge3/utils.py renamed to providers/edge3/src/airflow/providers/edge3/utils/types.py

@@ -16,22 +16,24 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Annotated, TypeAlias, TypeGuard, Union
+from typing import TYPE_CHECKING, Annotated, TypeAlias, TypeGuard
 
 from pydantic import Discriminator, Tag
 
-from airflow.executors import workloads
 from airflow.providers.edge3.version_compat import AIRFLOW_V_3_2_PLUS
 
+if TYPE_CHECKING:
+    from airflow.executors import workloads
+
 if not AIRFLOW_V_3_2_PLUS:
     from airflow.executors.workloads import ExecuteTask
 
     ExecuteTypeBody: TypeAlias = ExecuteTask
 else:
     from airflow.executors.workloads import ExecuteCallback, ExecuteTask
 
-    ExecuteTypeBody: TypeAlias = Annotated[
-        Union[Annotated[ExecuteTask, Tag("ExecuteTask")], Annotated[ExecuteCallback, Tag("ExecuteCallback")]],
+    ExecuteTypeBody: TypeAlias = Annotated[  # type: ignore[no-redef,misc]
+        Annotated[ExecuteTask, Tag("ExecuteTask")] | Annotated[ExecuteCallback, Tag("ExecuteCallback")],
         Discriminator("type"),
     ]
 

providers/edge3/src/airflow/providers/edge3/worker_api/datamodels.py

@@ -27,7 +27,7 @@
 
 from airflow.providers.common.compat.sdk import TaskInstanceKey
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerState  # noqa: TCH001
-from airflow.providers.edge3.utils import ExecuteTypeBody  # noqa: TCH001
+from airflow.providers.edge3.utils.types import ExecuteTypeBody  # noqa: TCH001
 
 
 class WorkerApiDocs:

providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import Annotated
+from typing import TYPE_CHECKING, Annotated
 
 from fastapi import Body, Depends, status
 from sqlalchemy import select, update
@@ -27,7 +27,6 @@
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.executors.workloads import ExecuteTask
 from airflow.providers.common.compat.sdk import Stats, timezone
-from airflow.providers.edge3.utils import ExecuteTypeBody
 from airflow.providers.edge3.version_compat import AIRFLOW_V_3_2_PLUS
 
 try:
@@ -43,6 +42,9 @@
 )
 from airflow.utils.state import TaskInstanceState
 
+if TYPE_CHECKING:
+    from airflow.providers.edge3.utils.types import ExecuteTypeBody
+
 jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs")
 
 
@@ -51,7 +53,7 @@ def parse_command(command: str, dag_id: str, run_id: str) -> ExecuteTypeBody:
         from airflow.executors.workloads import ExecuteCallback
 
         if dag_id == ExecuteCallback.TYPE and run_id.startswith(ExecuteCallback.TYPE):
-            return ExecuteCallback.model_validate_json(command)
+            return ExecuteCallback.model_validate_json(command)  # type: ignore[return-value]
 
     return ExecuteTask.model_validate_json(command)
 

providers/edge3/tests/unit/edge3/executors/test_edge_executor.py

@@ -26,7 +26,7 @@
 from sqlalchemy import delete, select
 
 from airflow.configuration import conf
-from airflow.executors.workloads import ExecuteTask, TaskInstanceDTO
+from airflow.executors.workloads import ExecuteTask
 from airflow.executors.workloads.base import BundleInfo
 from airflow.models.taskinstancekey import TaskInstanceKey
 from airflow.providers.common.compat.sdk import Stats, timezone
@@ -40,6 +40,7 @@
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
 
 if AIRFLOW_V_3_2_PLUS:
+    from airflow.executors.workloads import TaskInstanceDTO
     from airflow.executors.workloads.callback import CallbackDTO, CallbackFetchMethod, ExecuteCallback
 
 pytestmark = pytest.mark.db_test

providers/edge3/tests/unit/edge3/utils/__init__.py

@@ -0,0 +1,16 @@
+# 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.

providers/edge3/tests/unit/edge3/utils/test_types.py

@@ -0,0 +1,141 @@
+# 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 pathlib import Path
+from uuid import uuid4
+
+import pytest
+from pydantic import TypeAdapter
+
+from airflow.executors.workloads import ExecuteTask
+from airflow.executors.workloads.base import BundleInfo
+from airflow.providers.edge3.utils.types import ExecuteTypeBody, is_callback_execute
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
+
+if AIRFLOW_V_3_2_PLUS:
+    from airflow.executors.workloads import TaskInstanceDTO
+    from airflow.executors.workloads.callback import CallbackDTO, CallbackFetchMethod, ExecuteCallback
+
+
+def _make_execute_task() -> ExecuteTask:
+    ti = TaskInstanceDTO(
+        id=uuid4(),
+        dag_version_id=uuid4(),
+        task_id="test_task",
+        dag_id="test_dag",
+        run_id="test_run",
+        try_number=1,
+        map_index=-1,
+        pool_slots=1,
+        queue="default",
+        priority_weight=1,
+    )
+    return ExecuteTask(
+        ti=ti,
+        dag_rel_path=Path("test_dag.py"),
+        token="test_token",
+        bundle_info=BundleInfo(name="test_bundle", version="1.0"),
+        log_path="test.log",
+    )
+
+
+def _make_execute_callback() -> ExecuteCallback:
+    callback_data = CallbackDTO(
+        id=str(uuid4()),
+        fetch_method=CallbackFetchMethod.IMPORT_PATH,
+        data={
+            "path": "builtins.dict",
+            "kwargs": {"a": 1, "b": 2, "c": 3},
+        },
+    )
+    return ExecuteCallback(
+        callback=callback_data,
+        dag_rel_path=Path("test.py"),
+        bundle_info=BundleInfo(name="test_bundle", version="1.0"),
+        token="test_token",
+        log_path="test.log",
+    )
+
+
+class TestIsCallbackExecute:
+    def test_returns_false_for_execute_task(self):
+        workload = _make_execute_task()
+        assert is_callback_execute(workload) is False
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="ExecuteCallback requires Airflow 3.2+")
+    def test_returns_true_for_execute_callback(self):
+        workload = _make_execute_callback()
+        assert is_callback_execute(workload) is True
+
+
+@pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="ExecuteTypeBody union requires Airflow 3.2+")
+class TestExecuteTypeBody:
+    def setup_method(self):
+        self.adapter: TypeAdapter = TypeAdapter(ExecuteTypeBody)
+
+    def test_validate_execute_task_json(self):
+        workload = _make_execute_task()
+        json_str = workload.model_dump_json()
+
+        result = self.adapter.validate_json(json_str)
+
+        assert isinstance(result, ExecuteTask)
+        assert result.ti.dag_id == "test_dag"
+
+    def test_validate_execute_callback_json(self):
+        workload = _make_execute_callback()
+        json_str = workload.model_dump_json()
+
+        result = self.adapter.validate_json(json_str)
+
+        assert isinstance(result, ExecuteCallback)
+        assert result.callback.fetch_method == CallbackFetchMethod.IMPORT_PATH
+
+    def test_validate_execute_task_dict(self):
+        workload = _make_execute_task()
+        data = workload.model_dump()
+
+        result = self.adapter.validate_python(data)
+
+        assert isinstance(result, ExecuteTask)
+
+    def test_validate_execute_callback_dict(self):
+        workload = _make_execute_callback()
+        data = workload.model_dump()
+
+        result = self.adapter.validate_python(data)
+
+        assert isinstance(result, ExecuteCallback)
+
+    def test_roundtrip_execute_task(self):
+        original = _make_execute_task()
+        json_str = self.adapter.dump_json(original)
+        restored = self.adapter.validate_json(json_str)
+
+        assert isinstance(restored, ExecuteTask)
+        assert restored.ti.task_id == original.ti.task_id
+        assert restored.ti.dag_id == original.ti.dag_id
+
+    def test_roundtrip_execute_callback(self):
+        original = _make_execute_callback()
+        json_str = self.adapter.dump_json(original)
+        restored = self.adapter.validate_json(json_str)
+
+        assert isinstance(restored, ExecuteCallback)
+        assert restored.callback.id == original.callback.id

providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py

@@ -24,7 +24,7 @@
 import pytest
 from sqlalchemy import delete, select
 
-from airflow.executors.workloads import ExecuteTask, TaskInstanceDTO
+from airflow.executors.workloads import ExecuteTask
 from airflow.executors.workloads.base import BundleInfo
 from airflow.executors.workloads.callback import CallbackDTO
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
@@ -37,10 +37,11 @@
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
-    from airflow.executors.workloads import CallbackFetchMethod, ExecuteCallback
+    from airflow.executors.workloads import CallbackFetchMethod, ExecuteCallback, TaskInstanceDTO
 
 if AIRFLOW_V_3_2_PLUS:
-    from airflow.executors.workloads import CallbackFetchMethod, ExecuteCallback
+    from airflow.executors.workloads import CallbackFetchMethod, ExecuteCallback, TaskInstanceDTO
+
 
 try:
     from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager  # noqa: F401