airbyte-cdk 0.68.4__py3-none-any.whl → 0.69.1__py3-none-any.whl

unit_tests/sources/test_concurrent_source.py

@@ -12,7 +12,7 @@ from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPool
 from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import StreamAvailability, StreamAvailable, StreamUnavailable
-from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, NoopCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, FinalStateCursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_protocol.models import AirbyteStream
@@ -42,10 +42,11 @@ MESSAGE_FROM_REPOSITORY = Mock()
 
 
 class _MockStream(AbstractStream):
-    def __init__(self, name: str, available: bool = True, json_schema: Dict[str, Any] = {}):
+    def __init__(self, name: str, message_repository: MessageRepository, available: bool = True, json_schema: Dict[str, Any] = {}):
         self._name = name
         self._available = available
         self._json_schema = json_schema
+        self._message_repository = message_repository
 
     def generate_partitions(self) -> Iterable[Partition]:
         yield _MockPartition(self._name)
@@ -75,7 +76,7 @@ class _MockStream(AbstractStream):
 
     @property
     def cursor(self) -> Cursor:
-        return NoopCursor()
+        return FinalStateCursor(stream_name=self._name, stream_namespace=None, message_repository=self._message_repository)
 
 
 class _MockPartition(Partition):
@@ -103,8 +104,9 @@ class _MockPartition(Partition):
 
 
 def test_concurrent_source_reading_from_no_streams():
-    stream = _MockStream("my_stream", False, {})
-    source = _MockSource()
+    message_repository = InMemoryMessageRepository()
+    stream = _MockStream("my_stream", message_repository,False, {})
+    source = _MockSource(message_repository=message_repository)
     messages = []
     for m in source.read([stream]):
         messages.append(m)

unit_tests/sources/test_source_read.py

@@ -27,7 +27,7 @@ from airbyte_cdk.sources.concurrent_source.concurrent_source_adapter import Conc
 from airbyte_cdk.sources.message import InMemoryMessageRepository
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
-from airbyte_cdk.sources.streams.concurrent.cursor import NoopCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.utils import AirbyteTracedException
 from unit_tests.sources.streams.concurrent.scenarios.thread_based_concurrent_stream_source_builder import NeverLogSliceLogger
@@ -409,9 +409,8 @@ def _init_sources(stream_slice_to_partitions, state, logger):
 
 
 def _init_source(stream_slice_to_partitions, state, logger, source):
-    cursor = NoopCursor()
     streams = [
-        StreamFacade.create_from_stream(_MockStream(stream_slices, f"stream{i}"), source, logger, state, cursor)
+        StreamFacade.create_from_stream(_MockStream(stream_slices, f"stream{i}"), source, logger, state, FinalStateCursor(stream_name=f"stream{i}", stream_namespace=None, message_repository=InMemoryMessageRepository()))
         for i, stream_slices in enumerate(stream_slice_to_partitions)
     ]
     source.set_streams(streams)

unit_tests/test/test_entrypoint_wrapper.py

@@ -3,7 +3,7 @@
 import json
 import logging
 import os
-from typing import Any, Iterator, List
+from typing import Any, Iterator, List, Mapping
 from unittest import TestCase
 from unittest.mock import Mock, patch
 
@@ -16,7 +16,9 @@ from airbyte_protocol.models import (
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteRecordMessage,
+    AirbyteStateBlob,
     AirbyteStateMessage,
+    AirbyteStreamState,
     AirbyteStreamStatus,
     AirbyteStreamStatusTraceMessage,
     AirbyteTraceMessage,
@@ -28,8 +30,8 @@ from airbyte_protocol.models import (
 )
 
 
-def _a_state_message(state: Any) -> AirbyteMessage:
-    return AirbyteMessage(type=Type.STATE, state=AirbyteStateMessage(data=state))
+def _a_state_message(stream_name: str, stream_state: Mapping[str, Any]) -> AirbyteMessage:
+    return AirbyteMessage(type=Type.STATE, state=AirbyteStateMessage(stream=AirbyteStreamState(stream_descriptor=StreamDescriptor(name=stream_name), stream_state=AirbyteStateBlob(**stream_state))))
 
 
 def _a_status_message(stream_name: str, status: AirbyteStreamStatus) -> AirbyteMessage:
@@ -49,7 +51,7 @@ def _a_status_message(stream_name: str, status: AirbyteStreamStatus) -> AirbyteM
 _A_RECORD = AirbyteMessage(
     type=Type.RECORD, record=AirbyteRecordMessage(stream="stream", data={"record key": "record value"}, emitted_at=0)
 )
-_A_STATE_MESSAGE = _a_state_message({"state key": "state value for _A_STATE_MESSAGE"})
+_A_STATE_MESSAGE = _a_state_message("stream_name", {"state key": "state value for _A_STATE_MESSAGE"})
 _A_LOG = AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is an Airbyte log message"))
 _AN_ERROR_MESSAGE = AirbyteMessage(
     type=Type.TRACE,
@@ -176,8 +178,9 @@ class EntrypointWrapperTest(TestCase):
 
     @patch("airbyte_cdk.test.entrypoint_wrapper.AirbyteEntrypoint")
     def test_given_many_state_messages_and_records_when_read_then_output_has_records_and_state_message(self, entrypoint):
-        last_emitted_state = {"last state key": "last state value"}
-        entrypoint.return_value.run.return_value = _to_entrypoint_output([_A_STATE_MESSAGE, _a_state_message(last_emitted_state)])
+        state_value = {"state_key": "last state value"}
+        last_emitted_state = AirbyteStreamState(stream_descriptor=StreamDescriptor(name="stream_name"), stream_state=AirbyteStateBlob(**state_value))
+        entrypoint.return_value.run.return_value = _to_entrypoint_output([_A_STATE_MESSAGE, _a_state_message("stream_name", state_value)])
 
         output = read(self._a_source, _A_CONFIG, _A_CATALOG, _A_STATE)
 

unit_tests/utils/test_message_utils.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import pytest
+from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
+from airbyte_cdk.utils.message_utils import get_stream_descriptor
+from airbyte_protocol.models import (
+    AirbyteControlConnectorConfigMessage,
+    AirbyteControlMessage,
+    AirbyteMessage,
+    AirbyteRecordMessage,
+    AirbyteStateBlob,
+    AirbyteStateMessage,
+    AirbyteStateStats,
+    AirbyteStateType,
+    AirbyteStreamState,
+    OrchestratorType,
+    StreamDescriptor,
+    Type,
+)
+
+
+def test_get_record_message_stream_descriptor():
+    message = AirbyteMessage(
+        type=Type.RECORD,
+        record=AirbyteRecordMessage(
+            stream="test_stream",
+            namespace="test_namespace",
+            data={"id": "12345"},
+            emitted_at=1,
+        ),
+    )
+    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace="test_namespace")
+    assert get_stream_descriptor(message) == expected_descriptor
+
+
+def test_get_record_message_stream_descriptor_no_namespace():
+    message = AirbyteMessage(
+        type=Type.RECORD,
+        record=AirbyteRecordMessage(
+            stream="test_stream", data={"id": "12345"}, emitted_at=1
+        ),
+    )
+    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace=None)
+    assert get_stream_descriptor(message) == expected_descriptor
+
+
+def test_get_state_message_stream_descriptor():
+    message = AirbyteMessage(
+        type=Type.STATE,
+        state=AirbyteStateMessage(
+            type=AirbyteStateType.STREAM,
+            stream=AirbyteStreamState(
+                stream_descriptor=StreamDescriptor(
+                    name="test_stream", namespace="test_namespace"
+                ),
+                stream_state=AirbyteStateBlob(updated_at="2024-02-02"),
+            ),
+            sourceStats=AirbyteStateStats(recordCount=27.0),
+        ),
+    )
+    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace="test_namespace")
+    assert get_stream_descriptor(message) == expected_descriptor
+
+
+def test_get_state_message_stream_descriptor_no_namespace():
+    message = AirbyteMessage(
+        type=Type.STATE,
+        state=AirbyteStateMessage(
+            type=AirbyteStateType.STREAM,
+            stream=AirbyteStreamState(
+                stream_descriptor=StreamDescriptor(name="test_stream"),
+                stream_state=AirbyteStateBlob(updated_at="2024-02-02"),
+            ),
+            sourceStats=AirbyteStateStats(recordCount=27.0),
+        ),
+    )
+    expected_descriptor = HashableStreamDescriptor(name="test_stream", namespace=None)
+    assert get_stream_descriptor(message) == expected_descriptor
+
+
+def test_get_other_message_stream_descriptor_fails():
+    message = AirbyteMessage(
+        type=Type.CONTROL,
+        control=AirbyteControlMessage(
+            type=OrchestratorType.CONNECTOR_CONFIG,
+            emitted_at=10,
+            connectorConfig=AirbyteControlConnectorConfigMessage(config={"any config": "a config value"}),
+        ),
+    )
+    with pytest.raises(NotImplementedError):
+        get_stream_descriptor(message)