airbyte_cdk.sources.declarative.checks

 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, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
 69    ) -> Tuple[bool, Any]:
 70        """Checks the connection to the source and its streams."""
 71        try:
 72            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
 73            if not streams:
 74                return False, f"No streams to connect to from source {source}"
 75        except Exception as error:
 76            return self._log_error(logger, "discovering streams", error)
 77
 78        stream_name_to_stream = {s.name: s for s in streams}
 79        for stream_name in self.stream_names:
 80            if stream_name not in stream_name_to_stream:
 81                raise ValueError(
 82                    f"{stream_name} is not part of the catalog. Expected one of {list(stream_name_to_stream.keys())}."
 83                )
 84
 85            stream_availability, message = self._check_stream_availability(
 86                stream_name_to_stream, stream_name, logger
 87            )
 88            if not stream_availability:
 89                return stream_availability, message
 90
 91        should_check_dynamic_streams = (
 92            hasattr(source, "resolved_manifest")
 93            and hasattr(source, "dynamic_streams")
 94            and self.dynamic_streams_check_configs
 95        )
 96
 97        if should_check_dynamic_streams:
 98            return self._check_dynamic_streams_availability(source, stream_name_to_stream, logger)
 99
100        return True, None
101
102    def _check_stream_availability(
103        self,
104        stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]],
105        stream_name: str,
106        logger: logging.Logger,
107    ) -> Tuple[bool, Any]:
108        """Checks if streams are available."""
109        try:
110            stream = stream_name_to_stream[stream_name]
111            stream_is_available, reason = evaluate_availability(stream, logger)
112            if not stream_is_available:
113                message = f"Stream {stream_name} is not available: {reason}"
114                logger.warning(message)
115                return stream_is_available, message
116        except Exception as error:
117            return self._log_error(logger, f"checking availability of stream {stream_name}", error)
118        return True, None
119
120    def _check_dynamic_streams_availability(
121        self,
122        source: AbstractSource,
123        stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]],
124        logger: logging.Logger,
125    ) -> Tuple[bool, Any]:
126        """Checks the availability of dynamic streams."""
127        dynamic_streams = source.resolved_manifest.get("dynamic_streams", [])  # type: ignore[attr-defined] # The source's resolved_manifest manifest is checked before calling this method
128        dynamic_stream_name_to_dynamic_stream = {
129            ds.get("name", f"dynamic_stream_{i}"): ds for i, ds in enumerate(dynamic_streams)
130        }
131        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
132
133        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__
134            if check_config.dynamic_stream_name not in dynamic_stream_name_to_dynamic_stream:
135                return (
136                    False,
137                    f"Dynamic stream {check_config.dynamic_stream_name} is not found in manifest.",
138                )
139
140            generated = generated_streams.get(check_config.dynamic_stream_name, [])
141            stream_availability, message = self._check_generated_streams_availability(
142                generated, stream_name_to_stream, logger, check_config.stream_count
143            )
144            if not stream_availability:
145                return stream_availability, message
146
147        return True, None
148
149    def _map_generated_streams(
150        self, dynamic_streams: List[Dict[str, Any]]
151    ) -> Dict[str, List[Dict[str, Any]]]:
152        """Maps dynamic stream names to their corresponding generated streams."""
153        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
154        for stream in dynamic_streams:
155            mapped_streams.setdefault(stream["dynamic_stream_name"], []).append(stream)
156        return mapped_streams
157
158    def _check_generated_streams_availability(
159        self,
160        generated_streams: List[Dict[str, Any]],
161        stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]],
162        logger: logging.Logger,
163        max_count: int,
164    ) -> Tuple[bool, Any]:
165        """Checks availability of generated dynamic streams."""
166        for declarative_stream in generated_streams[: min(max_count, len(generated_streams))]:
167            stream = stream_name_to_stream[declarative_stream["name"]]
168            try:
169                stream_is_available, reason = evaluate_availability(stream, logger)
170                if not stream_is_available:
171                    message = f"Dynamic Stream {stream.name} is not available: {reason}"
172                    logger.warning(message)
173                    return False, message
174            except Exception as error:
175                return self._log_error(
176                    logger, f"checking availability of dynamic stream {stream.name}", error
177                )
178        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, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
38    ) -> Tuple[bool, Any]:
39        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
40
41        if len(streams) == 0:
42            return False, f"No streams to connect to from source {source}"
43        if not self.use_check_availability:
44            return True, None
45
46        try:
47            for stream in streams[: min(self.stream_count, len(streams))]:
48                stream_is_available, reason = evaluate_availability(stream, logger)
49                if not stream_is_available:
50                    logger.warning(f"Stream {stream.name} is not available: {reason}")
51                    return False, reason
52        except Exception as error:
53            error_message = (
54                f"Encountered an error trying to connect to stream {stream.name}. Error: {error}"
55            )
56            logger.error(error_message, exc_info=True)
57            return False, error_message
58
59        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, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
21    ) -> Tuple[bool, Any]:
22        """
23        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
24        to the Stripe API.
25
26        :param source: source
27        :param logger: source logger
28        :param config: The user-provided configuration as specified by the source's spec.
29          This usually contains information required to check connection e.g. tokens, secrets and keys etc.
30        :return: A tuple of (boolean, error). If boolean is true, then the connection check is successful
31          and we can connect to the underlying data source using the provided configuration.
32          Otherwise, the input config cannot be used to connect to the underlying data source,
33          and the "error" object should describe what went wrong.
34          The error object will be cast to string to display the problem to the user.
35        """
36        pass

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