jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2067321728
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,87 @@ 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: str(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
+
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.",
args.edge_hostname, getuser())
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def remove_remote_worker(args) -> None:
+ """Remove remote edge worker entry from db."""
+ 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 remove_worker
+
+ remove_worker(args.edge_hostname)
Review Comment:
The EdgeModel should prevent removing the worker if it is not offline /
maintenance offline.
Reason: When the worker is active and the record is deleted, then the
(online) worker fails reporting heartbeat as the DB entry is missing.
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +623,110 @@ def stop(args):
logger.info("Worker has been shut down.")
+@providers_configuration_loaded
+@provide_session
+def _fetch_edge_hosts_from_db(session: Session = NEW_SESSION) -> list:
+ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
+
+ return
session.scalars(select(EdgeWorkerModel).order_by(EdgeWorkerModel.worker_name)).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(session)
+ # Format and print worker info on the screen
+ fields = [
+ "worker_name",
+ "state",
+ "queues",
+ "jobs_active",
+ "jobs_success",
+ "jobs_failed",
+ "jobs_taken",
+ "maintenance_comment",
+ ]
+ all_hosts = [{f: str(host.__getattribute__(f)) for f in fields} for host
in all_hosts_iter]
+ # Filter hosts by state
+ if args.state:
+ all_hosts = [host for host in all_hosts if host["state"] in args.state]
+ AirflowConsole().print_as(data=all_hosts, output=args.output)
+
+
+@providers_configuration_loaded
+@provide_session
+def _valid_host(hostname, session: Session = NEW_SESSION) -> bool:
+ """
+ Query the database to check if provided hostname exists
+ """
+ all_hosts_iter = _fetch_edge_hosts_from_db(session)
+ for host in all_hosts_iter:
+ if hostname == host.worker_name:
+ return True
+ return False
Review Comment:
Would be great - else you can easily generate inconsistencies.
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,87 @@ 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",
Review Comment:
Queues are stored as string in the DB, to make it proper the strin must be
re-interpreted, else it is printing a string of the queue list:
Maybe this can be generalized because in the UI the same is made. So would
be best to add this to EdgeModel class and consume it in CLI only?
##########
providers/edge3/pyproject.toml:
##########
@@ -61,6 +61,8 @@ dependencies = [
"apache-airflow-providers-fab>=1.5.3",
"pydantic>=2.11.0",
"retryhttp>=1.2.0,!=1.3.0",
+ "psycopg2-binary>=2.9.10",
+ "asyncpg>=0.30.0"
Review Comment:
I would remove them, I checked also during istall the both packages are not
re-installed - as they are installed already when airflow-core is installed.
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,87 @@ 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: str(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
+
+ change_maintenance_comment(args.edge_hostname, args.comments)
+ logger.info("Maintenance comments updated for %s by %s.",
args.edge_hostname, getuser())
Review Comment:
Maintenance comment can also be set if not in maintenance. This might also
need to be checked with the status model.
##########
providers/edge3/src/airflow/providers/edge3/cli/edge_command.py:
##########
@@ -619,6 +625,87 @@ 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: str(host.__getattribute__(f)) for f in fields} for host
in all_hosts_iter]
+ AirflowConsole().print_as(data=all_hosts, output=args.output)
Review Comment:
Outsideof this PR... maybe this "standard" printing in any format needs to
be used in the `airflow edge status` reporting as well, did not know about this
before :-D
--
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]
