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

airflow/providers/edge3/worker_api/auth.py

@@ -0,0 +1,135 @@
+# 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
+
+import logging
+from functools import cache
+from uuid import uuid4
+
+from itsdangerous import BadSignature
+from jwt import (
+    ExpiredSignatureError,
+    ImmatureSignatureError,
+    InvalidAudienceError,
+    InvalidIssuedAtError,
+    InvalidSignatureError,
+)
+
+from airflow.configuration import conf
+from airflow.providers.edge3.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.edge3.worker_api.datamodels import JsonRpcRequestBase  # noqa: TCH001
+from airflow.providers.edge3.worker_api.routes._v2_compat import (
+    Header,
+    HTTPException,
+    Request,
+    status,
+)
+
+log = logging.getLogger(__name__)
+
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.api_fastapi.auth.tokens import JWTValidator
+
+    @cache
+    def jwt_validator() -> JWTValidator:
+        clock_grace = conf.getint("core", "internal_api_clock_grace", fallback=30)
+        return JWTValidator(
+            secret_key=conf.get("core", "internal_api_secret_key"),
+            leeway=clock_grace,
+            audience="api",
+        )
+
+    def jwt_validate(authorization: str) -> dict:
+        return jwt_validator().validated_claims(authorization)
+
+else:
+    # Airflow 2.10 compatibility
+    from airflow.utils.jwt_signer import JWTSigner  # type: ignore
+
+    @cache
+    def jwt_signer() -> JWTSigner:
+        clock_grace = conf.getint("core", "internal_api_clock_grace", fallback=30)
+        return JWTSigner(
+            secret_key=conf.get("core", "internal_api_secret_key"),
+            expiration_time_in_seconds=clock_grace,
+            leeway_in_seconds=clock_grace,
+            audience="api",
+        )
+
+    def jwt_validate(authorization: str) -> dict:
+        return jwt_signer().verify_token(authorization)
+
+
+def _forbidden_response(message: str):
+    """Log the error and return the response anonymized."""
+    error_id = uuid4()
+    log.exception("%s error_id=%s", message, error_id)
+    raise HTTPException(
+        status.HTTP_403_FORBIDDEN,
+        f"Forbidden. The server side traceback may be identified with error_id={error_id}",
+    )
+
+
+def jwt_token_authorization(method: str, authorization: str):
+    """Check if the JWT token is correct."""
+    try:
+        payload = jwt_validate(authorization)
+        signed_method = payload.get("method")
+        if not signed_method or signed_method != method:
+            _forbidden_response(
+                "Invalid method in token authorization. "
+                f"signed method='{signed_method}' "
+                f"called method='{method}'",
+            )
+    except BadSignature:
+        _forbidden_response("Bad Signature. Please use only the tokens provided by the API.")
+    except InvalidAudienceError:
+        _forbidden_response("Invalid audience for the request")
+    except InvalidSignatureError:
+        _forbidden_response("The signature of the request was wrong")
+    except ImmatureSignatureError:
+        _forbidden_response("The signature of the request was sent from the future")
+    except ExpiredSignatureError:
+        _forbidden_response(
+            "The signature of the request has expired. Make sure that all components "
+            "in your system have synchronized clocks.",
+        )
+    except InvalidIssuedAtError:
+        _forbidden_response(
+            "The request was issues in the future. Make sure that all components "
+            "in your system have synchronized clocks.",
+        )
+    except Exception:
+        _forbidden_response("Unable to authenticate API via token.")
+
+
+def jwt_token_authorization_rpc(
+    body: JsonRpcRequestBase, authorization: str = Header(description="JWT Authorization Token")
+):
+    """Check if the JWT token is correct for JSON PRC requests."""
+    jwt_token_authorization(body.method, authorization)
+
+
+def jwt_token_authorization_rest(
+    request: Request, authorization: str = Header(description="JWT Authorization Token")
+):
+    """Check if the JWT token is correct for REST API requests."""
+    PREFIX = "/edge_worker/v1/"
+    path = request.url.path
+    method_path = path[path.find(PREFIX) + len(PREFIX) :] if PREFIX in path else path
+    jwt_token_authorization(method_path, authorization)

airflow/providers/edge3/worker_api/datamodels.py

@@ -0,0 +1,190 @@
+# 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 datetime import datetime
+from typing import (
+    Annotated,
+    Any,
+)
+
+from pydantic import BaseModel, Field
+
+from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.providers.edge3.models.edge_worker import EdgeWorkerState  # noqa: TCH001
+from airflow.providers.edge3.worker_api.routes._v2_compat import ExecuteTask, Path
+
+
+class WorkerApiDocs:
+    """Documentation collection for the worker API."""
+
+    dag_id = Path(title="Dag ID", description="Identifier of the DAG to which the task belongs.")
+    task_id = Path(title="Task ID", description="Task name in the DAG.")
+    run_id = Path(title="Run ID", description="Run ID of the DAG execution.")
+    try_number = Path(title="Try Number", description="The number of attempt to execute this task.")
+    map_index = Path(
+        title="Map Index",
+        description="For dynamically mapped tasks the mapping number, -1 if the task is not mapped.",
+    )
+    state = Path(title="Task State", description="State of the assigned task under execution.")
+
+
+class JsonRpcRequestBase(BaseModel):
+    """Base JSON RPC request model to define just the method."""
+
+    method: Annotated[
+        str,
+        Field(description="Fully qualified python module method name that is called via JSON RPC."),
+    ]
+
+
+class JsonRpcRequest(JsonRpcRequestBase):
+    """JSON RPC request model."""
+
+    jsonrpc: Annotated[str, Field(description="JSON RPC Version", examples=["2.0"])]
+    params: Annotated[
+        dict[str, Any] | None,
+        Field(description="Dictionary of parameters passed to the method."),
+    ]
+
+
+class EdgeJobBase(BaseModel):
+    """Basic attributes of a job on the edge worker."""
+
+    dag_id: Annotated[
+        str, Field(title="Dag ID", description="Identifier of the DAG to which the task belongs.")
+    ]
+    task_id: Annotated[str, Field(title="Task ID", description="Task name in the DAG.")]
+    run_id: Annotated[str, Field(title="Run ID", description="Run ID of the DAG execution.")]
+    map_index: Annotated[
+        int,
+        Field(
+            title="Map Index",
+            description="For dynamically mapped tasks the mapping number, -1 if the task is not mapped.",
+        ),
+    ]
+    try_number: Annotated[
+        int, Field(title="Try Number", description="The number of attempt to execute this task.")
+    ]
+
+    @property
+    def key(self) -> TaskInstanceKey:
+        return TaskInstanceKey(self.dag_id, self.task_id, self.run_id, self.try_number, self.map_index)
+
+
+class EdgeJobFetched(EdgeJobBase):
+    """Job that is to be executed on the edge worker."""
+
+    command: Annotated[
+        ExecuteTask,
+        Field(
+            title="Command",
+            description="Command line to use to execute the job in Airflow 2. Task definition in Airflow 3",
+        ),
+    ]
+    concurrency_slots: Annotated[int, Field(description="Number of concurrency slots the job requires.")]
+
+
+class WorkerQueuesBase(BaseModel):
+    """Queues that a worker supports to run jobs on."""
+
+    queues: Annotated[
+        list[str] | None,
+        Field(
+            None,
+            description="List of queues the worker is pulling jobs from. If not provided, worker pulls from all queues.",
+        ),
+    ]
+
+
+class WorkerQueuesBody(WorkerQueuesBase):
+    """Queues that a worker supports to run jobs on."""
+
+    free_concurrency: Annotated[int, Field(description="Number of free concurrency slots on the worker.")]
+
+
+class WorkerStateBody(WorkerQueuesBase):
+    """Details of the worker state sent to the scheduler."""
+
+    state: Annotated[EdgeWorkerState, Field(description="State of the worker from the view of the worker.")]
+    jobs_active: Annotated[int, Field(description="Number of active jobs the worker is running.")] = 0
+    queues: Annotated[
+        list[str] | None,
+        Field(
+            description="List of queues the worker is pulling jobs from. If not provided, worker pulls from all queues."
+        ),
+    ] = None
+    sysinfo: Annotated[
+        dict[str, str | int],
+        Field(
+            description="System information of the worker.",
+            examples=[
+                {
+                    "concurrency": 4,
+                    "free_concurrency": 3,
+                    "airflow_version": "2.0.0",
+                    "edge_provider_version": "1.0.0",
+                }
+            ],
+        ),
+    ]
+    maintenance_comments: Annotated[
+        str | None,
+        Field(description="Comments about the maintenance state of the worker."),
+    ] = None
+
+
+class WorkerQueueUpdateBody(BaseModel):
+    """Changed queues for the worker."""
+
+    new_queues: Annotated[
+        list[str] | None,
+        Field(description="Additional queues to be added to worker."),
+    ]
+    remove_queues: Annotated[
+        list[str] | None,
+        Field(description="Queues to remove from worker."),
+    ]
+
+
+class PushLogsBody(BaseModel):
+    """Incremental new log content from worker."""
+
+    log_chunk_time: Annotated[datetime, Field(description="Time of the log chunk at point of sending.")]
+    log_chunk_data: Annotated[str, Field(description="Log chunk data as incremental log text.")]
+
+
+class WorkerRegistrationReturn(BaseModel):
+    """The return class for the worker registration."""
+
+    last_update: Annotated[datetime, Field(description="Time of the last update of the worker.")]
+
+
+class WorkerSetStateReturn(BaseModel):
+    """The return class for the worker set state."""
+
+    state: Annotated[EdgeWorkerState, Field(description="State of the worker from the view of the server.")]
+    queues: Annotated[
+        list[str] | None,
+        Field(
+            description="List of queues the worker is pulling jobs from. If not provided, worker pulls from all queues."
+        ),
+    ]
+    maintenance_comments: Annotated[
+        str | None,
+        Field(description="Comments about the maintenance state of the worker."),
+    ] = None

airflow/providers/edge3/worker_api/routes/__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.

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

@@ -0,0 +1,135 @@
+# 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 API to provide both FastAPI as well as Connexion based endpoints."""
+
+from __future__ import annotations
+
+from airflow.providers.edge3.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    # Just re-import the types from FastAPI and Airflow Core
+    from fastapi import Body, Depends, Header, HTTPException, Path, Request, status
+
+    from airflow.api_fastapi.common.db.common import SessionDep
+    from airflow.api_fastapi.common.router import AirflowRouter
+    from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+
+    # In Airflow 3 with AIP-72 we get workload addressed by ExecuteTask
+    from airflow.executors.workloads import ExecuteTask
+
+    def parse_command(command: str) -> ExecuteTask:
+        return ExecuteTask.model_validate_json(command)
+else:
+    # Mock the external dependnecies
+    from typing import Callable
+
+    from connexion import ProblemException
+
+    class Body:  # type: ignore[no-redef]
+        def __init__(self, *_, **__):
+            pass
+
+    class Depends:  # type: ignore[no-redef]
+        def __init__(self, *_, **__):
+            pass
+
+    class Header:  # type: ignore[no-redef]
+        def __init__(self, *_, **__):
+            pass
+
+    class Path:  # type: ignore[no-redef]
+        def __init__(self, *_, **__):
+            pass
+
+    class Request:  # type: ignore[no-redef]
+        pass
+
+    class SessionDep:  # type: ignore[no-redef]
+        pass
+
+    def create_openapi_http_exception_doc(responses_status_code: list[int]) -> dict:
+        return {}
+
+    class status:  # type: ignore[no-redef]
+        HTTP_204_NO_CONTENT = 204
+        HTTP_400_BAD_REQUEST = 400
+        HTTP_403_FORBIDDEN = 403
+        HTTP_500_INTERNAL_SERVER_ERROR = 500
+
+    class HTTPException(ProblemException):  # type: ignore[no-redef]
+        """Raise when the user does not have the required permissions."""
+
+        def __init__(
+            self,
+            status: int,
+            detail: str,
+        ) -> None:
+            from airflow.utils.docs import get_docs_url
+
+            doc_link = get_docs_url("stable-rest-api-ref.html")
+            EXCEPTIONS_LINK_MAP = {
+                400: f"{doc_link}#section/Errors/BadRequest",
+                403: f"{doc_link}#section/Errors/PermissionDenied",
+                500: f"{doc_link}#section/Errors/Unknown",
+            }
+            TITLE_MAP = {
+                400: "BadRequest",
+                403: "PermissionDenied",
+                500: "InternalServerError",
+            }
+            super().__init__(
+                status=status,
+                type=EXCEPTIONS_LINK_MAP[status],
+                title=TITLE_MAP[status],
+                detail=detail,
+            )
+
+        def to_response(self):
+            from flask import Response
+
+            return Response(response=self.detail, status=self.status)
+
+    class AirflowRouter:  # type: ignore[no-redef]
+        def __init__(self, *_, **__):
+            pass
+
+        def get(self, *_, **__):
+            def decorator(func: Callable) -> Callable:
+                return func
+
+            return decorator
+
+        def post(self, *_, **__):
+            def decorator(func: Callable) -> Callable:
+                return func
+
+            return decorator
+
+        def patch(self, *_, **__):
+            def decorator(func: Callable) -> Callable:
+                return func
+
+            return decorator
+
+    # In Airflow 3 with AIP-72 we get workload addressed by ExecuteTask
+    # But in Airflow 2.10 it is a command line array
+    ExecuteTask = list[str]  # type: ignore[no-redef,assignment,misc]
+
+    def parse_command(command: str) -> ExecuteTask:
+        from ast import literal_eval
+
+        return literal_eval(command)