apache-airflow-providers-edge3 1.0.0rc1__py3-none-any.whl

airflow/providers/edge3/worker_api/routes/_v2_routes.py

@@ -0,0 +1,237 @@
+# 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.
+"""Compatibility layer for Connexion API to Airflow v2.10 API routes."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from flask import Response, request
+
+from airflow.exceptions import AirflowException
+from airflow.providers.edge3.worker_api.auth import (
+    jwt_token_authorization,
+    jwt_token_authorization_rpc,
+)
+from airflow.providers.edge3.worker_api.datamodels import (
+    EdgeJobFetched,
+    JsonRpcRequest,
+    PushLogsBody,
+    WorkerQueuesBody,
+    WorkerStateBody,
+)
+from airflow.providers.edge3.worker_api.routes._v2_compat import HTTPException, status
+from airflow.providers.edge3.worker_api.routes.jobs import fetch, state as state_api
+from airflow.providers.edge3.worker_api.routes.logs import logfile_path, push_logs
+from airflow.providers.edge3.worker_api.routes.worker import register, set_state
+from airflow.serialization.serialized_objects import BaseSerialization
+from airflow.utils.session import NEW_SESSION, create_session, provide_session
+
+if TYPE_CHECKING:
+    from airflow.api_connexion.types import APIResponse
+    from airflow.utils.state import TaskInstanceState
+
+
+log = logging.getLogger(__name__)
+
+
+def error_response(message: str, status: int):
+    """Log the error and return the response as JSON object."""
+    error_id = uuid4()
+    server_message = f"{message} error_id={error_id}"
+    log.exception(server_message)
+    client_message = f"{message} The server side traceback may be identified with error_id={error_id}"
+    return HTTPException(status, client_message)
+
+
+def rpcapi_v2(body: dict[str, Any]) -> APIResponse:
+    """Handle Edge Worker API `/edge_worker/v1/rpcapi` endpoint for Airflow 2.10."""
+    # Note: Except the method map this _was_ a 100% copy of internal API module
+    #       airflow.api_internal.endpoints.rpc_api_endpoint.internal_airflow_api()
+    # As of rework for FastAPI in Airflow 3.0, this is updated and to be removed in the future.
+    from airflow.api_internal.endpoints.rpc_api_endpoint import (  # type: ignore[attr-defined]
+        # Note: This is just for compatibility with Airflow 2.10, not working for Airflow 3 / main as removed
+        initialize_method_map,
+    )
+
+    try:
+        if request.headers.get("Content-Type", "") != "application/json":
+            raise HTTPException(status.HTTP_403_FORBIDDEN, "Expected Content-Type: application/json")
+        if request.headers.get("Accept", "") != "application/json":
+            raise HTTPException(status.HTTP_403_FORBIDDEN, "Expected Accept: application/json")
+        auth = request.headers.get("Authorization", "")
+        request_obj = JsonRpcRequest(method=body["method"], jsonrpc=body["jsonrpc"], params=body["params"])
+        jwt_token_authorization_rpc(request_obj, auth)
+        if request_obj.jsonrpc != "2.0":
+            raise error_response("Expected jsonrpc 2.0 request.", status.HTTP_400_BAD_REQUEST)
+
+        log.debug("Got request for %s", request_obj.method)
+        methods_map = initialize_method_map()
+        if request_obj.method not in methods_map:
+            raise error_response(f"Unrecognized method: {request_obj.method}.", status.HTTP_400_BAD_REQUEST)
+
+        handler = methods_map[request_obj.method]
+        params = {}
+        try:
+            if request_obj.params:
+                # Note, this is Airflow 2.10 specific, as it uses Pydantic models for serialization
+                params = BaseSerialization.deserialize(request_obj.params, use_pydantic_models=True)  # type: ignore[call-arg]
+        except Exception:
+            raise error_response("Error deserializing parameters.", status.HTTP_400_BAD_REQUEST)
+
+        log.debug("Calling method %s\nparams: %s", request_obj.method, params)
+        try:
+            # Session must be created there as it may be needed by serializer for lazy-loaded fields.
+            with create_session() as session:
+                output = handler(**params, session=session)
+                # Note, this is Airflow 2.10 specific, as it uses Pydantic models for serialization
+                output_json = BaseSerialization.serialize(output, use_pydantic_models=True)  # type: ignore[call-arg]
+                log.debug(
+                    "Sending response: %s", json.dumps(output_json) if output_json is not None else None
+                )
+        # In case of AirflowException or other selective known types, transport the exception class back to caller
+        except (KeyError, AttributeError, AirflowException) as e:
+            # Note, this is Airflow 2.10 specific, as it uses Pydantic models for serialization
+            output_json = BaseSerialization.serialize(e, use_pydantic_models=True)  # type: ignore[call-arg]
+            log.debug(
+                "Sending exception response: %s", json.dumps(output_json) if output_json is not None else None
+            )
+        except Exception:
+            raise error_response(
+                f"Error executing method '{request_obj.method}'.", status.HTTP_500_INTERNAL_SERVER_ERROR
+            )
+        response = json.dumps(output_json) if output_json is not None else None
+        return Response(response=response, headers={"Content-Type": "application/json"})
+    except HTTPException as e:
+        return e.to_response()  # type: ignore[attr-defined]
+
+
+def jwt_token_authorization_v2(method: str, authorization: str):
+    """Proxy for v2 method path handling."""
+    PREFIX = "/edge_worker/v1/"
+    method_path = method[method.find(PREFIX) + len(PREFIX) :] if PREFIX in method else method
+    jwt_token_authorization(method_path, authorization)
+
+
+@provide_session
+def register_v2(worker_name: str, body: dict[str, Any], session=NEW_SESSION) -> Any:
+    """Handle Edge Worker API `/edge_worker/v1/worker/{worker_name}` endpoint for Airflow 2.10."""
+    try:
+        auth = request.headers.get("Authorization", "")
+        jwt_token_authorization_v2(request.path, auth)
+        request_obj = WorkerStateBody(
+            state=body["state"], jobs_active=0, queues=body["queues"], sysinfo=body["sysinfo"]
+        )
+        return register(worker_name, request_obj, session).model_dump()
+    except HTTPException as e:
+        return e.to_response()  # type: ignore[attr-defined]
+
+
+@provide_session
+def set_state_v2(worker_name: str, body: dict[str, Any], session=NEW_SESSION) -> Any:
+    """Handle Edge Worker API `/edge_worker/v1/worker/{worker_name}` endpoint for Airflow 2.10."""
+    try:
+        auth = request.headers.get("Authorization", "")
+        jwt_token_authorization_v2(request.path, auth)
+        request_obj = WorkerStateBody(
+            state=body["state"],
+            jobs_active=body["jobs_active"],
+            queues=body["queues"],
+            sysinfo=body["sysinfo"],
+            maintenance_comments=body.get("maintenance_comments"),
+        )
+        return set_state(worker_name, request_obj, session).model_dump()
+    except HTTPException as e:
+        return e.to_response()  # type: ignore[attr-defined]
+
+
+@provide_session
+def job_fetch_v2(worker_name: str, body: dict[str, Any], session=NEW_SESSION) -> Any:
+    """Handle Edge Worker API `/edge_worker/v1/jobs/fetch/{worker_name}` endpoint for Airflow 2.10."""
+    from flask import request
+
+    try:
+        auth = request.headers.get("Authorization", "")
+        jwt_token_authorization_v2(request.path, auth)
+        queues = body.get("queues")
+        free_concurrency = body.get("free_concurrency", 1)
+        request_obj = WorkerQueuesBody(queues=queues, free_concurrency=free_concurrency)
+        job: EdgeJobFetched | None = fetch(worker_name, request_obj, session)
+        return job.model_dump() if job is not None else None
+    except HTTPException as e:
+        return e.to_response()  # type: ignore[attr-defined]
+
+
+@provide_session
+def job_state_v2(
+    dag_id: str,
+    task_id: str,
+    run_id: str,
+    try_number: int,
+    map_index: str,  # Note: Connexion can not have negative numbers in path parameters, use string therefore
+    state: TaskInstanceState,
+    session=NEW_SESSION,
+) -> Any:
+    """Handle Edge Worker API `/jobs/state/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}/{state}` endpoint for Airflow 2.10."""
+    from flask import request
+
+    try:
+        auth = request.headers.get("Authorization", "")
+        jwt_token_authorization_v2(request.path, auth)
+        state_api(dag_id, task_id, run_id, try_number, int(map_index), state, session)
+    except HTTPException as e:
+        return e.to_response()  # type: ignore[attr-defined]
+
+
+def logfile_path_v2(
+    dag_id: str,
+    task_id: str,
+    run_id: str,
+    try_number: int,
+    map_index: str,  # Note: Connexion can not have negative numbers in path parameters, use string therefore
+) -> str:
+    """Handle Edge Worker API `/edge_worker/v1/logs/logfile_path/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}` endpoint for Airflow 2.10."""
+    try:
+        auth = request.headers.get("Authorization", "")
+        jwt_token_authorization_v2(request.path, auth)
+        return logfile_path(dag_id, task_id, run_id, try_number, int(map_index))
+    except HTTPException as e:
+        return e.to_response()  # type: ignore[attr-defined]
+
+
+def push_logs_v2(
+    dag_id: str,
+    task_id: str,
+    run_id: str,
+    try_number: int,
+    map_index: str,  # Note: Connexion can not have negative numbers in path parameters, use string therefore
+    body: dict[str, Any],
+) -> None:
+    """Handle Edge Worker API `/edge_worker/v1/logs/push/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}` endpoint for Airflow 2.10."""
+    try:
+        auth = request.headers.get("Authorization", "")
+        jwt_token_authorization_v2(request.path, auth)
+        request_obj = PushLogsBody(
+            log_chunk_data=body["log_chunk_data"], log_chunk_time=body["log_chunk_time"]
+        )
+        with create_session() as session:
+            push_logs(dag_id, task_id, run_id, try_number, int(map_index), request_obj, session)
+    except HTTPException as e:
+        return e.to_response()  # type: ignore[attr-defined]

airflow/providers/edge3/worker_api/routes/health.py

@@ -0,0 +1,28 @@
+# 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.providers.edge3.worker_api.routes._v2_compat import AirflowRouter
+
+health_router = AirflowRouter(tags=["Health"])
+
+
+@health_router.get("/health")
+def health() -> dict[str, str]:
+    """Report API Health."""
+    return {"status": "healthy"}

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

@@ -0,0 +1,162 @@
+# 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 typing import Annotated
+
+from sqlalchemy import select, update
+
+from airflow.providers.edge3.models.edge_job import EdgeJobModel
+from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
+from airflow.providers.edge3.worker_api.datamodels import (
+    EdgeJobFetched,
+    WorkerApiDocs,
+    WorkerQueuesBody,
+)
+from airflow.providers.edge3.worker_api.routes._v2_compat import (
+    AirflowRouter,
+    Body,
+    Depends,
+    SessionDep,
+    create_openapi_http_exception_doc,
+    parse_command,
+    status,
+)
+from airflow.stats import Stats
+from airflow.utils import timezone
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.state import TaskInstanceState
+
+jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs")
+
+
+@jobs_router.post(
+    "/fetch/{worker_name}",
+    dependencies=[Depends(jwt_token_authorization_rest)],
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_400_BAD_REQUEST,
+            status.HTTP_403_FORBIDDEN,
+        ]
+    ),
+)
+def fetch(
+    worker_name: str,
+    body: Annotated[
+        WorkerQueuesBody,
+        Body(
+            title="Log data chunks",
+            description="The queues and capacity from which the worker can fetch jobs.",
+        ),
+    ],
+    session: SessionDep,
+) -> EdgeJobFetched | None:
+    """Fetch a job to execute on the edge worker."""
+    query = (
+        select(EdgeJobModel)
+        .where(
+            EdgeJobModel.state == TaskInstanceState.QUEUED,
+            EdgeJobModel.concurrency_slots <= body.free_concurrency,
+        )
+        .order_by(EdgeJobModel.queued_dttm)
+    )
+    if body.queues:
+        query = query.where(EdgeJobModel.queue.in_(body.queues))
+    query = query.limit(1)
+    query = with_row_locks(query, of=EdgeJobModel, session=session, skip_locked=True)
+    job: EdgeJobModel = session.scalar(query)
+    if not job:
+        return None
+    job.state = TaskInstanceState.RUNNING
+    job.edge_worker = worker_name
+    job.last_update = timezone.utcnow()
+    session.commit()
+    # Edge worker does not backport emitted Airflow metrics, so export some metrics
+    tags = {"dag_id": job.dag_id, "task_id": job.task_id, "queue": job.queue}
+    Stats.incr(f"edge_worker.ti.start.{job.queue}.{job.dag_id}.{job.task_id}", tags=tags)
+    Stats.incr("edge_worker.ti.start", tags=tags)
+    return EdgeJobFetched(
+        dag_id=job.dag_id,
+        task_id=job.task_id,
+        run_id=job.run_id,
+        map_index=job.map_index,
+        try_number=job.try_number,
+        command=parse_command(job.command),
+        concurrency_slots=job.concurrency_slots,
+    )
+
+
+@jobs_router.patch(
+    "/state/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}/{state}",
+    dependencies=[Depends(jwt_token_authorization_rest)],
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_400_BAD_REQUEST,
+            status.HTTP_403_FORBIDDEN,
+        ]
+    ),
+)
+def state(
+    dag_id: Annotated[str, WorkerApiDocs.dag_id],
+    task_id: Annotated[str, WorkerApiDocs.task_id],
+    run_id: Annotated[str, WorkerApiDocs.run_id],
+    try_number: Annotated[int, WorkerApiDocs.try_number],
+    map_index: Annotated[int, WorkerApiDocs.map_index],
+    state: Annotated[TaskInstanceState, WorkerApiDocs.state],
+    session: SessionDep,
+) -> None:
+    """Update the state of a job running on the edge worker."""
+    # execute query to catch the queue and check if state toggles to success or failed
+    # otherwise possible that Executor resets orphaned jobs and stats are exported 2 times
+    if state in [TaskInstanceState.SUCCESS, state == TaskInstanceState.FAILED]:
+        query = select(EdgeJobModel).where(
+            EdgeJobModel.dag_id == dag_id,
+            EdgeJobModel.task_id == task_id,
+            EdgeJobModel.run_id == run_id,
+            EdgeJobModel.map_index == map_index,
+            EdgeJobModel.try_number == try_number,
+            EdgeJobModel.state == TaskInstanceState.RUNNING,
+        )
+        job = session.scalar(query)
+
+        if job:
+            # Edge worker does not backport emitted Airflow metrics, so export some metrics
+            tags = {
+                "dag_id": job.dag_id,
+                "task_id": job.task_id,
+                "queue": job.queue,
+                "state": str(state),
+            }
+            Stats.incr(
+                f"edge_worker.ti.finish.{job.queue}.{state}.{job.dag_id}.{job.task_id}",
+                tags=tags,
+            )
+            Stats.incr("edge_worker.ti.finish", tags=tags)
+
+    query = (
+        update(EdgeJobModel)
+        .where(
+            EdgeJobModel.dag_id == dag_id,
+            EdgeJobModel.task_id == task_id,
+            EdgeJobModel.run_id == run_id,
+            EdgeJobModel.map_index == map_index,
+            EdgeJobModel.try_number == try_number,
+        )
+        .values(state=state, last_update=timezone.utcnow())
+    )
+    session.execute(query)

airflow/providers/edge3/worker_api/routes/logs.py

@@ -0,0 +1,133 @@
+# 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 functools import cache
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated
+
+from airflow.configuration import conf
+from airflow.models.taskinstance import TaskInstance
+from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
+from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
+from airflow.providers.edge3.worker_api.datamodels import PushLogsBody, WorkerApiDocs
+from airflow.providers.edge3.worker_api.routes._v2_compat import (
+    AirflowRouter,
+    Body,
+    Depends,
+    SessionDep,
+    create_openapi_http_exception_doc,
+    status,
+)
+from airflow.utils.log.file_task_handler import FileTaskHandler
+from airflow.utils.session import NEW_SESSION, provide_session
+
+logs_router = AirflowRouter(tags=["Logs"], prefix="/logs")
+
+
+@cache
+@provide_session
+def _logfile_path(task: TaskInstanceKey, session=NEW_SESSION) -> str:
+    """Elaborate the (relative) path and filename to expect from task execution."""
+    ti = TaskInstance.get_task_instance(
+        dag_id=task.dag_id,
+        run_id=task.run_id,
+        task_id=task.task_id,
+        map_index=task.map_index,
+        session=session,
+    )
+    if TYPE_CHECKING:
+        assert ti
+        assert isinstance(ti, TaskInstance)
+    return FileTaskHandler(".")._render_filename(ti, task.try_number)
+
+
+@logs_router.get(
+    "/logfile_path/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}",
+    dependencies=[Depends(jwt_token_authorization_rest)],
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_400_BAD_REQUEST,
+            status.HTTP_403_FORBIDDEN,
+        ]
+    ),
+)
+def logfile_path(
+    dag_id: Annotated[str, WorkerApiDocs.dag_id],
+    task_id: Annotated[str, WorkerApiDocs.task_id],
+    run_id: Annotated[str, WorkerApiDocs.run_id],
+    try_number: Annotated[int, WorkerApiDocs.try_number],
+    map_index: Annotated[int, WorkerApiDocs.map_index],
+) -> str:
+    """Elaborate the path and filename to expect from task execution."""
+    task = TaskInstanceKey(
+        dag_id=dag_id, task_id=task_id, run_id=run_id, try_number=try_number, map_index=map_index
+    )
+    return _logfile_path(task)
+
+
+@logs_router.post(
+    "/push/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}",
+    dependencies=[Depends(jwt_token_authorization_rest)],
+    responses=create_openapi_http_exception_doc(
+        [
+            status.HTTP_400_BAD_REQUEST,
+            status.HTTP_403_FORBIDDEN,
+        ]
+    ),
+)
+def push_logs(
+    dag_id: Annotated[str, WorkerApiDocs.dag_id],
+    task_id: Annotated[str, WorkerApiDocs.task_id],
+    run_id: Annotated[str, WorkerApiDocs.run_id],
+    try_number: Annotated[int, WorkerApiDocs.try_number],
+    map_index: Annotated[int, WorkerApiDocs.map_index],
+    body: Annotated[
+        PushLogsBody,
+        Body(
+            title="Log data chunks",
+            description="The worker remote has no access to log sink and with this can send log chunks to the central site.",
+        ),
+    ],
+    session: SessionDep,
+) -> None:
+    """Push an incremental log chunk from Edge Worker to central site."""
+    log_chunk = EdgeLogsModel(
+        dag_id=dag_id,
+        task_id=task_id,
+        run_id=run_id,
+        map_index=map_index,
+        try_number=try_number,
+        log_chunk_time=body.log_chunk_time,
+        log_chunk_data=body.log_chunk_data,
+    )
+    session.add(log_chunk)
+    # Write logs to local file to make them accessible
+    task = TaskInstanceKey(
+        dag_id=dag_id, task_id=task_id, run_id=run_id, try_number=try_number, map_index=map_index
+    )
+    base_log_folder = conf.get("logging", "base_log_folder", fallback="NOT AVAILABLE")
+    logfile_path = Path(base_log_folder, _logfile_path(task))
+    if not logfile_path.exists():
+        new_folder_permissions = int(
+            conf.get("logging", "file_task_handler_new_folder_permissions", fallback="0o775"), 8
+        )
+        logfile_path.parent.mkdir(parents=True, exist_ok=True, mode=new_folder_permissions)
+    with logfile_path.open("a") as logfile:
+        logfile.write(body.log_chunk_data)