feat(taps,targets): Interruption and termination signals are handled in taps and targets

singer_sdk/plugin_base.py

@@ -5,13 +5,15 @@
 import abc
 import logging
 import os
+import signal
 import sys
+import threading
 import time
 import typing as t
 import warnings
 from importlib import metadata
 from pathlib import Path, PurePath
-from types import MappingProxyType
+from types import FrameType, MappingProxyType
 
 import click
 
@@ -217,6 +219,16 @@ def __init__(
         # Initialization timestamp
         self.__initialized_at = int(time.time() * 1000)
 
+        # Signal handling
+        self._setup_signal_handlers()
+
+    def _setup_signal_handlers(self) -> None:  # pragma: no cover
+        if threading.current_thread() == threading.main_thread():
+            if hasattr(signal, "SIGINT"):
+                signal.signal(signal.SIGINT, self._handle_termination)
+            if hasattr(signal, "SIGTERM"):
+                signal.signal(signal.SIGTERM, self._handle_termination)
+
     def setup_mapper(self) -> None:
         """Initialize the plugin mapper for this tap."""
         self._mapper = PluginMapper(
@@ -445,6 +457,20 @@ def _validate_config(self, *, raise_errors: bool = True) -> list[str]:
 
         return errors
 
+    def _handle_termination(  # pragma: no cover
+        self,
+        signum: int,  # noqa: ARG002
+        frame: FrameType | None,  # noqa: ARG002
+    ) -> None:
+        """Handle termination signal.
+
+        Args:
+            signum: Signal number.
+            frame: Frame.
+        """
+        self.logger.info("Gracefully shutting down...")
+        sys.exit(0)
+
     @classmethod
     def print_version(
         cls: type[PluginBase],

singer_sdk/tap_base.py

@@ -35,6 +35,7 @@
 
 if t.TYPE_CHECKING:
     from pathlib import PurePath
+    from types import FrameType
 
     from singer_sdk.connectors import SQLConnector
     from singer_sdk.mapper import PluginMapper
@@ -498,6 +499,24 @@ def sync_all(self) -> None:
 
     # Command Line Execution
 
+    def _handle_termination(  # pragma: no cover
+        self,
+        signum: int,
+        frame: FrameType | None,
+    ) -> None:
+        """Handle termination signal.
+
+        Args:
+            signum: Signal number.
+            frame: Frame.
+        """
+        # Emit a final state message to ensure the state is written to the output
+        # even if the process is terminated by a signal.
+        try:
+            self.write_message(StateMessage(value=self.state))
+        finally:
+            super()._handle_termination(signum, frame)
+
     @classmethod
     def invoke(  # type: ignore[override]
         cls: type[Tap],

singer_sdk/target_base.py

@@ -34,6 +34,7 @@
 if t.TYPE_CHECKING:
     from collections.abc import Iterable
     from pathlib import PurePath
+    from types import FrameType
 
     from singer_sdk.connectors import SQLConnector
     from singer_sdk.mapper import PluginMapper
@@ -537,6 +538,26 @@ def _write_state_message(self, state: dict) -> None:
 
     # CLI handler
 
+    def _handle_termination(  # pragma: no cover
+        self,
+        signum: int,
+        frame: FrameType | None,
+    ) -> None:
+        """Handle termination signals.
+
+        Args:
+            signum: Signal number.
+            frame: Frame object.
+        """
+        self.logger.info(
+            "Received termination signal %d, draining all sinks...",
+            signum,
+        )
+        try:
+            self.drain_all(is_endofpipe=True)
+        finally:
+            super()._handle_termination(signum, frame)
+
     @classmethod
     def invoke(  # type: ignore[override]
         cls: type[Target],