Change fetch_one to raise RecordNotFoundException instead of returning None for 404s

Co-Authored-By: AJ Steers <[email protected]>

Author: devin-ai-integration[bot]

airbyte_cdk/sources/declarative/exceptions.py

@@ -7,3 +7,7 @@ class ReadException(Exception):
     """
     Raise when there is an error reading data from an API Source
     """
+
+
+class RecordNotFoundException(ReadException):
+    """Raised when a requested record is not found (e.g., 404 response)."""

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -26,6 +26,7 @@
 from airbyte_cdk.legacy.sources.declarative.incremental import ResumableFullRefreshCursor
 from airbyte_cdk.legacy.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
 from airbyte_cdk.models import AirbyteMessage
+from airbyte_cdk.sources.declarative.exceptions import RecordNotFoundException
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
@@ -713,14 +714,17 @@ def fetch_one(
                 log_formatter=self.log_formatter,
             )
         except Exception as e:
-            # Check if this is a 404 (record not found) - return None
-            if hasattr(e, "response") and hasattr(e.response, "status_code"):
-                if e.response.status_code == 404:
-                    return None
+            # Check if this is a 404 (record not found) - raise RecordNotFoundException
+            if "404" in str(e) or (hasattr(e, "response") and hasattr(e.response, "status_code") and e.response.status_code == 404):
+                raise RecordNotFoundException(
+                    f"Record with primary key {pk_value} not found"
+                ) from e
             raise
 
         if not response:
-            return None
+            raise RecordNotFoundException(
+                f"Record with primary key {pk_value} not found (no response)"
+            )
 
         records = list(
             self._parse_response(
@@ -732,16 +736,20 @@ def fetch_one(
             )
         )
 
-        # Return the first record if found, None otherwise
+        # Return the first record if found, raise RecordNotFoundException otherwise
         if records:
             first_record = records[0]
             if isinstance(first_record, Record):
                 return dict(first_record.data)
             elif isinstance(first_record, Mapping):
                 return dict(first_record)
             else:
-                return None
-        return None
+                raise RecordNotFoundException(
+                    f"Record with primary key {pk_value} not found (invalid record type)"
+                )
+        raise RecordNotFoundException(
+            f"Record with primary key {pk_value} not found (empty response)"
+        )
 
 
 def _deep_merge(

unit_tests/sources/declarative/retrievers/test_simple_retriever.py

@@ -24,6 +24,7 @@
 )
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import NoAuth
 from airbyte_cdk.sources.declarative.decoders import JsonDecoder
+from airbyte_cdk.sources.declarative.exceptions import RecordNotFoundException
 from airbyte_cdk.sources.declarative.extractors import DpathExtractor, HttpSelector, RecordSelector
 from airbyte_cdk.sources.declarative.partition_routers import SinglePartitionRouter
 from airbyte_cdk.sources.declarative.requesters.paginators import DefaultPaginator, Paginator
@@ -1719,7 +1720,7 @@ def test_fetch_one_composite_pk():
 
 
 def test_fetch_one_not_found():
-    """Test fetch_one returns None when record is not found (404)."""
+    """Test fetch_one raises RecordNotFoundException when record is not found (404)."""
     requester = MagicMock()
     requester.get_path.return_value = "posts"
 
@@ -1739,9 +1740,10 @@ def test_fetch_one_not_found():
         config={},
     )
 
-    result = retriever.fetch_one("999", records_schema={})
+    with pytest.raises(RecordNotFoundException) as exc_info:
+        retriever.fetch_one("999", records_schema={})
 
-    assert result is None
+    assert "999" in str(exc_info.value)
 
 
 def test_fetch_one_server_error():
@@ -1794,7 +1796,7 @@ def test_fetch_one_invalid_pk_type():
 
 
 def test_fetch_one_no_response():
-    """Test fetch_one returns None when response is None."""
+    """Test fetch_one raises RecordNotFoundException when response is None."""
     requester = MagicMock()
     requester.get_path.return_value = "posts"
     requester.send_request.return_value = None
@@ -1810,13 +1812,14 @@ def test_fetch_one_no_response():
         config={},
     )
 
-    result = retriever.fetch_one("123", records_schema={})
+    with pytest.raises(RecordNotFoundException) as exc_info:
+        retriever.fetch_one("123", records_schema={})
 
-    assert result is None
+    assert "123" in str(exc_info.value)
 
 
 def test_fetch_one_empty_records():
-    """Test fetch_one returns None when no records are returned."""
+    """Test fetch_one raises RecordNotFoundException when no records are returned."""
     requester = MagicMock()
     requester.get_path.return_value = "posts"
 
@@ -1838,6 +1841,7 @@ def test_fetch_one_empty_records():
         config={},
     )
 
-    result = retriever.fetch_one("123", records_schema={})
+    with pytest.raises(RecordNotFoundException) as exc_info:
+        retriever.fetch_one("123", records_schema={})
 
-    assert result is None
+    assert "123" in str(exc_info.value)