apache-airflow-providers-edge3 1.3.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

airflow/providers/edge3/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "1.3.1"
+__version__ = "1.4.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/edge3/cli/worker.py

@@ -34,6 +34,7 @@ from requests import HTTPError
 
 from airflow import __version__ as airflow_version
 from airflow.configuration import conf
+from airflow.providers.common.compat.sdk import timezone
 from airflow.providers.edge3 import __version__ as edge_provider_version
 from airflow.providers.edge3.cli.api_client import (
     jobs_fetch,
@@ -52,7 +53,6 @@ from airflow.providers.edge3.cli.signalling import (
 )
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerState, EdgeWorkerVersionException
 from airflow.providers.edge3.version_compat import AIRFLOW_V_3_0_PLUS
-from airflow.utils import timezone
 from airflow.utils.net import getfqdn
 from airflow.utils.state import TaskInstanceState
 

airflow/providers/edge3/example_dags/win_test.py

@@ -67,7 +67,7 @@ if TYPE_CHECKING:
 try:
     from airflow.operators.python import PythonOperator
 except ImportError:
-    from airflow.providers.common.compat.standard.operators import PythonOperator
+    from airflow.providers.common.compat.standard.operators import PythonOperator  # type: ignore[no-redef]
 
 
 class CmdOperator(BaseOperator):

airflow/providers/edge3/executors/edge_executor.py

@@ -31,13 +31,13 @@ from airflow.cli.cli_config import GroupCommand
 from airflow.configuration import conf
 from airflow.executors.base_executor import BaseExecutor
 from airflow.models.taskinstance import TaskInstance, TaskInstanceState
+from airflow.providers.common.compat.sdk import timezone
 from airflow.providers.edge3.cli.edge_command import EDGE_COMMANDS
 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.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.stats import Stats
-from airflow.utils import timezone
 from airflow.utils.db import DBLocks, create_global_lock
 from airflow.utils.session import NEW_SESSION, provide_session
 
@@ -352,7 +352,7 @@ class EdgeExecutor(BaseExecutor):
                         del self.last_reported_state[job.key]
                     self.fail(job.key)
                 else:
-                    self.last_reported_state[job.key] = job.state
+                    self.last_reported_state[job.key] = TaskInstanceState(job.state)
             if (
                 job.state == TaskInstanceState.SUCCESS
                 and job.last_update_t < (datetime.now() - timedelta(minutes=job_success_purge)).timestamp()
@@ -399,6 +399,35 @@ class EdgeExecutor(BaseExecutor):
     def terminate(self):
         """Terminate the executor is not doing anything."""
 
+    @provide_session
+    def revoke_task(self, *, ti: TaskInstance, session: Session = NEW_SESSION):
+        """
+        Revoke a task instance from the executor.
+
+        This method removes the task from the executor's internal state and deletes
+        the corresponding EdgeJobModel record to prevent edge workers from picking it up.
+
+        :param ti: Task instance to revoke
+        :param session: Database session
+        """
+        # Remove from executor's internal state
+        self.running.discard(ti.key)
+        self.queued_tasks.pop(ti.key, None)
+        if ti.key in self.last_reported_state:
+            del self.last_reported_state[ti.key]
+
+        # Delete the job from the database to prevent edge workers from picking it up
+        session.execute(
+            delete(EdgeJobModel).where(
+                EdgeJobModel.dag_id == ti.dag_id,
+                EdgeJobModel.task_id == ti.task_id,
+                EdgeJobModel.run_id == ti.run_id,
+                EdgeJobModel.map_index == ti.map_index,
+                EdgeJobModel.try_number == ti.try_number,
+            )
+        )
+        self.log.info("Revoked task instance %s from EdgeExecutor", ti.key)
+
     def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[TaskInstance]:
         """
         Try to adopt running task instances that have been abandoned by a SchedulerJob dying.

airflow/providers/edge3/models/edge_job.py

@@ -19,16 +19,17 @@ from __future__ import annotations
 from datetime import datetime
 
 from sqlalchemy import (
-    Column,
     Index,
     Integer,
     String,
     text,
 )
+from sqlalchemy.orm import Mapped
 
 from airflow.models.base import Base, StringID
 from airflow.models.taskinstancekey import TaskInstanceKey
-from airflow.utils import timezone
+from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.common.compat.sqlalchemy.orm import mapped_column
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.sqlalchemy import UtcDateTime
 
@@ -41,18 +42,20 @@ class EdgeJobModel(Base, LoggingMixin):
     """
 
     __tablename__ = "edge_job"
-    dag_id = Column(StringID(), primary_key=True, nullable=False)
-    task_id = Column(StringID(), primary_key=True, nullable=False)
-    run_id = Column(StringID(), primary_key=True, nullable=False)
-    map_index = Column(Integer, primary_key=True, nullable=False, server_default=text("-1"))
-    try_number = Column(Integer, primary_key=True, default=0)
-    state = Column(String(20))
-    queue = Column(String(256))
-    concurrency_slots = Column(Integer)
-    command = Column(String(2048))
-    queued_dttm = Column(UtcDateTime)
-    edge_worker = Column(String(64))
-    last_update = Column(UtcDateTime)
+    dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False)
+    task_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False)
+    run_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False)
+    map_index: Mapped[int] = mapped_column(
+        Integer, primary_key=True, nullable=False, server_default=text("-1")
+    )
+    try_number: Mapped[int] = mapped_column(Integer, primary_key=True, default=0)
+    state: Mapped[str] = mapped_column(String(20))
+    queue: Mapped[str] = mapped_column(String(256))
+    concurrency_slots: Mapped[int] = mapped_column(Integer)
+    command: Mapped[str] = mapped_column(String(2048))
+    queued_dttm: Mapped[datetime | None] = mapped_column(UtcDateTime)
+    edge_worker: Mapped[str | None] = mapped_column(String(64))
+    last_update: Mapped[datetime | None] = mapped_column(UtcDateTime)
 
     def __init__(
         self,

airflow/providers/edge3/models/edge_logs.py

@@ -19,14 +19,15 @@ from __future__ import annotations
 from datetime import datetime
 
 from sqlalchemy import (
-    Column,
     Integer,
     Text,
     text,
 )
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlalchemy.orm import Mapped
 
 from airflow.models.base import Base, StringID
+from airflow.providers.common.compat.sqlalchemy.orm import mapped_column
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.sqlalchemy import UtcDateTime
 
@@ -45,13 +46,15 @@ class EdgeLogsModel(Base, LoggingMixin):
     """
 
     __tablename__ = "edge_logs"
-    dag_id = Column(StringID(), primary_key=True, nullable=False)
-    task_id = Column(StringID(), primary_key=True, nullable=False)
-    run_id = Column(StringID(), primary_key=True, nullable=False)
-    map_index = Column(Integer, primary_key=True, nullable=False, server_default=text("-1"))
-    try_number = Column(Integer, primary_key=True, default=0)
-    log_chunk_time = Column(UtcDateTime, primary_key=True, nullable=False)
-    log_chunk_data = Column(Text().with_variant(MEDIUMTEXT(), "mysql"), nullable=False)
+    dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False)
+    task_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False)
+    run_id: Mapped[str] = mapped_column(StringID(), primary_key=True, nullable=False)
+    map_index: Mapped[int] = mapped_column(
+        Integer, primary_key=True, nullable=False, server_default=text("-1")
+    )
+    try_number: Mapped[int] = mapped_column(Integer, primary_key=True, default=0)
+    log_chunk_time: Mapped[datetime] = mapped_column(UtcDateTime, primary_key=True, nullable=False)
+    log_chunk_data: Mapped[str] = mapped_column(Text().with_variant(MEDIUMTEXT(), "mysql"), nullable=False)
 
     def __init__(
         self,

airflow/providers/edge3/models/edge_worker.py

@@ -23,18 +23,22 @@ from datetime import datetime
 from enum import Enum
 from typing import TYPE_CHECKING
 
-from sqlalchemy import Column, Integer, String, delete, select
+from sqlalchemy import Integer, String, delete, select
+from sqlalchemy.orm import Mapped
 
 from airflow.exceptions import AirflowException
 from airflow.models.base import Base
+from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.common.compat.sqlalchemy.orm import mapped_column
 from airflow.stats import Stats
-from airflow.utils import timezone
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.providers_configuration_loader import providers_configuration_loaded
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.sqlalchemy import UtcDateTime
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from sqlalchemy.orm import Session
 
 logger = logging.getLogger(__name__)
@@ -79,29 +83,29 @@ class EdgeWorkerModel(Base, LoggingMixin):
     """A Edge Worker instance which reports the state and health."""
 
     __tablename__ = "edge_worker"
-    worker_name = Column(String(64), primary_key=True, nullable=False)
-    state = Column(String(20))
-    maintenance_comment = Column(String(1024))
-    _queues = Column("queues", String(256))
-    first_online = Column(UtcDateTime)
-    last_update = Column(UtcDateTime)
-    jobs_active = Column(Integer, default=0)
-    jobs_taken = Column(Integer, default=0)
-    jobs_success = Column(Integer, default=0)
-    jobs_failed = Column(Integer, default=0)
-    sysinfo = Column(String(256))
+    worker_name: Mapped[str] = mapped_column(String(64), primary_key=True, nullable=False)
+    state: Mapped[EdgeWorkerState] = mapped_column(String(20))
+    maintenance_comment: Mapped[str | None] = mapped_column(String(1024))
+    _queues: Mapped[str | None] = mapped_column("queues", String(256))
+    first_online: Mapped[datetime | None] = mapped_column(UtcDateTime)
+    last_update: Mapped[datetime | None] = mapped_column(UtcDateTime)
+    jobs_active: Mapped[int] = mapped_column(Integer, default=0)
+    jobs_taken: Mapped[int] = mapped_column(Integer, default=0)
+    jobs_success: Mapped[int] = mapped_column(Integer, default=0)
+    jobs_failed: Mapped[int] = mapped_column(Integer, default=0)
+    sysinfo: Mapped[str | None] = mapped_column(String(256))
 
     def __init__(
         self,
         worker_name: str,
-        state: str,
+        state: str | EdgeWorkerState,
         queues: list[str] | None,
         first_online: datetime | None = None,
         last_update: datetime | None = None,
         maintenance_comment: str | None = None,
     ):
         self.worker_name = worker_name
-        self.state = state
+        self.state = EdgeWorkerState(state)
         self.queues = queues
         self.first_online = first_online or timezone.utcnow()
         self.last_update = last_update
@@ -139,14 +143,14 @@ class EdgeWorkerModel(Base, LoggingMixin):
                 queues.remove(queue_name)
         self.queues = queues
 
-    def update_state(self, state: str) -> None:
+    def update_state(self, state: str | EdgeWorkerState) -> None:
         """Update state field."""
-        self.state = state
+        self.state = EdgeWorkerState(state)
 
 
 def set_metrics(
     worker_name: str,
-    state: EdgeWorkerState,
+    state: str | EdgeWorkerState,
     jobs_active: int,
     concurrency: int,
     free_concurrency: int,
@@ -204,7 +208,7 @@ def reset_metrics(worker_name: str) -> None:
 @provide_session
 def _fetch_edge_hosts_from_db(
     hostname: str | None = None, states: list | None = None, session: Session = NEW_SESSION
-) -> list:
+) -> Sequence[EdgeWorkerModel]:
     query = select(EdgeWorkerModel)
     if states:
         query = query.where(EdgeWorkerModel.state.in_(states))

airflow/providers/edge3/openapi/v2-edge-generated.yaml

@@ -886,8 +886,6 @@ components:
         token:
           type: string
           title: Token
-        ti:
-          $ref: '#/components/schemas/TaskInstance'
         dag_rel_path:
           type: string
           format: path
@@ -899,6 +897,8 @@ components:
           - type: string
           - type: 'null'
           title: Log Path
+        ti:
+          $ref: '#/components/schemas/TaskInstance'
         type:
           type: string
           const: ExecuteTask
@@ -907,10 +907,10 @@ components:
       type: object
       required:
       - token
-      - ti
       - dag_rel_path
       - bundle_info
       - log_path
+      - ti
       title: ExecuteTask
       description: Execute the given Task.
     HTTPExceptionResponse:

airflow/providers/edge3/plugins/edge_executor_plugin.py

@@ -227,6 +227,22 @@ else:
     RUNNING_ON_APISERVER = "gunicorn" in sys.argv[0] and "airflow-webserver" in sys.argv
 
 
+def _get_base_url_path(path: str) -> str:
+    """Construct URL path with webserver base_url prefix."""
+    base_url = conf.get("api", "base_url", fallback="/")
+    # Extract pathname from base_url (handles both full URLs and path-only)
+    if base_url.startswith(("http://", "https://")):
+        from urllib.parse import urlparse
+
+        base_path = urlparse(base_url).path
+    else:
+        base_path = base_url
+
+    # Normalize paths: remove trailing slash from base, ensure leading slash on path
+    base_path = base_path.rstrip("/")
+    return base_path + path
+
+
 class EdgeExecutorPlugin(AirflowPlugin):
     """EdgeExecutor Plugin - provides API endpoints for Edge Workers in Webserver."""
 
@@ -237,30 +253,30 @@ class EdgeExecutorPlugin(AirflowPlugin):
             react_apps = [
                 {
                     "name": "Edge Worker",
-                    "bundle_url": "/edge_worker/static/main.umd.cjs",
+                    "bundle_url": _get_base_url_path("/edge_worker/static/main.umd.cjs"),
                     "destination": "nav",
                     "url_route": "edge_worker",
                     "category": "admin",
-                    "icon": "/edge_worker/res/cloud-computer.svg",
-                    "icon_dark_mode": "/edge_worker/res/cloud-computer-dark.svg",
+                    "icon": _get_base_url_path("/edge_worker/res/cloud-computer.svg"),
+                    "icon_dark_mode": _get_base_url_path("/edge_worker/res/cloud-computer-dark.svg"),
                 },
                 {
                     "name": "Edge Worker Jobs",
-                    "bundle_url": "/edge_worker/static/main.umd.cjs",
+                    "bundle_url": _get_base_url_path("/edge_worker/static/main.umd.cjs"),
                     "url_route": "edge_jobs",
                     "category": "admin",
-                    "icon": "/edge_worker/res/cloud-computer.svg",
-                    "icon_dark_mode": "/edge_worker/res/cloud-computer-dark.svg",
+                    "icon": _get_base_url_path("/edge_worker/res/cloud-computer.svg"),
+                    "icon_dark_mode": _get_base_url_path("/edge_worker/res/cloud-computer-dark.svg"),
                 },
             ]
             external_views = [
                 {
                     "name": "Edge Worker API docs",
-                    "href": "/edge_worker/docs",
+                    "href": _get_base_url_path("/edge_worker/docs"),
                     "destination": "nav",
                     "category": "docs",
-                    "icon": "/edge_worker/res/cloud-computer.svg",
-                    "icon_dark_mode": "/edge_worker/res/cloud-computer-dark.svg",
+                    "icon": _get_base_url_path("/edge_worker/res/cloud-computer.svg"),
+                    "icon_dark_mode": _get_base_url_path("/edge_worker/res/cloud-computer-dark.svg"),
                     "url_route": "edge_worker_api_docs",
                 }
             ]
@@ -271,7 +287,7 @@ class EdgeExecutorPlugin(AirflowPlugin):
             appbuilder_menu_items = [
                 {
                     "name": "Edge Worker API docs",
-                    "href": "/edge_worker/v1/ui",
+                    "href": _get_base_url_path("/edge_worker/v1/ui"),
                     "category": "Docs",
                 }
             ]