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