Publish stream status messages in CDK

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

@@ -13,6 +13,7 @@
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteStateMessage,
+    AirbyteStreamStatus,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
     Level,
@@ -28,6 +29,7 @@
 from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_message
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig, split_config
 from airbyte_cdk.utils.event_timing import create_timer
+from airbyte_cdk.utils.stream_status_utils import as_airbyte_message as stream_status_as_airbyte_message
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
@@ -113,17 +115,20 @@ def read(
                     continue
                 try:
                     timer.start_event(f"Syncing stream {configured_stream.stream.name}")
+                    self._emit_stream_status(configured_stream, AirbyteStreamStatus.STARTED, None)
                     yield from self._read_stream(
                         logger=logger,
                         stream_instance=stream_instance,
                         configured_stream=configured_stream,
                         state_manager=state_manager,
                         internal_config=internal_config,
                     )
+                    self._emit_stream_status(configured_stream, AirbyteStreamStatus.STOPPED, True)
                 except AirbyteTracedException as e:
                     raise e
                 except Exception as e:
                     logger.exception(f"Encountered an exception while reading stream {configured_stream.stream.name}")
+                    self._emit_stream_status(configured_stream, AirbyteStreamStatus.STOPPED, False)
                     display_message = stream_instance.get_error_display_message(e)
                     if display_message:
                         raise AirbyteTracedException.from_exception(e, message=display_message) from e
@@ -185,6 +190,9 @@ def _read_stream(
         for record in record_iterator:
             if record.type == MessageType.RECORD:
                 record_counter += 1
+            if record_counter == 1:
+                # If we just read the first record of the stream, emit the transition to the RUNNING state
+                self._emit_stream_status(configured_stream, AirbyteStreamStatus.RUNNING, None)
             yield record
 
         logger.info(f"Read {record_counter} records from {stream_name} stream")
@@ -336,3 +344,9 @@ def _get_message(self, record_data_or_message: Union[StreamData, AirbyteMessage]
             return record_data_or_message
         else:
             return stream_data_to_airbyte_message(stream.name, record_data_or_message, stream.transformer, stream.get_json_schema())
+
+    def _emit_stream_status(self, stream: ConfiguredAirbyteStream, stream_status: AirbyteStreamStatus, successful: bool):
+        """
+        Emits a new AirbyteStreamStatusTraceMessage
+        """
+        print(stream_status_as_airbyte_message(stream, AirbyteStreamStatus.RUNNING, None).json(exclude_unset=True))

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

@@ -0,0 +1,29 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+from datetime import datetime
+
+from airbyte_cdk.models import AirbyteMessage, AirbyteStreamStatusTraceMessage, AirbyteTraceMessage, StreamDescriptor, TraceType
+from airbyte_cdk.models import Type as MessageType
+
+
+def as_airbyte_message(stream, current_status, successful) -> AirbyteMessage:
+    """
+    Builds an AirbyteStreamStatusTraceMessage for the provided stream
+    """
+
+    now_millis = datetime.now().timestamp() * 1000.0
+
+    trace_message = AirbyteTraceMessage(
+        type=TraceType.STREAM_STATUS,
+        emitted_at=now_millis,
+        stream_status=AirbyteStreamStatusTraceMessage(
+            stream_descriptor=StreamDescriptor(name=stream.stream.name, namespace=stream.stream.namespace),
+            status=current_status,
+            success=successful,
+        ),
+    )
+
+    return AirbyteMessage(type=MessageType.TRACE, trace=trace_message)

airbyte-cdk/python/setup.py

@@ -46,7 +46,7 @@
     packages=find_packages(exclude=("unit_tests",)),
     package_data={"airbyte_cdk": ["py.typed", "sources/declarative/declarative_component_schema.yaml"]},
     install_requires=[
-        "airbyte-protocol-models==1.0.0",
+        "airbyte-protocol-models==0.3.5",
         "backoff",
         "dpath~=2.0.1",
         "isodate~=0.6.1",

airbyte-cdk/python/unit_tests/utils/test_stream_status_utils.py

@@ -0,0 +1,74 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteStream,
+    AirbyteStreamStatus,
+    ConfiguredAirbyteStream,
+    DestinationSyncMode,
+    SyncMode,
+    TraceType,
+)
+from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.utils.stream_status_utils import as_airbyte_message as stream_status_as_airbyte_message
+
+stream = AirbyteStream(name="name", namespace="namespace", json_schema={}, supported_sync_modes=[SyncMode.full_refresh])
+configured_stream = ConfiguredAirbyteStream(stream=stream, sync_mode=SyncMode.full_refresh, destination_sync_mode=DestinationSyncMode.overwrite)
+
+
+def test_started_as_message():
+    stream_status = AirbyteStreamStatus.STARTED
+    airbyte_message = stream_status_as_airbyte_message(configured_stream, stream_status, None)
+
+    assert type(airbyte_message) == AirbyteMessage
+    assert airbyte_message.type == MessageType.TRACE
+    assert airbyte_message.trace.type == TraceType.STREAM_STATUS
+    assert airbyte_message.trace.emitted_at > 0
+    assert airbyte_message.trace.stream_status.stream_descriptor.name == configured_stream.stream.name
+    assert airbyte_message.trace.stream_status.stream_descriptor.namespace == configured_stream.stream.namespace
+    assert airbyte_message.trace.stream_status.status == stream_status
+    assert airbyte_message.trace.stream_status.success is None
+
+
+def test_running_as_message():
+    stream_status = AirbyteStreamStatus.RUNNING
+    airbyte_message = stream_status_as_airbyte_message(configured_stream, stream_status, None)
+
+    assert type(airbyte_message) == AirbyteMessage
+    assert airbyte_message.type == MessageType.TRACE
+    assert airbyte_message.trace.type == TraceType.STREAM_STATUS
+    assert airbyte_message.trace.emitted_at > 0
+    assert airbyte_message.trace.stream_status.stream_descriptor.name == configured_stream.stream.name
+    assert airbyte_message.trace.stream_status.stream_descriptor.namespace == configured_stream.stream.namespace
+    assert airbyte_message.trace.stream_status.status == stream_status
+    assert airbyte_message.trace.stream_status.success is None
+
+
+def test_stopped_successful_as_message():
+    stream_status = AirbyteStreamStatus.STOPPED
+    airbyte_message = stream_status_as_airbyte_message(configured_stream, stream_status, True)
+
+    assert type(airbyte_message) == AirbyteMessage
+    assert airbyte_message.type == MessageType.TRACE
+    assert airbyte_message.trace.type == TraceType.STREAM_STATUS
+    assert airbyte_message.trace.emitted_at > 0
+    assert airbyte_message.trace.stream_status.stream_descriptor.name == configured_stream.stream.name
+    assert airbyte_message.trace.stream_status.stream_descriptor.namespace == configured_stream.stream.namespace
+    assert airbyte_message.trace.stream_status.status == stream_status
+    assert airbyte_message.trace.stream_status.success is True
+
+
+def test_stopped_failed_as_message():
+    stream_status = AirbyteStreamStatus.STOPPED
+    airbyte_message = stream_status_as_airbyte_message(configured_stream, stream_status, False)
+
+    assert type(airbyte_message) == AirbyteMessage
+    assert airbyte_message.type == MessageType.TRACE
+    assert airbyte_message.trace.type == TraceType.STREAM_STATUS
+    assert airbyte_message.trace.emitted_at > 0
+    assert airbyte_message.trace.stream_status.stream_descriptor.name == configured_stream.stream.name
+    assert airbyte_message.trace.stream_status.stream_descriptor.namespace == configured_stream.stream.namespace
+    assert airbyte_message.trace.stream_status.status == stream_status
+    assert airbyte_message.trace.stream_status.success is False

airbyte-config-oss/config-models-oss/src/test/java/io/airbyte/configoss/DataTypeEnumTest.java

@@ -18,7 +18,7 @@ class DataTypeEnumTest {
   @Test
   void testConversionFromJsonSchemaPrimitiveToDataType() {
     assertEquals(5, DataType.class.getEnumConstants().length);
-    assertEquals(16, JsonSchemaPrimitive.class.getEnumConstants().length);
+    assertEquals(17, JsonSchemaPrimitive.class.getEnumConstants().length);
 
     assertEquals(DataType.STRING, DataType.fromValue(JsonSchemaPrimitive.STRING.toString().toLowerCase()));
     assertEquals(DataType.NUMBER, DataType.fromValue(JsonSchemaPrimitive.NUMBER.toString().toLowerCase()));

deps.toml

@@ -1,5 +1,5 @@
 [versions]
-airbyte-protocol = "1.0.0"
+airbyte-protocol = "0.3.5"
 commons_io = "2.7"
 connectors-destination-testcontainers-clickhouse = "1.17.3"
 connectors-destination-testcontainers-elasticsearch = "1.17.3"