airbyte_cdk.sources.declarative.checks

 45@dataclass
 46class CheckStream(ConnectionChecker):
 47    """
 48    Checks the connections by checking availability of one or many streams selected by the developer
 49
 50    Attributes:
 51        stream_name (List[str]): names of streams to check
 52    """
 53
 54    stream_names: List[str]
 55    parameters: InitVar[Mapping[str, Any]]
 56    dynamic_streams_check_configs: Optional[List[DynamicStreamCheckConfig]] = None
 57
 58    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 59        self._parameters = parameters
 60        if self.dynamic_streams_check_configs is None:
 61            self.dynamic_streams_check_configs = []
 62
 63    def _log_error(self, logger: logging.Logger, action: str, error: Exception) -> Tuple[bool, str]:
 64        """Logs an error and returns a formatted error message."""
 65        error_message = f"Encountered an error while {action}. Error: {error}"
 66        logger.error(error_message + f"Error traceback: \n {traceback.format_exc()}", exc_info=True)
 67        return False, error_message
 68
 69    def check_connection(
 70        self,
 71        source: Source,
 72        logger: logging.Logger,
 73        config: Mapping[str, Any],
 74    ) -> Tuple[bool, Any]:
 75        """Checks the connection to the source and its streams."""
 76        stream_names = self._get_stream_names(config)
 77        try:
 78            streams: List[Union[Stream, AbstractStream]] = self._build_streams(
 79                source, config, stream_names
 80            )
 81            if not streams:
 82                return False, f"No streams to connect to from source {source}"
 83        except Exception as error:
 84            return self._log_error(logger, "discovering streams", error)
 85
 86        stream_name_to_stream = {s.name: s for s in streams}
 87        for stream_name in stream_names:
 88            if stream_name not in stream_name_to_stream:
 89                raise ValueError(
 90                    f"{stream_name} is not part of the catalog. Expected one of {list(stream_name_to_stream.keys())}."
 91                )
 92
 93            stream_availability, message = self._check_stream_availability(
 94                stream_name_to_stream, stream_name, logger
 95            )
 96            if not stream_availability:
 97                return stream_availability, message
 98
 99        should_check_dynamic_streams = (
100            self.dynamic_streams_check_configs
101            and hasattr(source, "resolved_manifest")
102            and hasattr(source, "dynamic_streams")
103        )
104
105        if should_check_dynamic_streams:
106            return self._check_dynamic_streams_availability(source, stream_name_to_stream, logger)
107
108        return True, None
109
110    def _build_streams(
111        self,
112        source: Source,
113        config: Mapping[str, Any],
114        stream_names: List[str],
115    ) -> List[Union[Stream, AbstractStream]]:
116        from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
117            ConcurrentDeclarativeSource,
118        )
119
120        if (
121            isinstance(source, ConcurrentDeclarativeSource)
122            and not self.dynamic_streams_check_configs
123        ):
124            streams = source._streams_by_name(stream_names)
125            if streams is not None:
126                return cast(List[Union[Stream, AbstractStream]], streams)
127
128        return source.streams(config=config)  # type: ignore  # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker
129
130    def _get_stream_names(self, config: Mapping[str, Any]) -> List[str]:
131        if (
132            CHECK_STREAM_NAMES_CONFIG_KEY not in config
133            or config[CHECK_STREAM_NAMES_CONFIG_KEY] == []
134        ):
135            return self.stream_names
136        configured_stream_names = config[CHECK_STREAM_NAMES_CONFIG_KEY]
137        if not isinstance(configured_stream_names, list) or not all(
138            isinstance(stream_name, str) for stream_name in configured_stream_names
139        ):
140            raise ValueError(f"{CHECK_STREAM_NAMES_CONFIG_KEY} must be a list of strings.")
141        return configured_stream_names
142
143    def _check_stream_availability(
144        self,
145        stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]],
146        stream_name: str,
147        logger: logging.Logger,
148    ) -> Tuple[bool, Any]:
149        """Checks if streams are available."""
150        try:
151            stream = stream_name_to_stream[stream_name]
152            stream_is_available, reason = evaluate_availability(stream, logger)
153            if not stream_is_available:
154                message = f"Stream {stream_name} is not available: {reason}"
155                logger.warning(message)
156                return stream_is_available, message
157        except Exception as error:
158            return self._log_error(logger, f"checking availability of stream {stream_name}", error)
159        return True, None
160
161    def _check_dynamic_streams_availability(
162        self,
163        source: Source,
164        stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]],
165        logger: logging.Logger,
166    ) -> Tuple[bool, Any]:
167        """Checks the availability of dynamic streams."""
168        dynamic_streams = source.resolved_manifest.get("dynamic_streams", [])  # type: ignore[attr-defined] # The source's resolved_manifest manifest is checked before calling this method
169        dynamic_stream_name_to_dynamic_stream = {
170            ds.get("name", f"dynamic_stream_{i}"): ds for i, ds in enumerate(dynamic_streams)
171        }
172        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
173
174        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__
175            if check_config.dynamic_stream_name not in dynamic_stream_name_to_dynamic_stream:
176                return (
177                    False,
178                    f"Dynamic stream {check_config.dynamic_stream_name} is not found in manifest.",
179                )
180
181            generated = generated_streams.get(check_config.dynamic_stream_name, [])
182            stream_availability, message = self._check_generated_streams_availability(
183                generated, stream_name_to_stream, logger, check_config.stream_count
184            )
185            if not stream_availability:
186                return stream_availability, message
187
188        return True, None
189
190    def _map_generated_streams(
191        self, dynamic_streams: List[Dict[str, Any]]
192    ) -> Dict[str, List[Dict[str, Any]]]:
193        """Maps dynamic stream names to their corresponding generated streams."""
194        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
195        for stream in dynamic_streams:
196            mapped_streams.setdefault(stream["dynamic_stream_name"], []).append(stream)
197        return mapped_streams
198
199    def _check_generated_streams_availability(
200        self,
201        generated_streams: List[Dict[str, Any]],
202        stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]],
203        logger: logging.Logger,
204        max_count: Optional[int],
205    ) -> Tuple[bool, Any]:
206        """Checks availability of generated dynamic streams.
207
208        If `max_count` is `None`, all generated streams are checked. Otherwise, the
209        first `max_count` streams are checked (capped at the number of available streams).
210        """
211        streams_to_check = generated_streams if max_count is None else generated_streams[:max_count]
212        for declarative_stream in streams_to_check:
213            stream = stream_name_to_stream[declarative_stream["name"]]
214            try:
215                stream_is_available, reason = evaluate_availability(stream, logger)
216                if not stream_is_available:
217                    message = f"Dynamic Stream {stream.name} is not available: {reason}"
218                    logger.warning(message)
219                    return False, message
220            except Exception as error:
221                return self._log_error(
222                    logger, f"checking availability of dynamic stream {stream.name}", error
223                )
224        return True, None

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        from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
43            ConcurrentDeclarativeSource,
44        )
45
46        streams: Optional[List[Union[Stream, AbstractStream]]] = None
47        if (
48            isinstance(source, ConcurrentDeclarativeSource)
49            and self.use_check_availability
50            and self.stream_count >= 1
51        ):
52            streams = cast(
53                Optional[List[Union[Stream, AbstractStream]]],
54                source._first_streams(self.stream_count),
55            )
56        if streams is None:
57            streams = source.streams(config=config)  # type: ignore  # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker
58
59        if len(streams) == 0:
60            return False, f"No streams to connect to from source {source}"
61        if not self.use_check_availability:
62            return True, None
63
64        try:
65            for stream in streams[: min(self.stream_count, len(streams))]:
66                stream_is_available, reason = evaluate_availability(stream, logger)
67                if not stream_is_available:
68                    message = f"Stream {stream.name} is not available: {reason}"
69                    logger.warning(message)
70                    return False, message
71        except Exception as error:
72            error_message = (
73                f"Encountered an error trying to connect to stream {stream.name}. Error: {error}"
74            )
75            logger.error(error_message, exc_info=True)
76            return False, error_message
77
78        return True, None

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

35@dataclass(frozen=True)
36class DynamicStreamCheckConfig:
37    """Defines the configuration for dynamic stream during connection checking. This class specifies
38    what dynamic streams  in the stream template should be updated with value, supporting dynamic interpolation
39    and type enforcement."""
40
41    dynamic_stream_name: str
42    stream_count: Optional[int] = None