This is an automated email from the ASF dual-hosted git repository.
jscheffl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5e944e49001 Add queue and remove queue cli commands for EdgeExecutor
(#53505)
5e944e49001 is described below
commit 5e944e490012073fc28aff8cb44f27c3525c171c
Author: Dheeraj Turaga <[email protected]>
AuthorDate: Sat Jul 19 03:38:04 2025 -0500
Add queue and remove queue cli commands for EdgeExecutor (#53505)
* Add queue and remove queue cli commands for EdgeExecutor
* Add queue cli commands to docs
---
providers/edge3/docs/deployment.rst | 2 +
.../airflow/providers/edge3/cli/edge_command.py | 65 ++++++++++++++++++++++
.../airflow/providers/edge3/models/edge_worker.py | 34 +++++++++++
3 files changed, 101 insertions(+)
diff --git a/providers/edge3/docs/deployment.rst
b/providers/edge3/docs/deployment.rst
index b4fc9977c1f..3b90db06810 100644
--- a/providers/edge3/docs/deployment.rst
+++ b/providers/edge3/docs/deployment.rst
@@ -190,3 +190,5 @@ instance. The commands are:
- ``airflow edge remote-edge-worker-exit-maintenance``: Request a remote edge
worker to exit maintenance mode
- ``airflow edge shutdown-remote-edge-worker``: Shuts down a remote edge
worker gracefully
- ``airflow edge remove-remote-edge-worker``: Remove a worker instance from
the cluster
+- ``airflow edge add-worker-queues``: Add queues to an edge worker
+- ``airflow edge remove-worker-queues``: Remove queues from an edge worker
diff --git a/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
b/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
index 19573d7a96a..39bdb0c6575 100644
--- a/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
+++ b/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py
@@ -351,6 +351,48 @@ def remote_worker_request_shutdown(args) -> None:
logger.info("Requested shutdown of Edge Worker host %s by %s.",
args.edge_hostname, getuser())
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def add_worker_queues(args) -> None:
+ """Add queues to an edge worker."""
+ _check_valid_db_connection()
+ _check_if_registered_edge_host(hostname=args.edge_hostname)
+ from airflow.providers.edge3.models.edge_worker import add_worker_queues
+
+ queues = args.queues.split(",") if args.queues else []
+ if not queues:
+ raise SystemExit("Error: No queues specified to add.")
+
+ try:
+ add_worker_queues(args.edge_hostname, queues)
+ logger.info("Added queues %s to Edge Worker host %s by %s.", queues,
args.edge_hostname, getuser())
+ except TypeError as e:
+ logger.error(str(e))
+ raise SystemExit
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def remove_worker_queues(args) -> None:
+ """Remove queues from an edge worker."""
+ _check_valid_db_connection()
+ _check_if_registered_edge_host(hostname=args.edge_hostname)
+ from airflow.providers.edge3.models.edge_worker import remove_worker_queues
+
+ queues = args.queues.split(",") if args.queues else []
+ if not queues:
+ raise SystemExit("Error: No queues specified to remove.")
+
+ try:
+ remove_worker_queues(args.edge_hostname, queues)
+ logger.info(
+ "Removed queues %s from Edge Worker host %s by %s.", queues,
args.edge_hostname, getuser()
+ )
+ except TypeError as e:
+ logger.error(str(e))
+ raise SystemExit
+
+
ARG_CONCURRENCY = Arg(
("-c", "--concurrency"),
type=int,
@@ -380,6 +422,11 @@ ARG_REQUIRED_MAINTENANCE_COMMENT = Arg(
help="Maintenance comments to report reason. Required if enabling
maintenance",
required=True,
)
+ARG_QUEUES_MANAGE = Arg(
+ ("-q", "--queues"),
+ help="Comma delimited list of queues to add or remove.",
+ required=True,
+)
ARG_WAIT_MAINT = Arg(
("-w", "--wait"),
default=False,
@@ -516,4 +563,22 @@ EDGE_COMMANDS: list[ActionCommand] = [
func=remote_worker_request_shutdown,
args=(ARG_REQUIRED_EDGE_HOSTNAME,),
),
+ ActionCommand(
+ name="add-worker-queues",
+ help=add_worker_queues.__doc__,
+ func=add_worker_queues,
+ args=(
+ ARG_REQUIRED_EDGE_HOSTNAME,
+ ARG_QUEUES_MANAGE,
+ ),
+ ),
+ ActionCommand(
+ name="remove-worker-queues",
+ help=remove_worker_queues.__doc__,
+ func=remove_worker_queues,
+ args=(
+ ARG_REQUIRED_EDGE_HOSTNAME,
+ ARG_QUEUES_MANAGE,
+ ),
+ ),
]
diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
index bc0c426abdb..a6076b6a6fb 100644
--- a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
+++ b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py
@@ -283,3 +283,37 @@ def request_shutdown(worker_name: str, session: Session =
NEW_SESSION) -> None:
EdgeWorkerState.UNKNOWN,
):
worker.state = EdgeWorkerState.SHUTDOWN_REQUEST
+
+
+@provide_session
+def add_worker_queues(worker_name: str, queues: list[str], session: Session =
NEW_SESSION) -> None:
+ """Add queues to an edge worker."""
+ query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name ==
worker_name)
+ worker: EdgeWorkerModel = session.scalar(query)
+ if worker.state in (
+ EdgeWorkerState.OFFLINE,
+ EdgeWorkerState.OFFLINE_MAINTENANCE,
+ EdgeWorkerState.UNKNOWN,
+ ):
+ error_message = f"Cannot add queues to edge worker {worker_name} as it
is in {worker.state} state!"
+ logger.error(error_message)
+ raise TypeError(error_message)
+ worker.add_queues(queues)
+
+
+@provide_session
+def remove_worker_queues(worker_name: str, queues: list[str], session: Session
= NEW_SESSION) -> None:
+ """Remove queues from an edge worker."""
+ query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name ==
worker_name)
+ worker: EdgeWorkerModel = session.scalar(query)
+ if worker.state in (
+ EdgeWorkerState.OFFLINE,
+ EdgeWorkerState.OFFLINE_MAINTENANCE,
+ EdgeWorkerState.UNKNOWN,
+ ):
+ error_message = (
+ f"Cannot remove queues from edge worker {worker_name} as it is in
{worker.state} state!"
+ )
+ logger.error(error_message)
+ raise TypeError(error_message)
+ worker.remove_queues(queues)
