Merge branch 'main' into patch-1

Author: natikgadzhi

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -39,6 +39,7 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         token_expiry_date_format str: format of the datetime; provide it if expires_in is returned in datetime instead of seconds
         token_expiry_is_time_of_expiration bool: set True it if expires_in is returned as time of expiration instead of the number seconds until expiration
         refresh_request_body (Optional[Mapping[str, Any]]): The request body to send in the refresh request
+        refresh_request_headers (Optional[Mapping[str, Any]]): The request headers to send in the refresh request
         grant_type: The grant_type to request for access_token. If set to refresh_token, the refresh_token parameter has to be provided
         message_repository (MessageRepository): the message repository used to emit logs on HTTP requests
     """
@@ -56,8 +57,13 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
     token_expiry_is_time_of_expiration: bool = False
     access_token_name: Union[InterpolatedString, str] = "access_token"
     access_token_value: Optional[Union[InterpolatedString, str]] = None
+    client_id_name: Union[InterpolatedString, str] = "client_id"
+    client_secret_name: Union[InterpolatedString, str] = "client_secret"
     expires_in_name: Union[InterpolatedString, str] = "expires_in"
+    refresh_token_name: Union[InterpolatedString, str] = "refresh_token"
     refresh_request_body: Optional[Mapping[str, Any]] = None
+    refresh_request_headers: Optional[Mapping[str, Any]] = None
+    grant_type_name: Union[InterpolatedString, str] = "grant_type"
     grant_type: Union[InterpolatedString, str] = "refresh_token"
     message_repository: MessageRepository = NoopMessageRepository()
 
@@ -69,8 +75,15 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
             )
         else:
             self._token_refresh_endpoint = None
+        self._client_id_name = InterpolatedString.create(self.client_id_name, parameters=parameters)
         self._client_id = InterpolatedString.create(self.client_id, parameters=parameters)
+        self._client_secret_name = InterpolatedString.create(
+            self.client_secret_name, parameters=parameters
+        )
         self._client_secret = InterpolatedString.create(self.client_secret, parameters=parameters)
+        self._refresh_token_name = InterpolatedString.create(
+            self.refresh_token_name, parameters=parameters
+        )
         if self.refresh_token is not None:
             self._refresh_token: Optional[InterpolatedString] = InterpolatedString.create(
                 self.refresh_token, parameters=parameters
@@ -83,10 +96,16 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self.expires_in_name = InterpolatedString.create(
             self.expires_in_name, parameters=parameters
         )
+        self.grant_type_name = InterpolatedString.create(
+            self.grant_type_name, parameters=parameters
+        )
         self.grant_type = InterpolatedString.create(self.grant_type, parameters=parameters)
         self._refresh_request_body = InterpolatedMapping(
             self.refresh_request_body or {}, parameters=parameters
         )
+        self._refresh_request_headers = InterpolatedMapping(
+            self.refresh_request_headers or {}, parameters=parameters
+        )
         self._token_expiry_date: pendulum.DateTime = (
             pendulum.parse(
                 InterpolatedString.create(self.token_expiry_date, parameters=parameters).eval(
@@ -122,18 +141,27 @@ def get_token_refresh_endpoint(self) -> Optional[str]:
             return refresh_token_endpoint
         return None
 
+    def get_client_id_name(self) -> str:
+        return self._client_id_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
     def get_client_id(self) -> str:
         client_id: str = self._client_id.eval(self.config)
         if not client_id:
             raise ValueError("OAuthAuthenticator was unable to evaluate client_id parameter")
         return client_id
 
+    def get_client_secret_name(self) -> str:
+        return self._client_secret_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
     def get_client_secret(self) -> str:
         client_secret: str = self._client_secret.eval(self.config)
         if not client_secret:
             raise ValueError("OAuthAuthenticator was unable to evaluate client_secret parameter")
         return client_secret
 
+    def get_refresh_token_name(self) -> str:
+        return self._refresh_token_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
     def get_refresh_token(self) -> Optional[str]:
         return None if self._refresh_token is None else str(self._refresh_token.eval(self.config))
 
@@ -146,12 +174,18 @@ def get_access_token_name(self) -> str:
     def get_expires_in_name(self) -> str:
         return self.expires_in_name.eval(self.config)  # type: ignore # eval returns a string in this context
 
+    def get_grant_type_name(self) -> str:
+        return self.grant_type_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
     def get_grant_type(self) -> str:
         return self.grant_type.eval(self.config)  # type: ignore # eval returns a string in this context
 
     def get_refresh_request_body(self) -> Mapping[str, Any]:
         return self._refresh_request_body.eval(self.config)
 
+    def get_refresh_request_headers(self) -> Mapping[str, Any]:
+        return self._refresh_request_headers.eval(self.config)
+
     def get_token_expiry_date(self) -> pendulum.DateTime:
         return self._token_expiry_date  # type: ignore # _token_expiry_date is a pendulum.DateTime. It is never None despite what mypy thinks
 

airbyte_cdk/sources/declarative/checks/__init__.py

@@ -1,8 +1,24 @@
 #
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
 #
 
+from typing import Mapping
+
+from pydantic.v1 import BaseModel
+
+from airbyte_cdk.sources.declarative.checks.check_dynamic_stream import CheckDynamicStream
 from airbyte_cdk.sources.declarative.checks.check_stream import CheckStream
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
+from airbyte_cdk.sources.declarative.models import (
+    CheckDynamicStream as CheckDynamicStreamModel,
+)
+from airbyte_cdk.sources.declarative.models import (
+    CheckStream as CheckStreamModel,
+)
+
+COMPONENTS_CHECKER_TYPE_MAPPING: Mapping[str, type[BaseModel]] = {
+    "CheckStream": CheckStreamModel,
+    "CheckDynamicStream": CheckDynamicStreamModel,
+}
 
-__all__ = ["CheckStream", "ConnectionChecker"]
+__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker"]

airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py

@@ -0,0 +1,51 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+import traceback
+from dataclasses import InitVar, dataclass
+from typing import Any, List, Mapping, Tuple
+
+from airbyte_cdk import AbstractSource
+from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
+from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy
+
+
+@dataclass
+class CheckDynamicStream(ConnectionChecker):
+    """
+    Checks the connections by checking availability of one or many dynamic streams
+
+    Attributes:
+        stream_count (int): numbers of streams to check
+    """
+
+    stream_count: int
+    parameters: InitVar[Mapping[str, Any]]
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        self._parameters = parameters
+
+    def check_connection(
+        self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> Tuple[bool, Any]:
+        streams = source.streams(config=config)
+        if len(streams) == 0:
+            return False, f"No streams to connect to from source {source}"
+
+        for stream_index in range(min(self.stream_count, len(streams))):
+            stream = streams[stream_index]
+            availability_strategy = HttpAvailabilityStrategy()
+            try:
+                stream_is_available, reason = availability_strategy.check_availability(
+                    stream, logger
+                )
+                if not stream_is_available:
+                    return False, reason
+            except Exception as error:
+                logger.error(
+                    f"Encountered an error trying to connect to stream {stream.name}. Error: \n {traceback.format_exc()}"
+                )
+                return False, f"Unable to connect to stream {stream.name} - {error}"
+        return True, None