refactor(cdk): remove MemoryLimitExceeded subclass, raise AirbyteTracedException directly

Co-Authored-By: bot_apk <apk@cognition.ai>

Author: devin-ai-integration[bot]

airbyte_cdk/utils/memory_monitor.py

@@ -29,12 +29,6 @@
 _DEFAULT_CHECK_INTERVAL = 1000
 
 
-class MemoryLimitExceeded(AirbyteTracedException):
-    """Raised when connector memory usage exceeds critical threshold."""
-
-    pass
-
-
 class MemoryMonitor:
     """Monitors container memory usage via cgroup files and emits warnings before OOM kills.
 
@@ -119,8 +113,8 @@ def check_memory_usage(self) -> None:
         messages (default 1000) to minimise I/O overhead.
 
         At the warning threshold (default 85%), logs a warning message.
-        At the critical threshold (default 95%), raises MemoryLimitExceeded to
-        trigger a graceful shutdown with an actionable error message.
+        At the critical threshold (default 95%), raises ``AirbyteTracedException``
+        to trigger a graceful shutdown with an actionable error message.
 
         Each threshold triggers at most once per sync to avoid log spam.
         This method is a no-op if cgroup files are unavailable.
@@ -145,7 +139,7 @@ def check_memory_usage(self) -> None:
 
         if usage_ratio >= self._critical_threshold and not self._critical_raised:
             self._critical_raised = True
-            raise MemoryLimitExceeded(
+            raise AirbyteTracedException(
                 internal_message=f"Memory usage is {usage_percent}% ({usage_gb:.2f} / {limit_gb:.2f} GB). "
                 f"Critical threshold is {int(self._critical_threshold * 100)}%.",
                 message=f"Source exceeded memory limit ({usage_percent}% used) and must shut down to avoid an out-of-memory crash.",

unit_tests/test_entrypoint.py

@@ -836,10 +836,10 @@ def test_handle_record_counts(
 
 
 def test_memory_limit_exceeded_flushes_queued_messages(mocker, spec_mock, config_mock):
-    """When MemoryLimitExceeded is raised mid-read, queued messages should still be flushed.
+    """When AirbyteTracedException is raised mid-read, queued messages should still be flushed.
 
     The read() try/finally ensures _emit_queued_messages runs even when
-    MemoryLimitExceeded propagates.  The exception still surfaces to the
+    AirbyteTracedException propagates.  The exception still surfaces to the
     caller, but all messages yielded before (records) and during (finally-
     block state messages) the exception are available to the consumer.
     """
@@ -879,15 +879,15 @@ def test_memory_limit_exceeded_flushes_queued_messages(mocker, spec_mock, config
     mocker.patch.object(MockSource, "read_catalog", return_value={})
     mocker.patch.object(MockSource, "read", return_value=[record, record])
 
-    from airbyte_cdk.utils.memory_monitor import MemoryLimitExceeded
+    from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
     call_count = 0
 
     def _raise_on_second_call() -> None:
         nonlocal call_count
         call_count += 1
         if call_count >= 2:
-            raise MemoryLimitExceeded(
+            raise AirbyteTracedException(
                 internal_message="Memory at 96%",
                 message="Source exceeded memory limit (96% used) and must shut down to avoid an out-of-memory crash.",
                 failure_type=FailureType.system_error,
@@ -901,22 +901,22 @@ def _raise_on_second_call() -> None:
         command="read", config="config_path", state="statepath", catalog="catalogpath"
     )
 
-    # The generator yields messages until MemoryLimitExceeded propagates.
+    # The generator yields messages until AirbyteTracedException propagates.
     # Collect everything yielded before the exception surfaces.
     messages: list[str] = []
-    with pytest.raises(MemoryLimitExceeded):
+    with pytest.raises(AirbyteTracedException):
         for msg in entrypoint.run(parsed_args):
             messages.append(msg)
 
     # 1. Both records were yielded before the exception — the memory check
     #    runs after yield so every message pulled from the source is emitted.
     record_messages = [m for m in messages if "RECORD" in m]
-    assert len(record_messages) == 2, "Both records should be yielded before MemoryLimitExceeded"
+    assert len(record_messages) == 2, "Both records should be yielded before AirbyteTracedException"
 
     # 2. The queued state message was flushed by the finally block
     state_messages = [m for m in messages if "STATE" in m]
     assert len(state_messages) >= 1, (
-        "Queued state message should be flushed even after MemoryLimitExceeded"
+        "Queued state message should be flushed even after AirbyteTracedException"
     )
 
     # 3. The flushed state has sourceStats.recordCount set by handle_record_counts.

unit_tests/utils/test_memory_monitor.py

@@ -14,9 +14,9 @@
     _CGROUP_V1_USAGE,
     _CGROUP_V2_CURRENT,
     _CGROUP_V2_MAX,
-    MemoryLimitExceeded,
     MemoryMonitor,
 )
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 _MOCK_USAGE_BELOW = "500000000\n"  # 50% of 1 GB
 _MOCK_USAGE_WARNING = "870000000\n"  # 87% of 1 GB
@@ -160,27 +160,27 @@ def test_custom_thresholds_warning(caplog: pytest.LogCaptureFixture) -> None:
 
 
 def test_critical_at_95_percent_raises() -> None:
-    """MemoryLimitExceeded should be raised at 96% usage."""
+    """AirbyteTracedException should be raised at 96% usage."""
     monitor = MemoryMonitor(check_interval=1)
     with (
         patch.object(Path, "exists", _v2_exists),
         patch.object(Path, "read_text", _v2_mock_read(usage=_MOCK_USAGE_CRITICAL)),
     ):
-        with pytest.raises(MemoryLimitExceeded) as exc_info:
+        with pytest.raises(AirbyteTracedException) as exc_info:
             monitor.check_memory_usage()
 
     assert exc_info.value.failure_type == FailureType.system_error
     assert "96%" in (exc_info.value.message or "")
 
 
 def test_critical_raised_only_once() -> None:
-    """MemoryLimitExceeded should only be raised once."""
+    """AirbyteTracedException should only be raised once."""
     monitor = MemoryMonitor(check_interval=1)
     with (
         patch.object(Path, "exists", _v2_exists),
         patch.object(Path, "read_text", _v2_mock_read(usage=_MOCK_USAGE_CRITICAL)),
     ):
-        with pytest.raises(MemoryLimitExceeded):
+        with pytest.raises(AirbyteTracedException):
             monitor.check_memory_usage()
         # Second call should NOT raise again
         monitor.check_memory_usage()
@@ -197,7 +197,7 @@ def test_custom_thresholds_critical() -> None:
         patch.object(Path, "exists", _v2_exists),
         patch.object(Path, "read_text", _v2_mock_read(usage="850000000\n")),
     ):
-        with pytest.raises(MemoryLimitExceeded):
+        with pytest.raises(AirbyteTracedException):
             monitor.check_memory_usage()
 
 
@@ -293,32 +293,3 @@ def mock_read_text(self: Path) -> str:
     ):
         monitor.check_memory_usage()
     assert not caplog.records
-
-
-# ---------------------------------------------------------------------------
-# MemoryLimitExceeded exception
-# ---------------------------------------------------------------------------
-
-
-def test_memory_limit_exceeded_is_airbyte_traced_exception() -> None:
-    """MemoryLimitExceeded should be a subclass of AirbyteTracedException."""
-    from airbyte_cdk.utils.traced_exception import AirbyteTracedException
-
-    exc = MemoryLimitExceeded(
-        internal_message="test",
-        message="test message",
-        failure_type=FailureType.system_error,
-    )
-    assert isinstance(exc, AirbyteTracedException)
-
-
-def test_memory_limit_exceeded_attributes() -> None:
-    """MemoryLimitExceeded should have correct attributes."""
-    exc = MemoryLimitExceeded(
-        internal_message="Memory at 96%",
-        message="Source exceeded memory limit.",
-        failure_type=FailureType.system_error,
-    )
-    assert exc.failure_type == FailureType.system_error
-    assert exc.message == "Source exceeded memory limit."
-    assert exc.internal_message == "Memory at 96%"