airbyte_cdk.sources.declarative.checks
1# 2# Copyright (c) 2025 Airbyte, Inc., all rights reserved. 3# 4 5from typing import Mapping 6 7from pydantic.v1 import BaseModel 8 9from airbyte_cdk.sources.declarative.checks.check_dynamic_stream import CheckDynamicStream 10from airbyte_cdk.sources.declarative.checks.check_stream import ( 11 CheckStream, 12 DynamicStreamCheckConfig, 13) 14from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker 15from airbyte_cdk.sources.declarative.models import ( 16 CheckDynamicStream as CheckDynamicStreamModel, 17) 18from airbyte_cdk.sources.declarative.models import ( 19 CheckStream as CheckStreamModel, 20) 21 22COMPONENTS_CHECKER_TYPE_MAPPING: Mapping[str, type[BaseModel]] = { 23 "CheckStream": CheckStreamModel, 24 "CheckDynamicStream": CheckDynamicStreamModel, 25} 26 27__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker", "DynamicStreamCheckConfig"]
43@dataclass 44class CheckStream(ConnectionChecker): 45 """ 46 Checks the connections by checking availability of one or many streams selected by the developer 47 48 Attributes: 49 stream_name (List[str]): names of streams to check 50 """ 51 52 stream_names: List[str] 53 parameters: InitVar[Mapping[str, Any]] 54 dynamic_streams_check_configs: Optional[List[DynamicStreamCheckConfig]] = None 55 56 def __post_init__(self, parameters: Mapping[str, Any]) -> None: 57 self._parameters = parameters 58 if self.dynamic_streams_check_configs is None: 59 self.dynamic_streams_check_configs = [] 60 61 def _log_error(self, logger: logging.Logger, action: str, error: Exception) -> Tuple[bool, str]: 62 """Logs an error and returns a formatted error message.""" 63 error_message = f"Encountered an error while {action}. Error: {error}" 64 logger.error(error_message + f"Error traceback: \n {traceback.format_exc()}", exc_info=True) 65 return False, error_message 66 67 def check_connection( 68 self, 69 source: Source, 70 logger: logging.Logger, 71 config: Mapping[str, Any], 72 ) -> Tuple[bool, Any]: 73 """Checks the connection to the source and its streams.""" 74 try: 75 streams: List[Union[Stream, AbstractStream]] = source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker 76 if not streams: 77 return False, f"No streams to connect to from source {source}" 78 except Exception as error: 79 return self._log_error(logger, "discovering streams", error) 80 81 stream_name_to_stream = {s.name: s for s in streams} 82 for stream_name in self.stream_names: 83 if stream_name not in stream_name_to_stream: 84 raise ValueError( 85 f"{stream_name} is not part of the catalog. Expected one of {list(stream_name_to_stream.keys())}." 86 ) 87 88 stream_availability, message = self._check_stream_availability( 89 stream_name_to_stream, stream_name, logger 90 ) 91 if not stream_availability: 92 return stream_availability, message 93 94 should_check_dynamic_streams = ( 95 hasattr(source, "resolved_manifest") 96 and hasattr(source, "dynamic_streams") 97 and self.dynamic_streams_check_configs 98 ) 99 100 if should_check_dynamic_streams: 101 return self._check_dynamic_streams_availability(source, stream_name_to_stream, logger) 102 103 return True, None 104 105 def _check_stream_availability( 106 self, 107 stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]], 108 stream_name: str, 109 logger: logging.Logger, 110 ) -> Tuple[bool, Any]: 111 """Checks if streams are available.""" 112 try: 113 stream = stream_name_to_stream[stream_name] 114 stream_is_available, reason = evaluate_availability(stream, logger) 115 if not stream_is_available: 116 message = f"Stream {stream_name} is not available: {reason}" 117 logger.warning(message) 118 return stream_is_available, message 119 except Exception as error: 120 return self._log_error(logger, f"checking availability of stream {stream_name}", error) 121 return True, None 122 123 def _check_dynamic_streams_availability( 124 self, 125 source: Source, 126 stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]], 127 logger: logging.Logger, 128 ) -> Tuple[bool, Any]: 129 """Checks the availability of dynamic streams.""" 130 dynamic_streams = source.resolved_manifest.get("dynamic_streams", []) # type: ignore[attr-defined] # The source's resolved_manifest manifest is checked before calling this method 131 dynamic_stream_name_to_dynamic_stream = { 132 ds.get("name", f"dynamic_stream_{i}"): ds for i, ds in enumerate(dynamic_streams) 133 } 134 generated_streams = self._map_generated_streams(source.dynamic_streams) # type: ignore[attr-defined] # The source's dynamic_streams manifest is checked before calling this method 135 136 for check_config in self.dynamic_streams_check_configs: # type: ignore[union-attr] # None value for self.dynamic_streams_check_configs handled in __post_init__ 137 if check_config.dynamic_stream_name not in dynamic_stream_name_to_dynamic_stream: 138 return ( 139 False, 140 f"Dynamic stream {check_config.dynamic_stream_name} is not found in manifest.", 141 ) 142 143 generated = generated_streams.get(check_config.dynamic_stream_name, []) 144 stream_availability, message = self._check_generated_streams_availability( 145 generated, stream_name_to_stream, logger, check_config.stream_count 146 ) 147 if not stream_availability: 148 return stream_availability, message 149 150 return True, None 151 152 def _map_generated_streams( 153 self, dynamic_streams: List[Dict[str, Any]] 154 ) -> Dict[str, List[Dict[str, Any]]]: 155 """Maps dynamic stream names to their corresponding generated streams.""" 156 mapped_streams: Dict[str, List[Dict[str, Any]]] = {} 157 for stream in dynamic_streams: 158 mapped_streams.setdefault(stream["dynamic_stream_name"], []).append(stream) 159 return mapped_streams 160 161 def _check_generated_streams_availability( 162 self, 163 generated_streams: List[Dict[str, Any]], 164 stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]], 165 logger: logging.Logger, 166 max_count: int, 167 ) -> Tuple[bool, Any]: 168 """Checks availability of generated dynamic streams.""" 169 for declarative_stream in generated_streams[: min(max_count, len(generated_streams))]: 170 stream = stream_name_to_stream[declarative_stream["name"]] 171 try: 172 stream_is_available, reason = evaluate_availability(stream, logger) 173 if not stream_is_available: 174 message = f"Dynamic Stream {stream.name} is not available: {reason}" 175 logger.warning(message) 176 return False, message 177 except Exception as error: 178 return self._log_error( 179 logger, f"checking availability of dynamic stream {stream.name}", error 180 ) 181 return True, None
Checks the connections by checking availability of one or many streams selected by the developer
Attributes:
- stream_name (List[str]): names of streams to check
67 def check_connection( 68 self, 69 source: Source, 70 logger: logging.Logger, 71 config: Mapping[str, Any], 72 ) -> Tuple[bool, Any]: 73 """Checks the connection to the source and its streams.""" 74 try: 75 streams: List[Union[Stream, AbstractStream]] = source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker 76 if not streams: 77 return False, f"No streams to connect to from source {source}" 78 except Exception as error: 79 return self._log_error(logger, "discovering streams", error) 80 81 stream_name_to_stream = {s.name: s for s in streams} 82 for stream_name in self.stream_names: 83 if stream_name not in stream_name_to_stream: 84 raise ValueError( 85 f"{stream_name} is not part of the catalog. Expected one of {list(stream_name_to_stream.keys())}." 86 ) 87 88 stream_availability, message = self._check_stream_availability( 89 stream_name_to_stream, stream_name, logger 90 ) 91 if not stream_availability: 92 return stream_availability, message 93 94 should_check_dynamic_streams = ( 95 hasattr(source, "resolved_manifest") 96 and hasattr(source, "dynamic_streams") 97 and self.dynamic_streams_check_configs 98 ) 99 100 if should_check_dynamic_streams: 101 return self._check_dynamic_streams_availability(source, stream_name_to_stream, logger) 102 103 return True, None
Checks the connection to the source and its streams.
17@dataclass 18class CheckDynamicStream(ConnectionChecker): 19 """ 20 Checks the connections by checking availability of one or many dynamic streams 21 22 Attributes: 23 stream_count (int): numbers of streams to check 24 """ 25 26 # TODO: Add field stream_names to check_connection for static streams 27 # https://github.com/airbytehq/airbyte-python-cdk/pull/293#discussion_r1934933483 28 29 stream_count: int 30 parameters: InitVar[Mapping[str, Any]] 31 use_check_availability: bool = True 32 33 def __post_init__(self, parameters: Mapping[str, Any]) -> None: 34 self._parameters = parameters 35 36 def check_connection( 37 self, 38 source: Source, 39 logger: logging.Logger, 40 config: Mapping[str, Any], 41 ) -> Tuple[bool, Any]: 42 streams: List[Union[Stream, AbstractStream]] = source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker 43 44 if len(streams) == 0: 45 return False, f"No streams to connect to from source {source}" 46 if not self.use_check_availability: 47 return True, None 48 49 try: 50 for stream in streams[: min(self.stream_count, len(streams))]: 51 stream_is_available, reason = evaluate_availability(stream, logger) 52 if not stream_is_available: 53 logger.warning(f"Stream {stream.name} is not available: {reason}") 54 return False, reason 55 except Exception as error: 56 error_message = ( 57 f"Encountered an error trying to connect to stream {stream.name}. Error: {error}" 58 ) 59 logger.error(error_message, exc_info=True) 60 return False, error_message 61 62 return True, None
Checks the connections by checking availability of one or many dynamic streams
Attributes:
- stream_count (int): numbers of streams to check
36 def check_connection( 37 self, 38 source: Source, 39 logger: logging.Logger, 40 config: Mapping[str, Any], 41 ) -> Tuple[bool, Any]: 42 streams: List[Union[Stream, AbstractStream]] = source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker 43 44 if len(streams) == 0: 45 return False, f"No streams to connect to from source {source}" 46 if not self.use_check_availability: 47 return True, None 48 49 try: 50 for stream in streams[: min(self.stream_count, len(streams))]: 51 stream_is_available, reason = evaluate_availability(stream, logger) 52 if not stream_is_available: 53 logger.warning(f"Stream {stream.name} is not available: {reason}") 54 return False, reason 55 except Exception as error: 56 error_message = ( 57 f"Encountered an error trying to connect to stream {stream.name}. Error: {error}" 58 ) 59 logger.error(error_message, exc_info=True) 60 return False, error_message 61 62 return True, None
Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect to the Stripe API.
Parameters
- source: source
- 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.
13class ConnectionChecker(ABC): 14 """ 15 Abstract base class for checking a connection 16 """ 17 18 @abstractmethod 19 def check_connection( 20 self, 21 source: Source, 22 logger: logging.Logger, 23 config: Mapping[str, Any], 24 ) -> Tuple[bool, Any]: 25 """ 26 Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect 27 to the Stripe API. 28 29 :param source: source 30 :param logger: source logger 31 :param config: The user-provided configuration as specified by the source's spec. 32 This usually contains information required to check connection e.g. tokens, secrets and keys etc. 33 :return: A tuple of (boolean, error). If boolean is true, then the connection check is successful 34 and we can connect to the underlying data source using the provided configuration. 35 Otherwise, the input config cannot be used to connect to the underlying data source, 36 and the "error" object should describe what went wrong. 37 The error object will be cast to string to display the problem to the user. 38 """ 39 pass
Abstract base class for checking a connection
18 @abstractmethod 19 def check_connection( 20 self, 21 source: Source, 22 logger: logging.Logger, 23 config: Mapping[str, Any], 24 ) -> Tuple[bool, Any]: 25 """ 26 Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect 27 to the Stripe API. 28 29 :param source: source 30 :param logger: source logger 31 :param config: The user-provided configuration as specified by the source's spec. 32 This usually contains information required to check connection e.g. tokens, secrets and keys etc. 33 :return: A tuple of (boolean, error). If boolean is true, then the connection check is successful 34 and we can connect to the underlying data source using the provided configuration. 35 Otherwise, the input config cannot be used to connect to the underlying data source, 36 and the "error" object should describe what went wrong. 37 The error object will be cast to string to display the problem to the user. 38 """ 39 pass
Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect to the Stripe API.
Parameters
- source: source
- 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.
33@dataclass(frozen=True) 34class DynamicStreamCheckConfig: 35 """Defines the configuration for dynamic stream during connection checking. This class specifies 36 what dynamic streams in the stream template should be updated with value, supporting dynamic interpolation 37 and type enforcement.""" 38 39 dynamic_stream_name: str 40 stream_count: int = 0
Defines the configuration for dynamic stream during connection checking. This class specifies what dynamic streams in the stream template should be updated with value, supporting dynamic interpolation and type enforcement.