airbyte_cdk.sources

 54class AbstractSource(Source, ABC):
 55    """
 56    Abstract base class for an Airbyte Source. Consumers should implement any abstract methods
 57    in this class to create an Airbyte Specification compliant Source.
 58    """
 59
 60    @abstractmethod
 61    def check_connection(
 62        self, logger: logging.Logger, config: Mapping[str, Any]
 63    ) -> Tuple[bool, Optional[Any]]:
 64        """
 65        :param logger: source logger
 66        :param config: The user-provided configuration as specified by the source's spec.
 67          This usually contains information required to check connection e.g. tokens, secrets and keys etc.
 68        :return: A tuple of (boolean, error). If boolean is true, then the connection check is successful
 69          and we can connect to the underlying data source using the provided configuration.
 70          Otherwise, the input config cannot be used to connect to the underlying data source,
 71          and the "error" object should describe what went wrong.
 72          The error object will be cast to string to display the problem to the user.
 73        """
 74
 75    @abstractmethod
 76    def streams(self, config: Mapping[str, Any]) -> List[Stream]:
 77        """
 78        :param config: The user-provided configuration as specified by the source's spec.
 79        Any stream construction related operation should happen here.
 80        :return: A list of the streams in this source connector.
 81        """
 82
 83    # Stream name to instance map for applying output object transformation
 84    _stream_to_instance_map: Dict[str, Stream] = {}
 85    _slice_logger: SliceLogger = DebugSliceLogger()
 86
 87    def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog:
 88        """Implements the Discover operation from the Airbyte Specification.
 89        See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/#discover.
 90        """
 91        streams = [stream.as_airbyte_stream() for stream in self.streams(config=config)]
 92        return AirbyteCatalog(streams=streams)
 93
 94    def check(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
 95        """Implements the Check Connection operation from the Airbyte Specification.
 96        See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/#check.
 97        """
 98        check_succeeded, error = self.check_connection(logger, config)
 99        if not check_succeeded:
100            return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error))
101        return AirbyteConnectionStatus(status=Status.SUCCEEDED)
102
103    def check_section_connection(
104        self, logger: logging.Logger, section: str, config: Mapping[str, Any]
105    ) -> Tuple[bool, Optional[Any]]:
106        """
107        `check_connection` for one section of the stepped setup form: `config` only holds the fields of `section`
108        and of the sections before it. Returns (True, None) when they are valid, (False, error) otherwise.
109
110        Override it in a source whose specification declares sections.
111        """
112        return False, CHECK_SECTION_NOT_IMPLEMENTED.format(section=section)
113
114    def check_section(
115        self, logger: logging.Logger, section: str, config: Mapping[str, Any]
116    ) -> AirbyteConnectionStatus:
117        """Same as `check`, for one section: wraps `check_section_connection` into a connection status."""
118        check_succeeded, error = self.check_section_connection(logger, section, config)
119        if not check_succeeded:
120            return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error))
121        return AirbyteConnectionStatus(status=Status.SUCCEEDED)
122
123    def read(
124        self,
125        logger: logging.Logger,
126        config: Mapping[str, Any],
127        catalog: ConfiguredAirbyteCatalog,
128        state: Optional[List[AirbyteStateMessage]] = None,
129    ) -> Iterator[AirbyteMessage]:
130        """Implements the Read operation from the Airbyte Specification. See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/."""
131        logger.info(f"Starting syncing {self.name}")
132        config, internal_config = split_config(config)
133        # TODO assert all streams exist in the connector
134        # get the streams once in case the connector needs to make any queries to generate them
135        stream_instances = {s.name: s for s in self.streams(config)}
136        state_manager = ConnectorStateManager(state=state)
137        self._stream_to_instance_map = stream_instances
138
139        stream_name_to_exception: MutableMapping[str, AirbyteTracedException] = {}
140
141        with create_timer(self.name) as timer:
142            for configured_stream in catalog.streams:
143                stream_instance = stream_instances.get(configured_stream.stream.name)
144                is_stream_exist = bool(stream_instance)
145                try:
146                    # Used direct reference to `stream_instance` instead of `is_stream_exist` to avoid mypy type checking errors
147                    if not stream_instance:
148                        if not self.raise_exception_on_missing_stream:
149                            yield stream_status_as_airbyte_message(
150                                configured_stream.stream, AirbyteStreamStatus.INCOMPLETE
151                            )
152                            continue
153
154                        error_message = (
155                            f"The stream '{configured_stream.stream.name}' in your connection configuration was not found in the source. "
156                            f"Refresh the schema in your replication settings and remove this stream from future sync attempts."
157                        )
158
159                        # Use configured_stream as stream_instance to support references in error handling.
160                        stream_instance = configured_stream.stream
161
162                        raise AirbyteTracedException(
163                            message="A stream listed in your configuration was not found in the source. Please check the logs for more "
164                            "details.",
165                            internal_message=error_message,
166                            failure_type=FailureType.config_error,
167                        )
168
169                    timer.start_event(f"Syncing stream {configured_stream.stream.name}")
170                    logger.info(f"Marking stream {configured_stream.stream.name} as STARTED")
171                    yield stream_status_as_airbyte_message(
172                        configured_stream.stream, AirbyteStreamStatus.STARTED
173                    )
174                    yield from self._read_stream(
175                        logger=logger,
176                        stream_instance=stream_instance,
177                        configured_stream=configured_stream,
178                        state_manager=state_manager,
179                        internal_config=internal_config,
180                    )
181                    logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
182                    yield stream_status_as_airbyte_message(
183                        configured_stream.stream, AirbyteStreamStatus.COMPLETE
184                    )
185
186                except Exception as e:
187                    yield from self._emit_queued_messages()
188                    logger.exception(
189                        f"Encountered an exception while reading stream {configured_stream.stream.name}"
190                    )
191                    logger.info(f"Marking stream {configured_stream.stream.name} as STOPPED")
192                    yield stream_status_as_airbyte_message(
193                        configured_stream.stream, AirbyteStreamStatus.INCOMPLETE
194                    )
195
196                    stream_descriptor = StreamDescriptor(name=configured_stream.stream.name)
197
198                    if isinstance(e, AirbyteTracedException):
199                        traced_exception = e
200                        info_message = f"Stopping sync on error from stream {configured_stream.stream.name} because {self.name} does not support continuing syncs on error."
201                    else:
202                        traced_exception = self._serialize_exception(
203                            stream_descriptor, e, stream_instance=stream_instance
204                        )
205                        info_message = f"{self.name} does not support continuing syncs on error from stream {configured_stream.stream.name}"
206
207                    yield traced_exception.as_sanitized_airbyte_message(
208                        stream_descriptor=stream_descriptor
209                    )
210                    stream_name_to_exception[stream_instance.name] = traced_exception  # type: ignore # use configured_stream if stream_instance is None
211                    if self.stop_sync_on_stream_failure:
212                        logger.info(info_message)
213                        break
214                finally:
215                    # Finish read event only if the stream instance exists;
216                    # otherwise, there's no need as it never started
217                    if is_stream_exist:
218                        timer.finish_event()
219                        logger.info(f"Finished syncing {configured_stream.stream.name}")
220                        logger.info(timer.report())
221
222        if len(stream_name_to_exception) > 0:
223            error_message = generate_failed_streams_error_message(
224                {key: [value] for key, value in stream_name_to_exception.items()}
225            )
226            logger.info(error_message)
227            # We still raise at least one exception when a stream raises an exception because the platform currently relies
228            # on a non-zero exit code to determine if a sync attempt has failed. We also raise the exception as a config_error
229            # type because this combined error isn't actionable, but rather the previously emitted individual errors.
230            raise AirbyteTracedException(
231                message=error_message, failure_type=FailureType.config_error
232            )
233        logger.info(f"Finished syncing {self.name}")
234
235    @staticmethod
236    def _serialize_exception(
237        stream_descriptor: StreamDescriptor, e: Exception, stream_instance: Optional[Stream] = None
238    ) -> AirbyteTracedException:
239        display_message = stream_instance.get_error_display_message(e) if stream_instance else None
240        if display_message:
241            return AirbyteTracedException.from_exception(
242                e, message=display_message, stream_descriptor=stream_descriptor
243            )
244        return AirbyteTracedException.from_exception(e, stream_descriptor=stream_descriptor)
245
246    @property
247    def raise_exception_on_missing_stream(self) -> bool:
248        return False
249
250    def _read_stream(
251        self,
252        logger: logging.Logger,
253        stream_instance: Stream,
254        configured_stream: ConfiguredAirbyteStream,
255        state_manager: ConnectorStateManager,
256        internal_config: InternalConfig,
257    ) -> Iterator[AirbyteMessage]:
258        if internal_config.page_size and isinstance(stream_instance, HttpStream):
259            logger.info(
260                f"Setting page size for {stream_instance.name} to {internal_config.page_size}"
261            )
262            stream_instance.page_size = internal_config.page_size
263        logger.debug(
264            f"Syncing configured stream: {configured_stream.stream.name}",
265            extra={
266                "sync_mode": configured_stream.sync_mode,
267                "primary_key": configured_stream.primary_key,
268                "cursor_field": configured_stream.cursor_field,
269            },
270        )
271        stream_instance.log_stream_sync_configuration()
272
273        stream_name = configured_stream.stream.name
274        stream_state = state_manager.get_stream_state(stream_name, stream_instance.namespace)
275
276        # This is a hack. Existing full refresh streams that are converted into resumable full refresh need to discard
277        # the state because the terminal state for a full refresh sync is not compatible with substream resumable full
278        # refresh state. This is only required when running live traffic regression testing since the platform normally
279        # handles whether to pass state
280        if stream_state == {"__ab_no_cursor_state_message": True}:
281            stream_state = {}
282
283        if "state" in dir(stream_instance):
284            stream_instance.state = stream_state  # type: ignore # we check that state in the dir(stream_instance)
285            logger.info(f"Setting state of {self.name} stream to {stream_state}")
286
287        record_iterator = stream_instance.read(
288            configured_stream,
289            logger,
290            self._slice_logger,
291            stream_state,
292            state_manager,
293            internal_config,
294        )
295
296        record_counter = 0
297        logger.info(f"Syncing stream: {stream_name} ")
298        for record_data_or_message in record_iterator:
299            record = self._get_message(record_data_or_message, stream_instance)
300            if record.type == MessageType.RECORD:
301                record_counter += 1
302                if record_counter == 1:
303                    logger.info(f"Marking stream {stream_name} as RUNNING")
304                    # If we just read the first record of the stream, emit the transition to the RUNNING state
305                    yield stream_status_as_airbyte_message(
306                        configured_stream.stream, AirbyteStreamStatus.RUNNING
307                    )
308            yield from self._emit_queued_messages()
309            yield record
310
311        logger.info(f"Read {record_counter} records from {stream_name} stream")
312
313    def _emit_queued_messages(self) -> Iterable[AirbyteMessage]:
314        if self.message_repository:
315            yield from self.message_repository.consume_queue()
316        return
317
318    def _get_message(
319        self, record_data_or_message: Union[StreamData, AirbyteMessage], stream: Stream
320    ) -> AirbyteMessage:
321        """
322        Converts the input to an AirbyteMessage if it is a StreamData. Returns the input as is if it is already an AirbyteMessage
323        """
324        match record_data_or_message:
325            case AirbyteMessage():
326                return record_data_or_message
327            case _:
328                return stream_data_to_airbyte_message(
329                    stream.name,
330                    record_data_or_message,
331                    stream.transformer,
332                    stream.get_json_schema(),
333                )
334
335    @property
336    def message_repository(self) -> Union[None, MessageRepository]:
337        return _default_message_repository
338
339    @property
340    def stop_sync_on_stream_failure(self) -> bool:
341        """
342        WARNING: This function is in-development which means it is subject to change. Use at your own risk.
343
344        By default, when a source encounters an exception while syncing a stream, it will emit an error trace message and then
345        continue syncing the next stream. This can be overwritten on a per-source basis so that the source will stop the sync
346        on the first error seen and emit a single error trace message for that stream.
347        """
348        return False

13class BaseConfig(BaseModel):
14    """Base class for connector spec, adds the following behaviour:
15
16    - resolve $ref and replace it with definition
17    - replace all occurrences of anyOf with oneOf
18    - drop description
19    """
20
21    @classmethod
22    def schema(cls, *args: Any, **kwargs: Any) -> Dict[str, Any]:
23        """We're overriding the schema classmethod to enable some post-processing"""
24        schema = super().schema(*args, **kwargs)
25        rename_key(schema, old_key="anyOf", new_key="oneOf")  # UI supports only oneOf
26        expand_refs(schema)
27        schema.pop("description", None)  # description added from the docstring
28        return schema

56class Source(
57    DefaultConnectorMixin,
58    BaseSource[Mapping[str, Any], List[AirbyteStateMessage], ConfiguredAirbyteCatalog],
59    ABC,
60):
61    # can be overridden to change an input state.
62    @classmethod
63    def read_state(cls, state_path: str) -> List[AirbyteStateMessage]:
64        """
65        Retrieves the input state of a sync by reading from the specified JSON file. Incoming state can be deserialized into either
66        a JSON object for legacy state input or as a list of AirbyteStateMessages for the per-stream state format. Regardless of the
67        incoming input type, it will always be transformed and output as a list of AirbyteStateMessage(s).
68        :param state_path: The filepath to where the stream states are located
69        :return: The complete stream state based on the connector's previous sync
70        """
71        parsed_state_messages = []
72        if state_path:
73            state_obj = BaseConnector._read_json_file(state_path)
74            if state_obj:
75                for state in state_obj:  # type: ignore  # `isinstance(state_obj, List)` ensures that this is a list
76                    parsed_message = AirbyteStateMessageSerializer.load(state)
77                    if (
78                        not parsed_message.stream
79                        and not parsed_message.data
80                        and not parsed_message.global_
81                    ):
82                        raise ValueError(
83                            "AirbyteStateMessage should contain either a stream, global, or state field"
84                        )
85                    parsed_state_messages.append(parsed_message)
86        return parsed_state_messages
87
88    # can be overridden to change an input catalog
89    @classmethod
90    def read_catalog(cls, catalog_path: str) -> ConfiguredAirbyteCatalog:
91        return ConfiguredAirbyteCatalogSerializer.load(cls._read_json_file(catalog_path))
92
93    @property
94    def name(self) -> str:
95        """Source name"""
96        return self.__class__.__name__