jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2069486697
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(
Review Comment:
After a night of sleep I _think_ the playcement is okay-ish for the DB
activity. Just thinking (but not forcing) it might be a bit cleaner if all DB
related stuff is moced intoedge worker model and CLI just calls to retrieve -
then CLI would need to DB imports (and the risk of missing them on the client).
What do you think?
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,93 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(
+ hostname: str | None = None, states: list | None = None, session: Session
= NEW_SESSION
+) -> list:
+ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+
+ 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()
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+@provide_session
+def list_edge_workers(args, session: Session = NEW_SESSION) -> None:
+ """Query the db to list all registered edge workers."""
+ all_hosts_iter = _fetch_edge_hosts_from_db(states=args.state,
session=session)
+ # 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
+@providers_configuration_loaded
+def put_remote_worker_on_maintenance(args) -> None:
+ """Put remote edge worker on maintenance."""
+ if not _fetch_edge_hosts_from_db(hostname=args.edge_hostname):
+ raise SystemExit(f"Error: Edge Worker {args.edge_hostname} is
unknown!")
+ 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
+@providers_configuration_loaded
+def remove_remote_worker_from_maintenance(args) -> None:
+ """Remove remote edge worker from maintenance."""
+ if not _fetch_edge_hosts_from_db(hostname=args.edge_hostname):
+ raise SystemExit(f"Error: Edge Worker {args.edge_hostname} is
unknown!")
+ 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
+@providers_configuration_loaded
+def remote_worker_update_maintenance_comment(args) -> None:
+ """Update maintenance comments of the remote edge worker."""
+ if not _fetch_edge_hosts_from_db(hostname=args.edge_hostname):
+ raise SystemExit(f"Error: Edge Worker {args.edge_hostname} is
unknown!")
+ 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
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
Review Comment:
Forgot to test this - don't know if you considered this? What happens if you
call the commend on the edge where (most probably) DB is not configured? Will
it throw a crazy stack trace? Should we handle this? (not that it attempts to
setup a local sqlite in the file system)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
