connector builder: Set state on stream slices (#37109)

Author: girarda

airbyte-cdk/python/airbyte_cdk/connector_builder/message_grouper.py

@@ -164,6 +164,7 @@ def _get_message_groups(
         current_slice_pages: List[StreamReadPages] = []
         current_page_request: Optional[HttpRequest] = None
         current_page_response: Optional[HttpResponse] = None
+        latest_state_message: Optional[Dict[str, Any]] = None
 
         while records_count < limit and (message := next(messages, None)):
             json_object = self._parse_json(message.log) if message.type == MessageType.LOG else None
@@ -180,7 +181,7 @@ def _get_message_groups(
                 and message.type == MessageType.LOG
                 and message.log.message.startswith(SliceLogger.SLICE_LOG_PREFIX)
             ):
-                yield StreamReadSlices(pages=current_slice_pages, slice_descriptor=current_slice_descriptor)
+                yield StreamReadSlices(pages=current_slice_pages, slice_descriptor=current_slice_descriptor, state=latest_state_message)
                 current_slice_descriptor = self._parse_slice_description(message.log.message)
                 current_slice_pages = []
                 at_least_one_page_in_group = False
@@ -222,10 +223,12 @@ def _get_message_groups(
                 datetime_format_inferrer.accumulate(message.record)
             elif message.type == MessageType.CONTROL and message.control.type == OrchestratorType.CONNECTOR_CONFIG:
                 yield message.control
+            elif message.type == MessageType.STATE:
+                latest_state_message = message.state
         else:
             if current_page_request or current_page_response or current_page_records:
                 self._close_page(current_page_request, current_page_response, current_slice_pages, current_page_records)
-                yield StreamReadSlices(pages=current_slice_pages, slice_descriptor=current_slice_descriptor)
+                yield StreamReadSlices(pages=current_slice_pages, slice_descriptor=current_slice_descriptor, state=latest_state_message)
 
     @staticmethod
     def _need_to_close_page(at_least_one_page_in_group: bool, message: AirbyteMessage, json_message: Optional[Dict[str, Any]]) -> bool:

airbyte-cdk/python/unit_tests/connector_builder/test_message_grouper.py

@@ -15,8 +15,11 @@
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteRecordMessage,
+    AirbyteStateMessage,
+    AirbyteStreamState,
     Level,
     OrchestratorType,
+    StreamDescriptor,
 )
 from airbyte_cdk.models import Type as MessageType
 from unit_tests.connector_builder.utils import create_configured_catalog
@@ -470,6 +473,7 @@ def test_get_grouped_messages_with_many_slices(mock_entrypoint_read: Mock) -> No
                 request_response_log_message(request, response, url),
                 record_message("hashiras", {"name": "Obanai Iguro"}),
                 request_response_log_message(request, response, url),
+                state_message("hashiras", {"a_timestamp": 123}),
             ]
         ),
     )
@@ -486,13 +490,16 @@ def test_get_grouped_messages_with_many_slices(mock_entrypoint_read: Mock) -> No
     assert stream_read.slices[0].slice_descriptor == {"descriptor": "first_slice"}
     assert len(stream_read.slices[0].pages) == 1
     assert len(stream_read.slices[0].pages[0].records) == 1
+    assert stream_read.slices[0].state is None
 
     assert stream_read.slices[1].slice_descriptor == {"descriptor": "second_slice"}
     assert len(stream_read.slices[1].pages) == 3
     assert len(stream_read.slices[1].pages[0].records) == 2
     assert len(stream_read.slices[1].pages[1].records) == 1
     assert len(stream_read.slices[1].pages[2].records) == 0
 
+    assert stream_read.slices[1].state.stream.stream_state == {"a_timestamp": 123}
+
 
 @patch("airbyte_cdk.connector_builder.message_grouper.AirbyteEntrypoint.read")
 def test_get_grouped_messages_given_maximum_number_of_slices_then_test_read_limit_reached(mock_entrypoint_read: Mock) -> None:
@@ -698,6 +705,13 @@ def record_message(stream: str, data: Mapping[str, Any]) -> AirbyteMessage:
     return AirbyteMessage(type=MessageType.RECORD, record=AirbyteRecordMessage(stream=stream, data=data, emitted_at=1234))
 
 
+def state_message(stream: str, data: Mapping[str, Any]) -> AirbyteMessage:
+    return AirbyteMessage(type=MessageType.STATE, state=AirbyteStateMessage(stream=AirbyteStreamState(
+        stream_descriptor=StreamDescriptor(name=stream),
+        stream_state=data
+    )))
+
+
 def slice_message(slice_descriptor: str = '{"key": "value"}') -> AirbyteMessage:
     return AirbyteMessage(type=MessageType.LOG, log=AirbyteLogMessage(level=Level.INFO, message="slice:" + slice_descriptor))