|
17 | 17 | # under the License. |
18 | 18 | from __future__ import annotations |
19 | 19 |
|
20 | | -import logging |
21 | 20 | from functools import cache |
22 | | -from typing import TYPE_CHECKING |
23 | | - |
24 | | -import pluggy |
25 | 21 |
|
| 22 | +from airflow._shared.listeners.listener import ListenerManager |
| 23 | +from airflow._shared.listeners.spec import lifecycle, taskinstance |
| 24 | +from airflow.listeners.spec import asset, dagrun, importerrors |
26 | 25 | from airflow.plugins_manager import integrate_listener_plugins |
27 | 26 |
|
28 | | -if TYPE_CHECKING: |
29 | | - from pluggy._hooks import _HookRelay |
30 | | - |
31 | | -log = logging.getLogger(__name__) |
32 | | - |
33 | | - |
34 | | -def _before_hookcall(hook_name, hook_impls, kwargs): |
35 | | - log.debug("Calling %r with %r", hook_name, kwargs) |
36 | | - log.debug("Hook impls: %s", hook_impls) |
37 | | - |
38 | | - |
39 | | -def _after_hookcall(outcome, hook_name, hook_impls, kwargs): |
40 | | - log.debug("Result from %r: %s", hook_name, outcome.get_result()) |
41 | | - |
42 | | - |
43 | | -class ListenerManager: |
44 | | - """Manage listener registration and provides hook property for calling them.""" |
45 | | - |
46 | | - def __init__(self): |
47 | | - from airflow.listeners.spec import ( |
48 | | - asset, |
49 | | - dagrun, |
50 | | - importerrors, |
51 | | - lifecycle, |
52 | | - taskinstance, |
53 | | - ) |
54 | | - |
55 | | - self.pm = pluggy.PluginManager("airflow") |
56 | | - self.pm.add_hookcall_monitoring(_before_hookcall, _after_hookcall) |
57 | | - self.pm.add_hookspecs(lifecycle) |
58 | | - self.pm.add_hookspecs(dagrun) |
59 | | - self.pm.add_hookspecs(asset) |
60 | | - self.pm.add_hookspecs(taskinstance) |
61 | | - self.pm.add_hookspecs(importerrors) |
62 | | - |
63 | | - @property |
64 | | - def has_listeners(self) -> bool: |
65 | | - return bool(self.pm.get_plugins()) |
66 | | - |
67 | | - @property |
68 | | - def hook(self) -> _HookRelay: |
69 | | - """Return hook, on which plugin methods specified in spec can be called.""" |
70 | | - return self.pm.hook |
71 | | - |
72 | | - def add_listener(self, listener): |
73 | | - if self.pm.is_registered(listener): |
74 | | - return |
75 | | - self.pm.register(listener) |
76 | | - |
77 | | - def clear(self): |
78 | | - """Remove registered plugins.""" |
79 | | - for plugin in self.pm.get_plugins(): |
80 | | - self.pm.unregister(plugin) |
81 | | - |
82 | 27 |
|
83 | 28 | @cache |
84 | 29 | def get_listener_manager() -> ListenerManager: |
85 | | - """Get singleton listener manager.""" |
| 30 | + """ |
| 31 | + Get a listener manager for Airflow core. |
| 32 | +
|
| 33 | + Registers the following listeners: |
| 34 | + - lifecycle: on_starting, before_stopping |
| 35 | + - dagrun: on_dag_run_running, on_dag_run_success, on_dag_run_failed |
| 36 | + - taskinstance: on_task_instance_running, on_task_instance_success, etc. |
| 37 | + - asset: on_asset_created, on_asset_changed, etc. |
| 38 | + - importerrors: on_new_dag_import_error, on_existing_dag_import_error |
| 39 | + """ |
86 | 40 | _listener_manager = ListenerManager() |
| 41 | + |
| 42 | + _listener_manager.add_hookspecs(lifecycle) |
| 43 | + _listener_manager.add_hookspecs(dagrun) |
| 44 | + _listener_manager.add_hookspecs(taskinstance) |
| 45 | + _listener_manager.add_hookspecs(asset) |
| 46 | + _listener_manager.add_hookspecs(importerrors) |
| 47 | + |
87 | 48 | integrate_listener_plugins(_listener_manager) |
88 | 49 | return _listener_manager |
| 50 | + |
| 51 | + |
| 52 | +__all__ = ["get_listener_manager", "ListenerManager"] |
0 commit comments