Emit record counts in state messages for concurrent streams (#35907)

Co-authored-by: brianjlai <[email protected]>
Co-authored-by: Brian Lai <[email protected]>

Author: erohmensing

airbyte-cdk/python/airbyte_cdk/entrypoint.py

@@ -10,23 +10,24 @@
 import socket
 import sys
 import tempfile
+from collections import defaultdict
 from functools import wraps
-from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
+from typing import Any, DefaultDict, Iterable, List, Mapping, MutableMapping, Optional, Union
 from urllib.parse import urlparse
 
 import requests
 from airbyte_cdk.connector import TConfig
 from airbyte_cdk.exception_handler import init_uncaught_exception_handler
 from airbyte_cdk.logger import init_logger
-from airbyte_cdk.models import AirbyteMessage, Status, Type
-from airbyte_cdk.models.airbyte_protocol import ConnectorSpecification  # type: ignore [attr-defined]
+from airbyte_cdk.models import AirbyteMessage, FailureType, Status, Type
+from airbyte_cdk.models.airbyte_protocol import AirbyteStateStats, ConnectorSpecification  # type: ignore [attr-defined]
 from airbyte_cdk.sources import Source
+from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
 from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit, split_config
-from airbyte_cdk.utils import is_cloud_environment
+from airbyte_cdk.utils import is_cloud_environment, message_utils
 from airbyte_cdk.utils.airbyte_secrets_utils import get_secrets, update_secrets
 from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
-from airbyte_protocol.models import FailureType
 from requests import PreparedRequest, Response, Session
 
 logger = init_logger("airbyte")
@@ -160,8 +161,27 @@ def read(
         if self.source.check_config_against_spec:
             self.validate_connection(source_spec, config)
 
-        yield from self.source.read(self.logger, config, catalog, state)
-        yield from self._emit_queued_messages(self.source)
+        stream_message_counter: DefaultDict[HashableStreamDescriptor, int] = defaultdict(int)
+        for message in self.source.read(self.logger, config, catalog, state):
+            yield self.handle_record_counts(message, stream_message_counter)
+        for message in self._emit_queued_messages(self.source):
+            yield self.handle_record_counts(message, stream_message_counter)
+
+    @staticmethod
+    def handle_record_counts(message: AirbyteMessage, stream_message_count: DefaultDict[HashableStreamDescriptor, int]) -> AirbyteMessage:
+        if message.type == Type.RECORD:
+            stream_message_count[message_utils.get_stream_descriptor(message)] += 1
+
+        elif message.type == Type.STATE:
+            stream_descriptor = message_utils.get_stream_descriptor(message)
+
+            # Set record count from the counter onto the state message
+            message.state.sourceStats = message.state.sourceStats or AirbyteStateStats()
+            message.state.sourceStats.recordCount = stream_message_count.get(stream_descriptor, 0)
+
+            # Reset the counter
+            stream_message_count[stream_descriptor] = 0
+        return message
 
     @staticmethod
     def validate_connection(source_spec: ConnectorSpecification, config: TConfig) -> None:

airbyte-cdk/python/airbyte_cdk/sources/connector_state_manager.py

@@ -82,7 +82,6 @@ def create_state_message(self, stream_name: str, namespace: Optional[str]) -> Ai
         Generates an AirbyteMessage using the current per-stream state of a specified stream in either the per-stream or legacy format
         :param stream_name: The name of the stream for the message that is being created
         :param namespace: The namespace of the stream for the message that is being created
-        :param send_per_stream_state: Decides which state format the message should be generated as
         :return: The Airbyte state message to be emitted by the connector during a sync
         """
         hashable_descriptor = HashableStreamDescriptor(name=stream_name, namespace=namespace)

airbyte-cdk/python/airbyte_cdk/utils/message_utils.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
+from airbyte_protocol.models import AirbyteMessage, Type
+
+
+def get_stream_descriptor(message: AirbyteMessage) -> HashableStreamDescriptor:
+    if message.type == Type.RECORD:
+        return HashableStreamDescriptor(name=message.record.stream, namespace=message.record.namespace)
+    elif message.type == Type.STATE:
+        if not message.state.stream or not message.state.stream.stream_descriptor:
+            raise ValueError("State message was not in per-stream state format, which is required for record counts.")
+        return HashableStreamDescriptor(
+            name=message.state.stream.stream_descriptor.name, namespace=message.state.stream.stream_descriptor.namespace
+        )
+    else:
+        raise NotImplementedError(f"get_stream_descriptor is not implemented for message type '{message.type}'.")

airbyte-cdk/python/unit_tests/sources/file_based/test_scenarios.py

@@ -16,6 +16,7 @@
 from airbyte_cdk.sources.file_based.stream.concurrent.cursor import AbstractConcurrentFileBasedCursor
 from airbyte_cdk.test.entrypoint_wrapper import EntrypointOutput
 from airbyte_cdk.test.entrypoint_wrapper import read as entrypoint_read
+from airbyte_cdk.utils import message_utils
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 from airbyte_protocol.models import AirbyteLogMessage, AirbyteMessage, ConfiguredAirbyteCatalog
 from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenario
@@ -71,7 +72,7 @@ def assert_exception(expected_exception: type[BaseException], output: Entrypoint
 
 
 def _verify_read_output(output: EntrypointOutput, scenario: TestScenario[AbstractSource]) -> None:
-    records, log_messages = output.records_and_state_messages, output.logs
+    records_and_state_messages, log_messages = output.records_and_state_messages, output.logs
     logs = [message.log for message in log_messages if message.log.level.value in scenario.log_levels]
     if scenario.expected_records is None:
         return
@@ -85,7 +86,7 @@ def _verify_read_output(output: EntrypointOutput, scenario: TestScenario[Abstrac
         ),
     )
     sorted_records = sorted(
-        filter(lambda r: r.record, records),
+        filter(lambda r: r.record, records_and_state_messages),
         key=lambda record: ",".join(
             f"{k}={v}" for k, v in sorted(record.record.data.items(), key=lambda items: (items[0], items[1])) if k != "emitted_at"
         ),
@@ -104,8 +105,9 @@ def _verify_read_output(output: EntrypointOutput, scenario: TestScenario[Abstrac
             assert actual.record.stream == expected["stream"]
 
     expected_states = list(filter(lambda e: "data" not in e, expected_records))
-    states = list(filter(lambda r: r.state, records))
+    states = list(filter(lambda r: r.state, records_and_state_messages))
     assert len(states) > 0, "No state messages emitted. Successful syncs should emit at least one stream state."
+    _verify_state_record_counts(sorted_records, states)
 
     if hasattr(scenario.source, "cursor_cls") and issubclass(scenario.source.cursor_cls, AbstractConcurrentFileBasedCursor):
         # Only check the last state emitted because we don't know the order the others will be in.
@@ -126,9 +128,33 @@ def _verify_read_output(output: EntrypointOutput, scenario: TestScenario[Abstrac
         _verify_analytics(analytics, scenario.expected_analytics)
 
 
+def _verify_state_record_counts(records: List[AirbyteMessage], states: List[AirbyteMessage]) -> None:
+    actual_record_counts = {}
+    for record in records:
+        stream_descriptor = message_utils.get_stream_descriptor(record)
+        actual_record_counts[stream_descriptor] = actual_record_counts.get(stream_descriptor, 0) + 1
+
+    state_record_count_sums = {}
+    for state_message in states:
+        stream_descriptor = message_utils.get_stream_descriptor(state_message)
+        state_record_count_sums[stream_descriptor] = (
+            state_record_count_sums.get(stream_descriptor, 0)
+            + state_message.state.sourceStats.recordCount
+        )
+
+    for stream, actual_count in actual_record_counts.items():
+        assert state_record_count_sums.get(stream) == actual_count
+
+    # We can have extra keys in state_record_count_sums if we processed a stream and reported 0 records
+    extra_keys = state_record_count_sums.keys() - actual_record_counts.keys()
+    for stream in extra_keys:
+        assert state_record_count_sums[stream] == 0
+
+
 def _verify_analytics(analytics: List[AirbyteMessage], expected_analytics: Optional[List[AirbyteAnalyticsTraceMessage]]) -> None:
     if expected_analytics:
-        assert len(analytics) == len(expected_analytics), \
+        assert len(analytics) == len(
+            expected_analytics), \
             f"Number of actual analytics messages ({len(analytics)}) did not match expected ({len(expected_analytics)})"
         for actual, expected in zip(analytics, expected_analytics):
             actual_type, actual_value = actual.trace.analytics.type, actual.trace.analytics.value

airbyte-cdk/python/unit_tests/sources/mock_server_tests/test_mock_server_abstract_source.py

@@ -205,6 +205,7 @@ def test_full_refresh_sync(self, http_mocker):
         validate_message_order([Type.RECORD, Type.RECORD, Type.STATE], actual_messages.records_and_state_messages)
         assert actual_messages.state_messages[0].state.stream.stream_descriptor.name == "users"
         assert actual_messages.state_messages[0].state.stream.stream_state == {"__ab_full_refresh_state_message": True}
+        assert actual_messages.state_messages[0].state.sourceStats.recordCount == 2
 
     @HttpMocker()
     def test_full_refresh_with_slices(self, http_mocker):
@@ -232,6 +233,7 @@ def test_full_refresh_with_slices(self, http_mocker):
         validate_message_order([Type.RECORD, Type.RECORD, Type.RECORD, Type.RECORD, Type.STATE], actual_messages.records_and_state_messages)
         assert actual_messages.state_messages[0].state.stream.stream_descriptor.name == "dividers"
         assert actual_messages.state_messages[0].state.stream.stream_state == {"__ab_full_refresh_state_message": True}
+        assert actual_messages.state_messages[0].state.sourceStats.recordCount == 4
 
 
 @freezegun.freeze_time(_NOW)
@@ -264,8 +266,10 @@ def test_incremental_sync(self, http_mocker):
         validate_message_order([Type.RECORD, Type.RECORD, Type.RECORD, Type.STATE, Type.RECORD, Type.RECORD, Type.STATE], actual_messages.records_and_state_messages)
         assert actual_messages.state_messages[0].state.stream.stream_descriptor.name == "planets"
         assert actual_messages.state_messages[0].state.stream.stream_state == {"created_at": last_record_date_0}
+        assert actual_messages.state_messages[0].state.sourceStats.recordCount == 3
         assert actual_messages.state_messages[1].state.stream.stream_descriptor.name == "planets"
         assert actual_messages.state_messages[1].state.stream.stream_state == {"created_at": last_record_date_1}
+        assert actual_messages.state_messages[1].state.sourceStats.recordCount == 2
 
     @HttpMocker()
     def test_incremental_running_as_full_refresh(self, http_mocker):
@@ -295,6 +299,7 @@ def test_incremental_running_as_full_refresh(self, http_mocker):
         validate_message_order([Type.RECORD, Type.RECORD, Type.RECORD, Type.RECORD, Type.RECORD, Type.STATE], actual_messages.records_and_state_messages)
         assert actual_messages.state_messages[0].state.stream.stream_descriptor.name == "planets"
         assert actual_messages.state_messages[0].state.stream.stream_state == {"created_at": last_record_date_1}
+        assert actual_messages.state_messages[0].state.sourceStats.recordCount == 5
 
     @HttpMocker()
     def test_legacy_incremental_sync(self, http_mocker):
@@ -324,8 +329,10 @@ def test_legacy_incremental_sync(self, http_mocker):
         validate_message_order([Type.RECORD, Type.RECORD, Type.RECORD, Type.STATE, Type.RECORD, Type.RECORD, Type.STATE], actual_messages.records_and_state_messages)
         assert actual_messages.state_messages[0].state.stream.stream_descriptor.name == "legacies"
         assert actual_messages.state_messages[0].state.stream.stream_state == {"created_at": last_record_date_0}
+        assert actual_messages.state_messages[0].state.sourceStats.recordCount == 3
         assert actual_messages.state_messages[1].state.stream.stream_descriptor.name == "legacies"
         assert actual_messages.state_messages[1].state.stream.stream_state == {"created_at": last_record_date_1}
+        assert actual_messages.state_messages[1].state.sourceStats.recordCount == 2
 
 
 @freezegun.freeze_time(_NOW)
@@ -395,12 +402,16 @@ def test_incremental_and_full_refresh_streams(self, http_mocker):
         ], actual_messages.records_and_state_messages)
         assert actual_messages.state_messages[0].state.stream.stream_descriptor.name == "users"
         assert actual_messages.state_messages[0].state.stream.stream_state == {"__ab_full_refresh_state_message": True}
+        assert actual_messages.state_messages[0].state.sourceStats.recordCount == 2
         assert actual_messages.state_messages[1].state.stream.stream_descriptor.name == "planets"
         assert actual_messages.state_messages[1].state.stream.stream_state == {"created_at": last_record_date_0}
+        assert actual_messages.state_messages[1].state.sourceStats.recordCount == 3
         assert actual_messages.state_messages[2].state.stream.stream_descriptor.name == "planets"
         assert actual_messages.state_messages[2].state.stream.stream_state == {"created_at": last_record_date_1}
+        assert actual_messages.state_messages[2].state.sourceStats.recordCount == 2
         assert actual_messages.state_messages[3].state.stream.stream_descriptor.name == "dividers"
         assert actual_messages.state_messages[3].state.stream.stream_state == {"__ab_full_refresh_state_message": True}
+        assert actual_messages.state_messages[3].state.sourceStats.recordCount == 4
 
 
 def emits_successful_sync_status_messages(status_messages: List[AirbyteStreamStatus]) -> bool:

airbyte-cdk/python/unit_tests/sources/streams/concurrent/scenarios/thread_based_concurrent_stream_scenarios.py

@@ -255,7 +255,7 @@
                     primary_key=[],
                     cursor_field=None,
                     logger=logging.getLogger("test_logger"),
-                    cursor=FinalStateCursor(stream_name="stream1", stream_namespace=None, message_repository=_message_repository),
+                    cursor=FinalStateCursor(stream_name="stream2", stream_namespace=None, message_repository=_message_repository),
                 ),
             ]
         )