diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py index c110e0ad07f1e..7ed84b180d771 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py @@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Annotated -from fastapi import Body, Depends, status +from fastapi import Body, Depends, HTTPException, status from sqlalchemy import select, update from airflow.api_fastapi.common.db.common import SessionDep # noqa: TC001 @@ -28,6 +28,7 @@ from airflow.executors.workloads import ExecuteTask from airflow.providers.common.compat.sdk import Stats, timezone from airflow.providers.edge3.models.edge_job import EdgeJobModel +from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel from airflow.providers.edge3.version_compat import AIRFLOW_V_3_3_PLUS from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest from airflow.providers.edge3.worker_api.datamodels import ( @@ -61,6 +62,7 @@ def parse_command(command: str, dag_id: str, run_id: str) -> ExecuteTypeBody: [ status.HTTP_400_BAD_REQUEST, status.HTTP_403_FORBIDDEN, + status.HTTP_404_NOT_FOUND, ] ), ) @@ -76,6 +78,10 @@ def fetch( session: SessionDep, ) -> EdgeJobFetched | None: """Fetch a job to execute on the edge worker.""" + worker = session.scalar(select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)) + if not worker: + raise HTTPException(status.HTTP_404_NOT_FOUND, "Worker not found") + query = ( select(EdgeJobModel) .where( @@ -86,7 +92,8 @@ def fetch( ) if body.queues: query = query.where(EdgeJobModel.queue.in_(body.queues)) - query = query.where(EdgeJobModel.team_name == body.team_name) + if worker.team_name is not None: + query = query.where(EdgeJobModel.team_name == worker.team_name) query = query.limit(1) query = query.with_for_update(skip_locked=True) job: EdgeJobModel | None = session.scalar(query) diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml index 8c0c6a3eacdd1..408cb14078381 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml @@ -65,6 +65,12 @@ paths: schema: $ref: '#/components/schemas/HTTPExceptionResponse' description: Forbidden + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found '422': description: Validation Error content: diff --git a/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py b/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py index c120e86f08861..26a06ba470f4e 100644 --- a/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py +++ b/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py @@ -23,11 +23,13 @@ from uuid import uuid4 import pytest +from fastapi import HTTPException, status from sqlalchemy import delete, select from airflow.executors.workloads import BundleInfo, ExecuteTask from airflow.providers.common.compat.sdk import Stats from airflow.providers.edge3.models.edge_job import EdgeJobModel +from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState from airflow.providers.edge3.models.types import EXECUTE_CALLBACK_TAG from airflow.providers.edge3.worker_api.datamodels import WorkerQueuesBody from airflow.providers.edge3.worker_api.routes.jobs import fetch, parse_command, state @@ -80,6 +82,7 @@ class TestJobsApiRoutes: @pytest.fixture(autouse=True) def setup_test_cases(self, dag_maker, session: Session): session.execute(delete(EdgeJobModel)) + session.execute(delete(EdgeWorkerModel)) session.commit() @patch(f"{Stats.__module__}.Stats.incr") @@ -135,8 +138,13 @@ def test_state(self, mock_stats_incr, session: Session): assert db_job is not None assert db_job.state == TaskInstanceState.SUCCESS - def test_fetch_filters_by_team_name(self, session: Session): + def test_fetch_filters_by_worker_team_name(self, session: Session): with create_session() as session: + session.add( + EdgeWorkerModel( + worker_name="worker1", state=EdgeWorkerState.IDLE, queues=[QUEUE], team_name="team_a" + ) + ) job_team_a = EdgeJobModel( dag_id="dag_a", task_id="task_a", @@ -164,15 +172,29 @@ def test_fetch_filters_by_team_name(self, session: Session): session.add_all([job_team_a, job_team_b]) session.commit() - body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE], team_name="team_a") + body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE], team_name="team_b") result = fetch("worker1", body, session) assert result is not None assert result.dag_id == "dag_a" assert result.task_id == "task_a" + def test_fetch_unknown_worker_raises_404(self, session: Session): + body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE], team_name="team_a") + + with pytest.raises(HTTPException) as exc_info: + fetch("unknown-worker", body, session) + + assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND + assert exc_info.value.detail == "Worker not found" + def test_fetch_without_team_name_returns_any_team(self, session: Session): - """When team_name is None, no team filter is applied so any queued job can be returned.""" + """When a worker has no team_name, no team filter is applied so any queued job can be returned.""" with create_session() as session: + session.add( + EdgeWorkerModel( + worker_name="worker1", state=EdgeWorkerState.IDLE, queues=[QUEUE], team_name=None + ) + ) job_team_a = EdgeJobModel( dag_id="dag_a", task_id="task_a",