ashb commented on code in PR #27113:
URL: https://github.com/apache/airflow/pull/27113#discussion_r1028105968
##########
airflow/jobs/base_job.py:
##########
@@ -149,6 +151,7 @@ def is_alive(self, grace_multiplier=2.1):
@provide_session
def kill(self, session=None):
"""Handles on_kill callback and updates state in database."""
+ get_listener_manager().hook.before_stopping()
Review Comment:
We shouldn't need this here and in the `finally` block below do we?
##########
airflow/listeners/spec/lifecycle.py:
##########
@@ -0,0 +1,38 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from pluggy import HookspecMarker
+
+hookspec = HookspecMarker("airflow")
+
+
+@hookspec
+def on_starting():
+ """
+ Called when Airflow component (scheduler, worker) starts.
+ It's guaranteed this will be called before any other plugin method.
+ """
+
+
+@hookspec
+def before_stopping():
+ """
+ Called before Airflow component (scheduler, worker) stops.
Review Comment:
```suggestion
Called before Airflow component (scheduler, worker) stops.
```
##########
airflow/jobs/base_job.py:
##########
@@ -252,6 +255,7 @@ def run(self):
self.state = State.FAILED
raise
finally:
+ get_listener_manager().hook.before_stopping()
Review Comment:
```suggestion
get_listener_manager().hook.before_stopping(self)
```
##########
airflow/jobs/base_job.py:
##########
@@ -110,6 +111,7 @@ def __init__(self, executor=None, heartrate=None, *args,
**kwargs):
self.heartrate = heartrate
self.unixname = getuser()
self.max_tis_per_query: int = conf.getint("scheduler",
"max_tis_per_query")
+ get_listener_manager().hook.on_starting()
Review Comment:
```suggestion
get_listener_manager().hook.on_starting(self)
```
(Because this is also called for LocalTaskJob, so being able to tell them
apart is useful)
##########
airflow/jobs/backfill_job.py:
##########
@@ -946,3 +952,12 @@ def query(result, items):
self.log.info("Reset the following %s TaskInstances:\n\t%s",
len(reset_tis), task_instance_str)
return len(reset_tis)
+
+ def notify_dagrun_state_changed(self, dag_run: DagRun, msg: str = ""):
+ if dag_run.state == DagRunState.RUNNING:
+ get_listener_manager().hook.on_dag_run_running(dag_run=dag_run,
msg=msg)
+ elif dag_run.state == DagRunState.SUCCESS:
+ get_listener_manager().hook.on_dag_run_success(dag_run=dag_run,
msg=msg)
+ elif dag_run.state == DagRunState.FAILED:
+ get_listener_manager().hook.on_dag_run_failed(dag_run=dag_run,
msg=msg)
+ # deliberately not notifying on QUEUED
Review Comment:
This shouldn't be job specific right? Can't this live directly in DagRun
itself?
##########
airflow/jobs/backfill_job.py:
##########
@@ -946,3 +952,12 @@ def query(result, items):
self.log.info("Reset the following %s TaskInstances:\n\t%s",
len(reset_tis), task_instance_str)
return len(reset_tis)
+
+ def notify_dagrun_state_changed(self, dag_run: DagRun, msg: str = ""):
+ if dag_run.state == DagRunState.RUNNING:
+ get_listener_manager().hook.on_dag_run_running(dag_run=dag_run,
msg=msg)
+ elif dag_run.state == DagRunState.SUCCESS:
+ get_listener_manager().hook.on_dag_run_success(dag_run=dag_run,
msg=msg)
+ elif dag_run.state == DagRunState.FAILED:
+ get_listener_manager().hook.on_dag_run_failed(dag_run=dag_run,
msg=msg)
+ # deliberately not notifying on QUEUED
Review Comment:
We should have queued on the spec, and it's up to the listener to ignore it
if it wants.
Additionally I wonder if we should just have a single
`dag_run_state_changed` type event instead of one for each and every state we
might have in a dag run?
##########
airflow/listeners/spec/lifecycle.py:
##########
@@ -0,0 +1,38 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from pluggy import HookspecMarker
+
+hookspec = HookspecMarker("airflow")
+
+
+@hookspec
+def on_starting():
+ """
+ Called when Airflow component (scheduler, worker) starts.
Review Comment:
```suggestion
Called when Airflow component (scheduler, worker) starts.
```
##########
airflow/listeners/spec/lifecycle.py:
##########
@@ -0,0 +1,38 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from pluggy import HookspecMarker
+
+hookspec = HookspecMarker("airflow")
+
+
+@hookspec
+def on_starting():
+ """
+ Called when Airflow component (scheduler, worker) starts.
Review Comment:
This should probably recieve the `Job` object so the plugin/listener can
tell what kind of thing it's being run inside of
--
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]
