pipe_callback_sink.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. #
  2. # Licensed to the Apache Software Foundation (ASF) under one
  3. # or more contributor license agreements. See the NOTICE file
  4. # distributed with this work for additional information
  5. # regarding copyright ownership. The ASF licenses this file
  6. # to you under the Apache License, Version 2.0 (the
  7. # "License"); you may not use this file except in compliance
  8. # with the License. You may obtain a copy of the License at
  9. #
  10. # http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. # Unless required by applicable law or agreed to in writing,
  13. # software distributed under the License is distributed on an
  14. # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15. # KIND, either express or implied. See the License for the
  16. # specific language governing permissions and limitations
  17. # under the License.
  18. from __future__ import annotations
  19. from typing import TYPE_CHECKING, Callable
  20. from airflow.callbacks.base_callback_sink import BaseCallbackSink
  21. if TYPE_CHECKING:
  22. from multiprocessing.connection import Connection as MultiprocessingConnection
  23. from airflow.callbacks.callback_requests import CallbackRequest
  24. class PipeCallbackSink(BaseCallbackSink):
  25. """
  26. Class for sending callbacks to DagProcessor using pipe.
  27. It is used when DagProcessor is not executed in standalone mode.
  28. """
  29. def __init__(self, get_sink_pipe: Callable[[], MultiprocessingConnection]):
  30. self._get_sink_pipe = get_sink_pipe
  31. def send(self, callback: CallbackRequest):
  32. """
  33. Send information about the callback to be executed by Pipe.
  34. :param callback: Callback request to be executed.
  35. """
  36. try:
  37. self._get_sink_pipe().send(callback)
  38. except ConnectionError:
  39. # If this died cos of an error then we will noticed and restarted
  40. # when harvest_serialized_dags calls _heartbeat_manager.
  41. pass