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

airflow/providers/edge3/__init__.py

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

airflow/providers/edge3/cli/api_client.py

@@ -82,10 +82,9 @@ def _make_generic_request(method: str, rest_path: str, data: str | None = None)
 
         @cache
         def jwt_generator() -> JWTGenerator:
-            clock_grace = conf.getint("core", "internal_api_clock_grace", fallback=30)
             return JWTGenerator(
-                secret_key=conf.get("core", "internal_api_secret_key"),
-                valid_for=clock_grace,
+                secret_key=conf.get("api_auth", "jwt_secret"),
+                valid_for=conf.getint("api_auth", "jwt_leeway", fallback=30),
                 audience="api",
             )
 

airflow/providers/edge3/cli/edge_command.py

@@ -23,6 +23,7 @@ import signal
 import sys
 from dataclasses import asdict
 from datetime import datetime
+from getpass import getuser
 from http import HTTPStatus
 from multiprocessing import Process
 from pathlib import Path
@@ -36,6 +37,8 @@ from requests import HTTPError
 
 from airflow import __version__ as airflow_version, settings
 from airflow.cli.cli_config import ARG_PID, ARG_VERBOSE, ActionCommand, Arg
+from airflow.cli.commands.daemon_utils import run_command_with_daemon_option
+from airflow.cli.simple_table import AirflowConsole
 from airflow.configuration import conf
 from airflow.providers.edge3 import __version__ as edge_provider_version
 from airflow.providers.edge3.cli.api_client import (
@@ -164,7 +167,6 @@ class _EdgeWorkerCli:
     """Flag if job processing should be completed and no new jobs fetched for maintenance mode. """
     maintenance_comments: str | None = None
     """Comments for maintenance mode."""
-
     edge_instance: _EdgeWorkerCli | None = None
     """Singleton instance of the worker."""
 
@@ -176,6 +178,7 @@ class _EdgeWorkerCli:
         concurrency: int,
         job_poll_interval: int,
         heartbeat_interval: int,
+        daemon: bool = False,
     ):
         self.pid_file_path = pid_file_path
         self.job_poll_interval = job_poll_interval
@@ -184,6 +187,7 @@ class _EdgeWorkerCli:
         self.queues = queues
         self.concurrency = concurrency
         self.free_concurrency = concurrency
+        self.daemon = daemon
 
         _EdgeWorkerCli.edge_instance = self
 
@@ -271,6 +275,12 @@ class _EdgeWorkerCli:
             setproctitle(f"airflow edge worker: {workload.ti.key}")
 
             try:
+                base_url = conf.get("api", "base_url", fallback="/")
+                # If it's a relative URL, use localhost:8080 as the default
+                if base_url.startswith("/"):
+                    base_url = f"http://localhost:8080{base_url}"
+                default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+
                 supervise(
                     # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this.
                     # Same like in airflow/executors/local_executor.py:_execute_work()
@@ -278,7 +288,9 @@ class _EdgeWorkerCli:
                     dag_rel_path=workload.dag_rel_path,
                     bundle_info=workload.bundle_info,
                     token=workload.token,
-                    server=conf.get("core", "execution_api_server_url"),
+                    server=conf.get(
+                        "core", "execution_api_server_url", fallback=default_execution_api_server
+                    ),
                     log_path=workload.log_path,
                 )
                 return 0
@@ -332,7 +344,8 @@ class _EdgeWorkerCli:
             if e.response.status_code == HTTPStatus.NOT_FOUND:
                 raise SystemExit("Error: API endpoint is not ready, please set [edge] api_enabled=True.")
             raise SystemExit(str(e))
-        _write_pid_to_pidfile(self.pid_file_path)
+        if not self.daemon:
+            _write_pid_to_pidfile(self.pid_file_path)
         signal.signal(signal.SIGINT, _EdgeWorkerCli.signal_handler)
         signal.signal(SIG_STATUS, _EdgeWorkerCli.signal_handler)
         signal.signal(signal.SIGTERM, self.shutdown_handler)
@@ -358,7 +371,8 @@ class _EdgeWorkerCli:
             except EdgeWorkerVersionException:
                 logger.info("Version mismatch of Edge worker and Core. Quitting worker anyway.")
         finally:
-            remove_existing_pidfile(self.pid_file_path)
+            if not self.daemon:
+                remove_existing_pidfile(self.pid_file_path)
 
     def loop(self):
         """Run a loop of scheduling and monitoring tasks."""
@@ -463,6 +477,9 @@ class _EdgeWorkerCli:
                 _EdgeWorkerCli.maintenance_comments = worker_info.maintenance_comments
             else:
                 _EdgeWorkerCli.maintenance_comments = None
+            if worker_info.state == EdgeWorkerState.SHUTDOWN_REQUEST:
+                logger.info("Shutdown requested!")
+                _EdgeWorkerCli.drain = True
 
             worker_state_changed = worker_info.state != state
         except EdgeWorkerVersionException:
@@ -479,10 +496,8 @@ class _EdgeWorkerCli:
                 return
 
 
-@cli_utils.action_cli(check_db=False)
 @providers_configuration_loaded
-def worker(args):
-    """Start Airflow Edge Worker."""
+def _launch_worker(args):
     print(settings.HEADER)
     print(EDGE_WORKER_HEADER)
 
@@ -493,14 +508,31 @@ def worker(args):
         concurrency=args.concurrency,
         job_poll_interval=conf.getint("edge", "job_poll_interval"),
         heartbeat_interval=conf.getint("edge", "heartbeat_interval"),
+        daemon=args.daemon,
     )
     edge_worker.start()
 
 
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def worker(args):
+    """Start Airflow Edge Worker."""
+    umask = args.umask or conf.get("edge", "worker_umask", fallback=settings.DAEMON_UMASK)
+
+    run_command_with_daemon_option(
+        args=args,
+        process_name=EDGE_WORKER_PROCESS_NAME,
+        callback=lambda: _launch_worker(args),
+        should_setup_logging=True,
+        pid_file=_pid_file_path(args.pid),
+        umask=umask,
+    )
+
+
 @cli_utils.action_cli(check_db=False)
 @providers_configuration_loaded
 def status(args):
-    """Check for Airflow Edge Worker status."""
+    """Check for Airflow Local Edge Worker status."""
     pid = _get_pid(args.pid)
 
     # Send Signal as notification to drop status JSON
@@ -526,7 +558,7 @@ def status(args):
 @cli_utils.action_cli(check_db=False)
 @providers_configuration_loaded
 def maintenance(args):
-    """Set or Unset maintenance mode of worker."""
+    """Set or Unset maintenance mode of local edge worker."""
     if args.maintenance == "on" and not args.comments:
         logger.error("Comments are required when setting maintenance mode.")
         sys.exit(4)
@@ -597,7 +629,7 @@ def maintenance(args):
 @cli_utils.action_cli(check_db=False)
 @providers_configuration_loaded
 def stop(args):
-    """Stop a running Airflow Edge Worker."""
+    """Stop a running local Airflow Edge Worker."""
     pid = _get_pid(args.pid)
     # Send SIGINT
     logger.info("Sending SIGINT to worker pid %i.", pid)
@@ -611,6 +643,109 @@ def stop(args):
         logger.info("Worker has been shut down.")
 
 
+def _check_valid_db_connection():
+    """Check for a valid db connection before executing db dependent cli commands."""
+    db_conn = conf.get("database", "sql_alchemy_conn")
+    db_default = conf.get_default_value("database", "sql_alchemy_conn")
+    if db_conn == db_default:
+        raise SystemExit(
+            "Error: The database connection is not set. Please set the connection in the configuration file."
+        )
+
+
+def _check_if_registered_edge_host(hostname: str):
+    """Check if edge worker is registered with the db before executing dependent cli commands."""
+    from airflow.providers.edge3.models.edge_worker import _fetch_edge_hosts_from_db
+
+    if not _fetch_edge_hosts_from_db(hostname=hostname):
+        raise SystemExit(f"Error: Edge Worker {hostname} is unknown!")
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def list_edge_workers(args) -> None:
+    """Query the db to list all registered edge workers."""
+    _check_valid_db_connection()
+    from airflow.providers.edge3.models.edge_worker import get_registered_edge_hosts
+
+    all_hosts_iter = get_registered_edge_hosts(states=args.state)
+    # Format and print worker info on the screen
+    fields = [
+        "worker_name",
+        "state",
+        "queues",
+        "maintenance_comment",
+    ]
+    all_hosts = [{f: host.__getattribute__(f) for f in fields} for host in all_hosts_iter]
+    AirflowConsole().print_as(data=all_hosts, output=args.output)
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def put_remote_worker_on_maintenance(args) -> None:
+    """Put remote edge worker on maintenance."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import request_maintenance
+
+    request_maintenance(args.edge_hostname, args.comments)
+    logger.info("%s has been put on maintenance by %s.", args.edge_hostname, getuser())
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def remove_remote_worker_from_maintenance(args) -> None:
+    """Remove remote edge worker from maintenance."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import exit_maintenance
+
+    exit_maintenance(args.edge_hostname)
+    logger.info("%s has been removed from maintenance by %s.", args.edge_hostname, getuser())
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def remote_worker_update_maintenance_comment(args) -> None:
+    """Update maintenance comments of the remote edge worker."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import change_maintenance_comment
+
+    try:
+        change_maintenance_comment(args.edge_hostname, args.comments)
+        logger.info("Maintenance comments updated for %s by %s.", args.edge_hostname, getuser())
+    except TypeError:
+        raise SystemExit
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
+    """Remove remote edge worker entry from db."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import remove_worker
+
+    try:
+        remove_worker(args.edge_hostname)
+        logger.info("Edge Worker host %s removed by %s.", args.edge_hostname, getuser())
+    except TypeError:
+        raise SystemExit
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def remote_worker_request_shutdown(args) -> None:
+    """Initiate the shutdown of the remote edge worker."""
+    _check_valid_db_connection()
+    _check_if_registered_edge_host(hostname=args.edge_hostname)
+    from airflow.providers.edge3.models.edge_worker import request_shutdown
+
+    request_shutdown(args.edge_hostname)
+    logger.info("Requested shutdown of Edge Worker host %s by %s.", args.edge_hostname, getuser())
+
+
 ARG_CONCURRENCY = Arg(
     ("-c", "--concurrency"),
     type=int,
@@ -625,11 +760,21 @@ ARG_EDGE_HOSTNAME = Arg(
     ("-H", "--edge-hostname"),
     help="Set the hostname of worker if you have multiple workers on a single machine",
 )
+ARG_REQUIRED_EDGE_HOSTNAME = Arg(
+    ("-H", "--edge-hostname"),
+    help="Set the hostname of worker if you have multiple workers on a single machine",
+    required=True,
+)
 ARG_MAINTENANCE = Arg(("maintenance",), help="Desired maintenance state", choices=("on", "off"))
 ARG_MAINTENANCE_COMMENT = Arg(
     ("-c", "--comments"),
     help="Maintenance comments to report reason. Required if maintenance is turned on.",
 )
+ARG_REQUIRED_MAINTENANCE_COMMENT = Arg(
+    ("-c", "--comments"),
+    help="Maintenance comments to report reason. Required if enabling maintenance",
+    required=True,
+)
 ARG_WAIT_MAINT = Arg(
     ("-w", "--wait"),
     default=False,
@@ -642,6 +787,36 @@ ARG_WAIT_STOP = Arg(
     help="Wait until edge worker is shut down.",
     action="store_true",
 )
+ARG_OUTPUT = Arg(
+    (
+        "-o",
+        "--output",
+    ),
+    help="Output format. Allowed values: json, yaml, plain, table (default: table)",
+    metavar="(table, json, yaml, plain)",
+    choices=("table", "json", "yaml", "plain"),
+    default="table",
+)
+ARG_STATE = Arg(
+    (
+        "-s",
+        "--state",
+    ),
+    nargs="+",
+    help="State of the edge worker",
+)
+
+ARG_DAEMON = Arg(
+    ("-D", "--daemon"), help="Daemonize instead of running in the foreground", action="store_true"
+)
+ARG_UMASK = Arg(
+    ("-u", "--umask"),
+    help="Set the umask of edge worker in daemon mode",
+)
+ARG_STDERR = Arg(("--stderr",), help="Redirect stderr to this file if run in daemon mode")
+ARG_STDOUT = Arg(("--stdout",), help="Redirect stdout to this file if run in daemon mode")
+ARG_LOG_FILE = Arg(("-l", "--log-file"), help="Location of the log file if run in daemon mode")
+
 EDGE_COMMANDS: list[ActionCommand] = [
     ActionCommand(
         name=worker.__name__,
@@ -653,6 +828,11 @@ EDGE_COMMANDS: list[ActionCommand] = [
             ARG_EDGE_HOSTNAME,
             ARG_PID,
             ARG_VERBOSE,
+            ARG_DAEMON,
+            ARG_STDOUT,
+            ARG_STDERR,
+            ARG_LOG_FILE,
+            ARG_UMASK,
         ),
     ),
     ActionCommand(
@@ -686,4 +866,49 @@ EDGE_COMMANDS: list[ActionCommand] = [
             ARG_VERBOSE,
         ),
     ),
+    ActionCommand(
+        name="list-workers",
+        help=list_edge_workers.__doc__,
+        func=list_edge_workers,
+        args=(
+            ARG_OUTPUT,
+            ARG_STATE,
+        ),
+    ),
+    ActionCommand(
+        name="remote-edge-worker-request-maintenance",
+        help=put_remote_worker_on_maintenance.__doc__,
+        func=put_remote_worker_on_maintenance,
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_REQUIRED_MAINTENANCE_COMMENT,
+        ),
+    ),
+    ActionCommand(
+        name="remote-edge-worker-exit-maintenance",
+        help=remove_remote_worker_from_maintenance.__doc__,
+        func=remove_remote_worker_from_maintenance,
+        args=(ARG_REQUIRED_EDGE_HOSTNAME,),
+    ),
+    ActionCommand(
+        name="remote-edge-worker-update-maintenance-comment",
+        help=remote_worker_update_maintenance_comment.__doc__,
+        func=remote_worker_update_maintenance_comment,
+        args=(
+            ARG_REQUIRED_EDGE_HOSTNAME,
+            ARG_REQUIRED_MAINTENANCE_COMMENT,
+        ),
+    ),
+    ActionCommand(
+        name="remove-remote-edge-worker",
+        help=remove_remote_worker.__doc__,
+        func=remove_remote_worker,
+        args=(ARG_REQUIRED_EDGE_HOSTNAME,),
+    ),
+    ActionCommand(
+        name="shutdown-remote-edge-worker",
+        help=remote_worker_request_shutdown.__doc__,
+        func=remote_worker_request_shutdown,
+        args=(ARG_REQUIRED_EDGE_HOSTNAME,),
+    ),
 ]

airflow/providers/edge3/example_dags/integration_test.py

@@ -26,20 +26,24 @@ from __future__ import annotations
 from datetime import datetime
 from time import sleep
 
-from airflow.decorators import task, task_group
 from airflow.exceptions import AirflowNotFoundException
 from airflow.hooks.base import BaseHook
-from airflow.models.dag import DAG
-from airflow.models.variable import Variable
-from airflow.providers.common.compat.standard.operators import PythonOperator
-from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.sdk import Param
 from airflow.utils.trigger_rule import TriggerRule
 
 try:
     from airflow.providers.standard.operators.bash import BashOperator
+    from airflow.providers.standard.operators.empty import EmptyOperator
+    from airflow.providers.standard.operators.python import PythonOperator
+    from airflow.sdk import DAG, Param, Variable, task, task_group
 except ImportError:
+    # Airflow 2.10 compat
+    from airflow.decorators import task, task_group  # type: ignore[no-redef,attr-defined]
+    from airflow.models.dag import DAG  # type: ignore[no-redef,attr-defined,assignment]
+    from airflow.models.param import Param  # type: ignore[no-redef,attr-defined]
+    from airflow.models.variable import Variable  # type: ignore[no-redef,attr-defined]
     from airflow.operators.bash import BashOperator  # type: ignore[no-redef,attr-defined]
+    from airflow.operators.empty import EmptyOperator  # type: ignore[no-redef,attr-defined]
+    from airflow.operators.python import PythonOperator  # type: ignore[no-redef,attr-defined]
 
 with DAG(
     dag_id="integration_test",
@@ -53,9 +57,18 @@ with DAG(
         "mapping_count": Param(
             4,
             type="integer",
+            minimum=1,
+            maximum=1024,
             title="Mapping Count",
             description="Amount of tasks that should be mapped",
         ),
+        "minutes_to_run": Param(
+            15,
+            type="integer",
+            minimum=1,
+            title="Minutes to run",
+            description="Duration in minutes the long running task should run",
+        ),
     },
 ) as dag:
 
@@ -146,9 +159,10 @@ with DAG(
         [plan_to_fail(), needs_retry()] >> capture_fail()
 
     @task
-    def long_running():
-        print("This task runs for 15 minutes")
-        for i in range(15):
+    def long_running(**context):
+        minutes_to_run: int = context["params"]["minutes_to_run"]
+        print(f"This task runs for {minutes_to_run} minute{'s' if minutes_to_run > 1 else ''}.")
+        for i in range(minutes_to_run):
             sleep(60)
             print(f"Running for {i + 1} minutes now.")
         print("Long running task completed.")

airflow/providers/edge3/executors/edge_executor.py

@@ -177,7 +177,11 @@ class EdgeExecutor(BaseExecutor):
             .with_for_update(skip_locked=True)
             .filter(
                 EdgeWorkerModel.state.not_in(
-                    [EdgeWorkerState.UNKNOWN, EdgeWorkerState.OFFLINE, EdgeWorkerState.OFFLINE_MAINTENANCE]
+                    [
+                        EdgeWorkerState.UNKNOWN,
+                        EdgeWorkerState.OFFLINE,
+                        EdgeWorkerState.OFFLINE_MAINTENANCE,
+                    ]
                 ),
                 EdgeWorkerModel.last_update < (timezone.utcnow() - timedelta(seconds=heartbeat_interval * 5)),
             )
@@ -186,7 +190,17 @@ class EdgeExecutor(BaseExecutor):
 
         for worker in lifeless_workers:
             changed = True
-            worker.state = EdgeWorkerState.UNKNOWN
+            #  If the worker dies in maintenance mode we want to remember it, so it can start in maintenance mode
+            worker.state = (
+                EdgeWorkerState.OFFLINE_MAINTENANCE
+                if worker.state
+                in (
+                    EdgeWorkerState.MAINTENANCE_MODE,
+                    EdgeWorkerState.MAINTENANCE_PENDING,
+                    EdgeWorkerState.MAINTENANCE_REQUEST,
+                )
+                else EdgeWorkerState.UNKNOWN
+            )
             reset_metrics(worker.worker_name)
 
         return changed

airflow/providers/edge3/get_provider_info.py

@@ -25,7 +25,7 @@ def get_provider_info():
     return {
         "package-name": "apache-airflow-providers-edge3",
         "name": "Edge Executor",
-        "description": "Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites\n",
+        "description": "Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites.\n\nWhen tasks need to be executed on remote sites where the connection need to pass through\nfirewalls or other network restrictions, the Edge Worker can be deployed. The Edge Worker\nis a lightweight process with reduced dependencies. The worker only needs to be able to\ncommunicate with the central Airflow site via HTTPS.\n\nIn the central Airflow site the EdgeExecutor is used to orchestrate the work. The EdgeExecutor\nis a custom executor which is used to schedule tasks on the edge workers. The EdgeExecutor can co-exist\nwith other executors (for example CeleryExecutor or KubernetesExecutor) in the same Airflow site.\n\nAdditional REST API endpoints are provided to distribute tasks and manage the edge workers. The endpoints\nare provided by the API server.\n",
         "plugins": [
             {
                 "name": "edge_executor",
@@ -93,6 +93,13 @@ def get_provider_info():
                         "example": None,
                         "default": "524288",
                     },
+                    "worker_umask": {
+                        "description": "The default umask to use for edge worker when run in daemon mode\n\nThis controls the file-creation mode mask which determines the initial value of file permission bits\nfor newly created files.\n\nThis value is treated as an octal-integer.\n",
+                        "version_added": None,
+                        "type": "string",
+                        "default": None,
+                        "example": None,
+                    },
                 },
             }
         },

airflow/providers/edge3/models/edge_worker.py

@@ -18,6 +18,7 @@ from __future__ import annotations
 
 import ast
 import json
+import logging
 from datetime import datetime
 from enum import Enum
 from typing import TYPE_CHECKING
@@ -29,12 +30,15 @@ from airflow.models.base import Base
 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 sqlalchemy.orm import Session
 
+logger = logging.getLogger(__name__)
+
 
 class EdgeWorkerVersionException(AirflowException):
     """Signal a version mismatch between core and Edge Site."""
@@ -51,6 +55,8 @@ class EdgeWorkerState(str, Enum):
     """Edge Worker is actively running a task."""
     IDLE = "idle"
     """Edge Worker is active and waiting for a task."""
+    SHUTDOWN_REQUEST = "shutdown request"
+    """Request to shutdown Edge Worker."""
     TERMINATING = "terminating"
     """Edge Worker is completing work and stopping."""
     OFFLINE = "offline"
@@ -194,6 +200,26 @@ def reset_metrics(worker_name: str) -> None:
     )
 
 
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(
+    hostname: str | None = None, states: list | None = None, session: Session = NEW_SESSION
+) -> list:
+    query = select(EdgeWorkerModel)
+    if states:
+        query = query.where(EdgeWorkerModel.state.in_(states))
+    if hostname:
+        query = query.where(EdgeWorkerModel.worker_name == hostname)
+    query = query.order_by(EdgeWorkerModel.worker_name)
+    return session.scalars(query).all()
+
+
+@providers_configuration_loaded
+@provide_session
+def get_registered_edge_hosts(states: list | None = None, session: Session = NEW_SESSION):
+    return _fetch_edge_hosts_from_db(states=states, session=session)
+
+
 @provide_session
 def request_maintenance(
     worker_name: str, maintenance_comment: str | None, session: Session = NEW_SESSION
@@ -217,7 +243,14 @@ def exit_maintenance(worker_name: str, session: Session = NEW_SESSION) -> None:
 @provide_session
 def remove_worker(worker_name: str, session: Session = NEW_SESSION) -> None:
     """Remove a worker that is offline or just gone from DB."""
-    session.execute(delete(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name))
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if worker.state in (EdgeWorkerState.OFFLINE, EdgeWorkerState.OFFLINE_MAINTENANCE):
+        session.execute(delete(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name))
+    else:
+        error_message = f"Cannot remove edge worker {worker_name} as it is in {worker.state} state!"
+        logger.error(error_message)
+        raise TypeError(error_message)
 
 
 @provide_session
@@ -227,4 +260,22 @@ def change_maintenance_comment(
     """Write maintenance comment in the db."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
     worker: EdgeWorkerModel = session.scalar(query)
-    worker.maintenance_comment = maintenance_comment
+    if worker.state in (EdgeWorkerState.MAINTENANCE_MODE, EdgeWorkerState.OFFLINE_MAINTENANCE):
+        worker.maintenance_comment = maintenance_comment
+    else:
+        error_message = f"Cannot change maintenance comment as {worker_name} is not in maintenance!"
+        logger.error(error_message)
+        raise TypeError(error_message)
+
+
+@provide_session
+def request_shutdown(worker_name: str, session: Session = NEW_SESSION) -> None:
+    """Request to shutdown the edge worker."""
+    query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
+    worker: EdgeWorkerModel = session.scalar(query)
+    if worker.state not in (
+        EdgeWorkerState.OFFLINE,
+        EdgeWorkerState.OFFLINE_MAINTENANCE,
+        EdgeWorkerState.UNKNOWN,
+    ):
+        worker.state = EdgeWorkerState.SHUTDOWN_REQUEST

airflow/providers/edge3/version_compat.py

@@ -32,5 +32,4 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
     return airflow_version.major, airflow_version.minor, airflow_version.micro
 
 
-AIRFLOW_V_2_10_PLUS = get_base_airflow_version_tuple() >= (2, 10, 0)
 AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)

airflow/providers/edge3/worker_api/auth.py

@@ -47,10 +47,9 @@ if AIRFLOW_V_3_0_PLUS:
 
     @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,
+            secret_key=conf.get("api_auth", "jwt_secret"),
+            leeway=conf.getint("api_auth", "jwt_leeway", fallback=30),
             audience="api",
         )
 

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

@@ -114,7 +114,7 @@ _worker_queue_doc = Body(
 
 
 def redefine_state(worker_state: EdgeWorkerState, body_state: EdgeWorkerState) -> EdgeWorkerState:
-    """Redefine the state of the worker based on maintenance request."""
+    """Redefine the state of the worker based on maintenance or shutdown request."""
     if (
         worker_state == EdgeWorkerState.MAINTENANCE_REQUEST
         and body_state
@@ -122,7 +122,12 @@ def redefine_state(worker_state: EdgeWorkerState, body_state: EdgeWorkerState) -
             EdgeWorkerState.MAINTENANCE_PENDING,
             EdgeWorkerState.MAINTENANCE_MODE,
         )
-        or worker_state == EdgeWorkerState.OFFLINE_MAINTENANCE
+        or worker_state
+        in (
+            EdgeWorkerState.OFFLINE_MAINTENANCE,
+            EdgeWorkerState.MAINTENANCE_MODE,
+            EdgeWorkerState.MAINTENANCE_PENDING,
+        )
         and body_state == EdgeWorkerState.STARTING
     ):
         return EdgeWorkerState.MAINTENANCE_REQUEST
@@ -133,6 +138,14 @@ def redefine_state(worker_state: EdgeWorkerState, body_state: EdgeWorkerState) -
         if body_state == EdgeWorkerState.MAINTENANCE_MODE:
             return EdgeWorkerState.IDLE
 
+    if worker_state == EdgeWorkerState.SHUTDOWN_REQUEST:
+        if body_state not in (
+            EdgeWorkerState.OFFLINE_MAINTENANCE,
+            EdgeWorkerState.OFFLINE,
+            EdgeWorkerState.UNKNOWN,
+        ):
+            return EdgeWorkerState.SHUTDOWN_REQUEST
+
     return body_state
 
 

{apache_airflow_providers_edge3-1.0.0rc1.dist-info → apache_airflow_providers_edge3-1.1.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-edge3
-Version: 1.0.0rc1
+Version: 1.1.0rc1
 Summary: Provider package apache-airflow-providers-edge3 for Apache Airflow
 Keywords: airflow-provider,edge3,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -20,18 +20,17 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
-Requires-Dist: apache-airflow>=2.10.0rc0
+Requires-Dist: apache-airflow>=2.10.0rc1
+Requires-Dist: apache-airflow-providers-fab>=1.5.3rc1
 Requires-Dist: pydantic>=2.11.0
 Requires-Dist: retryhttp>=1.2.0,!=1.3.0
-Requires-Dist: apache-airflow-providers-fab ; extra == "fab"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-edge3/1.1.0/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-edge3/1.1.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
-Provides-Extra: fab
 
 
 .. Licensed to the Apache Software Foundation (ASF) under one
@@ -58,10 +57,22 @@ Provides-Extra: fab
 
 Package ``apache-airflow-providers-edge3``
 
-Release: ``1.0.0``
+Release: ``1.1.0``
 
 
-Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites
+Handle edge workers on remote sites via HTTP(s) connection and orchestrates work over distributed sites.
+
+When tasks need to be executed on remote sites where the connection need to pass through
+firewalls or other network restrictions, the Edge Worker can be deployed. The Edge Worker
+is a lightweight process with reduced dependencies. The worker only needs to be able to
+communicate with the central Airflow site via HTTPS.
+
+In the central Airflow site the EdgeExecutor is used to orchestrate the work. The EdgeExecutor
+is a custom executor which is used to schedule tasks on the edge workers. The EdgeExecutor can co-exist
+with other executors (for example CeleryExecutor or KubernetesExecutor) in the same Airflow site.
+
+Additional REST API endpoints are provided to distribute tasks and manage the edge workers. The endpoints
+are provided by the API server.
 
 
 Provider package
@@ -71,7 +82,7 @@ This is a provider package for ``edge3`` provider. All classes for this provider
 are in ``airflow.providers.edge3`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.1.0/>`_.
 
 Installation
 ------------
@@ -85,13 +96,14 @@ The package supports the following python versions: 3.9,3.10,3.11,3.12
 Requirements
 ------------
 
-==================  ===================
-PIP package         Version required
-==================  ===================
-``apache-airflow``  ``>=2.10.0``
-``pydantic``        ``>=2.11.0``
-``retryhttp``       ``>=1.2.0,!=1.3.0``
-==================  ===================
+================================  ===================
+PIP package                       Version required
+================================  ===================
+``apache-airflow``                ``>=2.10.0``
+``apache-airflow-providers-fab``  ``>=1.5.3``
+``pydantic``                      ``>=2.11.0``
+``retryhttp``                     ``>=1.2.0,!=1.3.0``
+================================  ===================
 
 Cross provider package dependencies
 -----------------------------------
@@ -113,5 +125,5 @@ Dependent package
 ==============================================================================================  =======
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-edge3/1.1.0/changelog.html>`_.
 

{apache_airflow_providers_edge3-1.0.0rc1.dist-info → apache_airflow_providers_edge3-1.1.0rc1.dist-info}/RECORD

@@ -1,21 +1,21 @@
 airflow/providers/edge3/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
-airflow/providers/edge3/__init__.py,sha256=hb8OUkltNOLHQS7pnEmS9LOc9NJyajbwwZY_51qn6kk,1494
-airflow/providers/edge3/get_provider_info.py,sha256=T8CMBInzQT7TRJFEFf3tbDK6otJ210gXj4qOaOLyfnE,5548
-airflow/providers/edge3/version_compat.py,sha256=aHg90_DtgoSnQvILFICexMyNlHlALBdaeWqkX3dFDug,1605
+airflow/providers/edge3/__init__.py,sha256=NrV-3Uc00pr80NbPVUZWK9JYceO_jK26xp9BmvwFGnU,1494
+airflow/providers/edge3/get_provider_info.py,sha256=Ek27-dB4UALHUFYoYjtoQIGq0p7zeHcEgmELHvpVmCU,6836
+airflow/providers/edge3/version_compat.py,sha256=j5PCtXvZ71aBjixu-EFTNtVDPsngzzs7os0ZQDgFVDk,1536
 airflow/providers/edge3/cli/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/edge3/cli/api_client.py,sha256=J5l99mwwKXC7Ub_kO6fe_56A8Ue3ag8mIYxD_qPLfuQ,7497
+airflow/providers/edge3/cli/api_client.py,sha256=334KHVB4eMSzRpQ5emS56o-RTUJQprxf5Q3xQldCHDQ,7440
 airflow/providers/edge3/cli/dataclasses.py,sha256=JUuvvmzSVWvG9uOEfzLIiXrTZ-HbESvu50jkPpVIYVw,2895
-airflow/providers/edge3/cli/edge_command.py,sha256=xWNcgvUJyHqT51_Q-6bVPjCd1lrI7KpxSrFfejPvgjA,27466
+airflow/providers/edge3/cli/edge_command.py,sha256=dRKC1VPAUIvJnzvPlmcT98oBGlG0_MK7TBTNvjDEI2k,35658
 airflow/providers/edge3/example_dags/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/edge3/example_dags/integration_test.py,sha256=pXf6rmS4crHKfCBT73LvPmhJS8Rmri8pUlfIvuLja7E,5240
+airflow/providers/edge3/example_dags/integration_test.py,sha256=kevxwDePjTRqLWJGV-nKEQQWb7qb-y2oedm3mdqi6-8,6127
 airflow/providers/edge3/example_dags/win_notepad.py,sha256=2evbqiupi5Ko4tyRkIEC5TPbc2c0wyR2YpsDYsBLMMM,2828
 airflow/providers/edge3/example_dags/win_test.py,sha256=GegWqjvbsSdbsA_f3S9_FRYftVO0pggXwQQggB9Vvz4,13220
 airflow/providers/edge3/executors/__init__.py,sha256=y830gGSKCvjOcLwLuCDp84NCrHWWB9RSSH1qvJpFhyY,923
-airflow/providers/edge3/executors/edge_executor.py,sha256=d05OpRKd82mSkKeybNZPXw1yiRdZAKYLZGh15eIgaiM,15254
+airflow/providers/edge3/executors/edge_executor.py,sha256=Z7q1Ph_SvAuT1cR4YYcEmKYt4Yi711q65VrzofQ5t1A,15791
 airflow/providers/edge3/models/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/edge3/models/edge_job.py,sha256=rdl9cH1bBrc7id8zkZ7uxsCJNPLG-8o9cnspGwBPBcQ,3167
 airflow/providers/edge3/models/edge_logs.py,sha256=bNstp7gR54O2vbxzz4NTL0erbifFbGUjZ-YOM0I4sqk,2768
-airflow/providers/edge3/models/edge_worker.py,sha256=MsWzvJRTdvCRLZAMDy-nYYkMOl6Ji6AUeNfSQOeO6do,8587
+airflow/providers/edge3/models/edge_worker.py,sha256=7U74k24xw36X8gQkrac3yWbtFBHbopvmRIWsjpFY8Og,10693
 airflow/providers/edge3/openapi/__init__.py,sha256=0O-WvmDx8GeKSoECpHYrbe0hW-LgjlKny3VqTCpBQeQ,927
 airflow/providers/edge3/openapi/edge_worker_api_v1.yaml,sha256=GAE2IdOXmcUueNy5KFkLBgNpoWnOjnHT9TrW5NZEWpI,24938
 airflow/providers/edge3/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
@@ -24,7 +24,7 @@ airflow/providers/edge3/plugins/templates/edge_worker_hosts.html,sha256=0_P2yfZw
 airflow/providers/edge3/plugins/templates/edge_worker_jobs.html,sha256=bZ-6ysmIy6j4eR_TPHiqbgb3qpNMKCcEEB-SpxuxNgc,2831
 airflow/providers/edge3/worker_api/__init__.py,sha256=nnPvxWGTEKZ9YyB1Yd7P9IvDOenK01LVHm22Owwxj3g,839
 airflow/providers/edge3/worker_api/app.py,sha256=Dda2VjkzgBtbQbSWSVEAoqd22RlqvBMyiPau65uKkv4,2006
-airflow/providers/edge3/worker_api/auth.py,sha256=4UFKc6wQPt9AqoMXMxt1rXljEy8xofgaIhzIdLR54gc,4936
+airflow/providers/edge3/worker_api/auth.py,sha256=XVTfL-c0JYUhpVkKdqqaxZACZXqvnPp_3W6q2TK0Bjc,4883
 airflow/providers/edge3/worker_api/datamodels.py,sha256=FAiXqnrSN8zH4YE2fUMjXfXcH9cHlhRh4uZvvr936Ys,6696
 airflow/providers/edge3/worker_api/routes/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/edge3/worker_api/routes/_v2_compat.py,sha256=Q4b2Io0yoK5V_hbgk6fiFviTeT6CFbFMOGYRZgLEeR4,4543
@@ -32,8 +32,8 @@ airflow/providers/edge3/worker_api/routes/_v2_routes.py,sha256=-WAofvXJpOYpTyh98
 airflow/providers/edge3/worker_api/routes/health.py,sha256=XxqIppnRA138Q6mAHCdyL2JvoeeganUiI-TXyXSPTGo,1075
 airflow/providers/edge3/worker_api/routes/jobs.py,sha256=UK1w6nXEUadOLwE9abZ4jHH4KtbvXcwaAF0EnwSa3y4,5733
 airflow/providers/edge3/worker_api/routes/logs.py,sha256=uk0SZ5hAimj3sAcq1FYCDu0AXYNeTeyjZDGBvw-986E,4945
-airflow/providers/edge3/worker_api/routes/worker.py,sha256=jb6m9eE9J1fyejf-DuEtGPzGrynyUmf5qkZxLeN-DQo,8554
-apache_airflow_providers_edge3-1.0.0rc1.dist-info/entry_points.txt,sha256=7WUIGfd3o9NvvbK5trbZxNXTgYGc6pqg74wZPigbx5o,206
-apache_airflow_providers_edge3-1.0.0rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-apache_airflow_providers_edge3-1.0.0rc1.dist-info/METADATA,sha256=rJGN0FI-zkY78-jsCoUM-Vwk8yfAigxt9FoSkReWF-w,4986
-apache_airflow_providers_edge3-1.0.0rc1.dist-info/RECORD,,
+airflow/providers/edge3/worker_api/routes/worker.py,sha256=BGARu1RZ74lW9X-ltuMYbbVXczm_MZdqHaai2MhDWtY,8969
+apache_airflow_providers_edge3-1.1.0rc1.dist-info/entry_points.txt,sha256=7WUIGfd3o9NvvbK5trbZxNXTgYGc6pqg74wZPigbx5o,206
+apache_airflow_providers_edge3-1.1.0rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_edge3-1.1.0rc1.dist-info/METADATA,sha256=3kJ3dEUkqjhppUJvNuD0I4r7MOg0pZ1kjmWcOrOARS0,5876
+apache_airflow_providers_edge3-1.1.0rc1.dist-info/RECORD,,