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,
 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

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                    message = f"Stream {stream.name} is not available: {reason}"
54                    logger.warning(message)
55                    return False, message
56        except Exception as error:
57            error_message = (
58                f"Encountered an error trying to connect to stream {stream.name}. Error: {error}"
59            )
60            logger.error(error_message, exc_info=True)
61            return False, error_message
62
63        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

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