1
1
# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
2
2
3
3
4
- from dataclasses import InitVar , dataclass
4
+ from dataclasses import InitVar , dataclass , field
5
5
from typing import Any , Iterable , Mapping , Optional
6
6
7
7
from typing_extensions import deprecated
8
8
9
9
from airbyte_cdk .sources .declarative .async_job .job import AsyncJob
10
- from airbyte_cdk .sources .declarative .async_job .job_orchestrator import AsyncPartition
11
10
from airbyte_cdk .sources .declarative .extractors .record_selector import RecordSelector
12
11
from airbyte_cdk .sources .declarative .partition_routers .async_job_partition_router import (
13
12
AsyncJobPartitionRouter ,
16
15
from airbyte_cdk .sources .source import ExperimentalClassWarning
17
16
from airbyte_cdk .sources .streams .core import StreamData
18
17
from airbyte_cdk .sources .types import Config , StreamSlice , StreamState
18
+ from airbyte_cdk .sources .utils .slice_logger import AlwaysLogSliceLogger
19
19
20
20
21
21
@deprecated (
@@ -28,6 +28,10 @@ class AsyncRetriever(Retriever):
28
28
parameters : InitVar [Mapping [str , Any ]]
29
29
record_selector : RecordSelector
30
30
stream_slicer : AsyncJobPartitionRouter
31
+ slice_logger : AlwaysLogSliceLogger = field (
32
+ init = False ,
33
+ default_factory = lambda : AlwaysLogSliceLogger (),
34
+ )
31
35
32
36
def __post_init__ (self , parameters : Mapping [str , Any ]) -> None :
33
37
self ._parameters = parameters
@@ -74,14 +78,36 @@ def _validate_and_get_stream_slice_jobs(
74
78
"""
75
79
return stream_slice .extra_fields .get ("jobs" , []) if stream_slice else []
76
80
81
+ def _get_cursor_slice_info (
82
+ self ,
83
+ stream_slice : Optional [StreamSlice ] = None ,
84
+ ) -> Mapping [str , Any ]:
85
+ """
86
+ Retrieve the cursor slice information from a provided stream slice.
87
+
88
+ This method checks if a stream slice is provided and contains a 'cursor_slice' attribute.
89
+ If present, it returns the value of 'cursor_slice' as a dictionary. Otherwise, it returns an empty dictionary.
90
+
91
+ Args:
92
+ stream_slice (Optional[StreamSlice]): An optional stream slice object that may include a 'cursor_slice' attribute.
93
+
94
+ Returns:
95
+ Mapping[str, Any]: A dictionary containing the cursor slice information if available, otherwise an empty dictionary.
96
+ """
97
+ return stream_slice .cursor_slice if stream_slice and stream_slice .cursor_slice else {}
98
+
77
99
def stream_slices (self ) -> Iterable [Optional [StreamSlice ]]:
78
- return self .stream_slicer .stream_slices ()
100
+ yield from self .stream_slicer .stream_slices ()
79
101
80
102
def read_records (
81
103
self ,
82
104
records_schema : Mapping [str , Any ],
83
105
stream_slice : Optional [StreamSlice ] = None ,
84
106
) -> Iterable [StreamData ]:
107
+ # emit the slice_descriptor log message, for connector builder TestRead
108
+ yield self .slice_logger .create_slice_log_message (
109
+ self ._get_cursor_slice_info (stream_slice ),
110
+ )
85
111
stream_state : StreamState = self ._get_stream_state ()
86
112
jobs : Iterable [AsyncJob ] = self ._validate_and_get_stream_slice_jobs (stream_slice )
87
113
records : Iterable [Mapping [str , Any ]] = self .stream_slicer .fetch_records (jobs )
0 commit comments