jscheffl commented on code in PR #49915:
URL: https://github.com/apache/airflow/pull/49915#discussion_r2070249485
##########
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:
Okay, just checked this. If you attempt to call this on the edge-side w/o db
connection configured you generate a lot of errors:
```
root@cd4307c6ba06:/opt/airflow# airflow edge list-workers
Please confirm database initialize (or wait 4 seconds to skip it). Are you
sure? [y/N]
[2025-05-01T12:48:53.345+0000] {db.py:924} INFO - Log template table does
not exist (added in 2.3.0); skipping log template sync.
Traceback (most recent call last):
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py",
line 1910, in _execute_context
self.dialect.do_execute(
File
"/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line
736, in do_execute
cursor.execute(statement, parameters)
sqlite3.OperationalError: no such table: edge_worker
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/usr/local/bin/airflow", line 10, in <module>
sys.exit(main())
^^^^^^
File "/opt/airflow/airflow-core/src/airflow/__main__.py", line 55, in main
args.func(args)
File "/opt/airflow/airflow-core/src/airflow/utils/cli.py", line 111, in
wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File
"/opt/airflow/airflow-core/src/airflow/utils/providers_configuration_loader.py",
line 55, in wrapped_function
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/airflow/airflow-core/src/airflow/utils/session.py", line 101,
in wrapper
return func(*args, session=session, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File
"/opt/airflow/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py",
line 649, in list_edge_workers
all_hosts_iter = _fetch_edge_hosts_from_db(states=args.state,
session=session)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File
"/opt/airflow/airflow-core/src/airflow/utils/providers_configuration_loader.py",
line 55, in wrapped_function
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/airflow/airflow-core/src/airflow/utils/session.py", line 99, in
wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File
"/opt/airflow/providers/edge3/src/airflow/providers/edge3/cli/edge_command.py",
line 641, in _fetch_edge_hosts_from_db
return session.scalars(query).all()
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/orm/session.py",
line 1778, in scalars
return self.execute(
^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/orm/session.py",
line 1717, in execute
result = conn._execute_20(statement, params or {}, execution_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py",
line 1710, in _execute_20
return meth(self, args_10style, kwargs_10style, execution_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/sql/elements.py",
line 334, in _execute_on_connection
return connection._execute_clauseelement(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py",
line 1577, in _execute_clauseelement
ret = self._execute_context(
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py",
line 1953, in _execute_context
self._handle_dbapi_exception(
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py",
line 2134, in _handle_dbapi_exception
util.raise_(
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/util/compat.py",
line 211, in raise_
raise exception
File "/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/base.py",
line 1910, in _execute_context
self.dialect.do_execute(
File
"/usr/local/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line
736, in do_execute
cursor.execute(statement, parameters)
sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) no such table:
edge_worker
[SQL: SELECT edge_worker.queues, edge_worker.worker_name, edge_worker.state,
edge_worker.maintenance_comment, edge_worker.first_online,
edge_worker.last_update, edge_worker.jobs_active, edge_worker.jobs_taken,
edge_worker.jobs_success, edge_worker.jobs_failed, edge_worker.sysinfo
FROM edge_worker ORDER BY edge_worker.worker_name]
(Background on this error at: https://sqlalche.me/e/14/e3q8)
```
Can you please add a small validation preventing sqlite DB creation and the
error?
Proposal as code to check whether a valid DB is configured, maybe put this
into a utility?:
```
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."
)
```
--
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]
