airbyte-cdk 0.72.1__py3-none-any.whl → 6.17.1.dev0__py3-none-any.whl

airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py

@@ -10,12 +10,13 @@ from typing import Any, List, Mapping, MutableMapping, Optional, Tuple, Union
 import backoff
 import pendulum
 import requests
+from requests.auth import AuthBase
+
 from airbyte_cdk.models import FailureType, Level
 from airbyte_cdk.sources.http_logger import format_http_message
 from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
 from airbyte_cdk.utils import AirbyteTracedException
 from airbyte_cdk.utils.airbyte_secrets_utils import add_to_secrets
-from requests.auth import AuthBase
 
 from ..exceptions import DefaultBackoffException
 
@@ -53,7 +54,12 @@ class AbstractOauth2Authenticator(AuthBase):
 
     def get_auth_header(self) -> Mapping[str, Any]:
         """HTTP header to set on the requests"""
-        return {"Authorization": f"Bearer {self.get_access_token()}"}
+        token = self.access_token if self._is_access_token_flow else self.get_access_token()
+        return {"Authorization": f"Bearer {token}"}
+
+    @property
+    def _is_access_token_flow(self) -> bool:
+        return self.get_token_refresh_endpoint() is None and self.access_token is not None
 
     def get_access_token(self) -> str:
         """Returns the access token"""
@@ -92,7 +98,9 @@ class AbstractOauth2Authenticator(AuthBase):
 
         return payload
 
-    def _wrap_refresh_token_exception(self, exception: requests.exceptions.RequestException) -> bool:
+    def _wrap_refresh_token_exception(
+        self, exception: requests.exceptions.RequestException
+    ) -> bool:
         try:
             if exception.response is not None:
                 exception_content = exception.response.json()
@@ -102,7 +110,8 @@ class AbstractOauth2Authenticator(AuthBase):
             return False
         return (
             exception.response.status_code in self._refresh_token_error_status_codes
-            and exception_content.get(self._refresh_token_error_key) in self._refresh_token_error_values
+            and exception_content.get(self._refresh_token_error_key)
+            in self._refresh_token_error_values
         )
 
     @backoff.on_exception(
@@ -115,14 +124,20 @@ class AbstractOauth2Authenticator(AuthBase):
     )
     def _get_refresh_access_token_response(self) -> Any:
         try:
-            response = requests.request(method="POST", url=self.get_token_refresh_endpoint(), data=self.build_refresh_request_body())
+            response = requests.request(
+                method="POST",
+                url=self.get_token_refresh_endpoint(),  # type: ignore # returns None, if not provided, but str | bytes is expected.
+                data=self.build_refresh_request_body(),
+            )
             if response.ok:
                 response_json = response.json()
                 # Add the access token to the list of secrets so it is replaced before logging the response
                 # An argument could be made to remove the prevous access key from the list of secrets, but unmasking values seems like a security incident waiting to happen...
                 access_key = response_json.get(self.get_access_token_name())
                 if not access_key:
-                    raise Exception("Token refresh API response was missing access token {self.get_access_token_name()}")
+                    raise Exception(
+                        "Token refresh API response was missing access token {self.get_access_token_name()}"
+                    )
                 add_to_secrets(access_key)
                 self._log_response(response)
                 return response_json
@@ -136,7 +151,9 @@ class AbstractOauth2Authenticator(AuthBase):
                     raise DefaultBackoffException(request=e.response.request, response=e.response)
             if self._wrap_refresh_token_exception(e):
                 message = "Refresh token is invalid or expired. Please re-authenticate from Sources/<your source>/Settings."
-                raise AirbyteTracedException(internal_message=message, message=message, failure_type=FailureType.config_error)
+                raise AirbyteTracedException(
+                    internal_message=message, message=message, failure_type=FailureType.config_error
+                )
             raise
         except Exception as e:
             raise Exception(f"Error while refreshing access token: {e}") from e
@@ -149,7 +166,9 @@ class AbstractOauth2Authenticator(AuthBase):
         """
         response_json = self._get_refresh_access_token_response()
 
-        return response_json[self.get_access_token_name()], response_json[self.get_expires_in_name()]
+        return response_json[self.get_access_token_name()], response_json[
+            self.get_expires_in_name()
+        ]
 
     def _parse_token_expiration_date(self, value: Union[str, int]) -> pendulum.DateTime:
         """
@@ -184,7 +203,7 @@ class AbstractOauth2Authenticator(AuthBase):
         return None
 
     @abstractmethod
-    def get_token_refresh_endpoint(self) -> str:
+    def get_token_refresh_endpoint(self) -> Optional[str]:
         """Returns the endpoint to refresh the access token"""
 
     @abstractmethod

airbyte_cdk/sources/streams/http/requests_native_auth/abstract_token.py

@@ -5,13 +5,14 @@
 from abc import abstractmethod
 from typing import Any, Mapping
 
+import requests
 from requests.auth import AuthBase
 
 
 class AbstractHeaderAuthenticator(AuthBase):
     """Abstract class for an header-based authenticators that add a header to outgoing HTTP requests."""
 
-    def __call__(self, request):
+    def __call__(self, request: requests.PreparedRequest) -> Any:
         """Attach the HTTP headers required to authenticate on the HTTP request"""
         request.headers.update(self.get_auth_header())
         return request

airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py

@@ -6,9 +6,15 @@ from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union
 
 import dpath
 import pendulum
-from airbyte_cdk.config_observation import create_connector_config_control_message, emit_configuration_as_airbyte_control_message
+
+from airbyte_cdk.config_observation import (
+    create_connector_config_control_message,
+    emit_configuration_as_airbyte_control_message,
+)
 from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
-from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_oauth import AbstractOauth2Authenticator
+from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_oauth import (
+    AbstractOauth2Authenticator,
+)
 
 
 class Oauth2Authenticator(AbstractOauth2Authenticator):
@@ -24,12 +30,12 @@ class Oauth2Authenticator(AbstractOauth2Authenticator):
         client_id: str,
         client_secret: str,
         refresh_token: str,
-        scopes: List[str] = None,
-        token_expiry_date: pendulum.DateTime = None,
-        token_expiry_date_format: str = None,
+        scopes: List[str] | None = None,
+        token_expiry_date: pendulum.DateTime | None = None,
+        token_expiry_date_format: str | None = None,
         access_token_name: str = "access_token",
         expires_in_name: str = "expires_in",
-        refresh_request_body: Mapping[str, Any] = None,
+        refresh_request_body: Mapping[str, Any] | None = None,
         grant_type: str = "refresh_token",
         token_expiry_is_time_of_expiration: bool = False,
         refresh_token_error_status_codes: Tuple[int, ...] = (),
@@ -46,11 +52,13 @@ class Oauth2Authenticator(AbstractOauth2Authenticator):
         self._refresh_request_body = refresh_request_body
         self._grant_type = grant_type
 
-        self._token_expiry_date = token_expiry_date or pendulum.now().subtract(days=1)
+        self._token_expiry_date = token_expiry_date or pendulum.now().subtract(days=1)  # type: ignore [no-untyped-call]
         self._token_expiry_date_format = token_expiry_date_format
         self._token_expiry_is_time_of_expiration = token_expiry_is_time_of_expiration
         self._access_token = None
-        super().__init__(refresh_token_error_status_codes, refresh_token_error_key, refresh_token_error_values)
+        super().__init__(
+            refresh_token_error_status_codes, refresh_token_error_key, refresh_token_error_values
+        )
 
     def get_token_refresh_endpoint(self) -> str:
         return self._token_refresh_endpoint
@@ -67,14 +75,14 @@ class Oauth2Authenticator(AbstractOauth2Authenticator):
     def get_access_token_name(self) -> str:
         return self._access_token_name
 
-    def get_scopes(self) -> [str]:
-        return self._scopes
+    def get_scopes(self) -> list[str]:
+        return self._scopes  # type: ignore [return-value]
 
     def get_expires_in_name(self) -> str:
         return self._expires_in_name
 
     def get_refresh_request_body(self) -> Mapping[str, Any]:
-        return self._refresh_request_body
+        return self._refresh_request_body  # type: ignore [return-value]
 
     def get_grant_type(self) -> str:
         return self._grant_type
@@ -82,7 +90,7 @@ class Oauth2Authenticator(AbstractOauth2Authenticator):
     def get_token_expiry_date(self) -> pendulum.DateTime:
         return self._token_expiry_date
 
-    def set_token_expiry_date(self, value: Union[str, int]):
+    def set_token_expiry_date(self, value: Union[str, int]) -> None:
         self._token_expiry_date = self._parse_token_expiration_date(value)
 
     @property
@@ -95,11 +103,11 @@ class Oauth2Authenticator(AbstractOauth2Authenticator):
 
     @property
     def access_token(self) -> str:
-        return self._access_token
+        return self._access_token  # type: ignore [return-value]
 
     @access_token.setter
-    def access_token(self, value: str):
-        self._access_token = value
+    def access_token(self, value: str) -> None:
+        self._access_token = value  # type: ignore [assignment]  # Incorrect type for assignment
 
 
 class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
@@ -116,11 +124,11 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
         self,
         connector_config: Mapping[str, Any],
         token_refresh_endpoint: str,
-        scopes: List[str] = None,
+        scopes: List[str] | None = None,
         access_token_name: str = "access_token",
         expires_in_name: str = "expires_in",
         refresh_token_name: str = "refresh_token",
-        refresh_request_body: Mapping[str, Any] = None,
+        refresh_request_body: Mapping[str, Any] | None = None,
         grant_type: str = "refresh_token",
         client_id: Optional[str] = None,
         client_secret: Optional[str] = None,
@@ -153,9 +161,18 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
             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
             message_repository (MessageRepository): the message repository used to emit logs on HTTP requests and control message on config update
         """
-        self._client_id = client_id if client_id is not None else dpath.util.get(connector_config, ("credentials", "client_id"))
+        self._client_id = (
+            client_id  # type: ignore [assignment]  # Incorrect type for assignment
+            if client_id is not None
+            else dpath.get(connector_config, ("credentials", "client_id"))  # type: ignore [arg-type]
+        )
         self._client_secret = (
-            client_secret if client_secret is not None else dpath.util.get(connector_config, ("credentials", "client_secret"))
+            client_secret  # type: ignore [assignment]  # Incorrect type for assignment
+            if client_secret is not None
+            else dpath.get(
+                connector_config,  # type: ignore [arg-type]
+                ("credentials", "client_secret"),
+            )
         )
         self._access_token_config_path = access_token_config_path
         self._refresh_token_config_path = refresh_token_config_path
@@ -193,31 +210,61 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
 
     @property
     def access_token(self) -> str:
-        return dpath.util.get(self._connector_config, self._access_token_config_path, default="")
+        return dpath.get(  # type: ignore [return-value]
+            self._connector_config,  # type: ignore [arg-type]
+            self._access_token_config_path,
+            default="",
+        )
 
     @access_token.setter
-    def access_token(self, new_access_token: str):
-        dpath.util.new(self._connector_config, self._access_token_config_path, new_access_token)
+    def access_token(self, new_access_token: str) -> None:
+        dpath.new(
+            self._connector_config,  # type: ignore [arg-type]
+            self._access_token_config_path,
+            new_access_token,
+        )
 
     def get_refresh_token(self) -> str:
-        return dpath.util.get(self._connector_config, self._refresh_token_config_path, default="")
+        return dpath.get(  # type: ignore [return-value]
+            self._connector_config,  # type: ignore [arg-type]
+            self._refresh_token_config_path,
+            default="",
+        )
 
-    def set_refresh_token(self, new_refresh_token: str):
-        dpath.util.new(self._connector_config, self._refresh_token_config_path, new_refresh_token)
+    def set_refresh_token(self, new_refresh_token: str) -> None:
+        dpath.new(
+            self._connector_config,  # type: ignore [arg-type]
+            self._refresh_token_config_path,
+            new_refresh_token,
+        )
 
     def get_token_expiry_date(self) -> pendulum.DateTime:
-        expiry_date = dpath.util.get(self._connector_config, self._token_expiry_date_config_path, default="")
-        return pendulum.now().subtract(days=1) if expiry_date == "" else pendulum.parse(expiry_date)
+        expiry_date = dpath.get(
+            self._connector_config,  # type: ignore [arg-type]
+            self._token_expiry_date_config_path,
+            default="",
+        )
+        return pendulum.now().subtract(days=1) if expiry_date == "" else pendulum.parse(expiry_date)  # type: ignore [arg-type, return-value, no-untyped-call]
 
-    def set_token_expiry_date(self, new_token_expiry_date):
-        dpath.util.new(self._connector_config, self._token_expiry_date_config_path, str(new_token_expiry_date))
+    def set_token_expiry_date(  # type: ignore[override]
+        self,
+        new_token_expiry_date: pendulum.DateTime,
+    ) -> None:
+        dpath.new(
+            self._connector_config,  # type: ignore [arg-type]
+            self._token_expiry_date_config_path,
+            str(new_token_expiry_date),
+        )
 
     def token_has_expired(self) -> bool:
         """Returns True if the token is expired"""
         return pendulum.now("UTC") > self.get_token_expiry_date()
 
     @staticmethod
-    def get_new_token_expiry_date(access_token_expires_in: str, token_expiry_date_format: str = None) -> pendulum.DateTime:
+    def get_new_token_expiry_date(
+        access_token_expires_in: str,
+        token_expiry_date_format: str | None = None,
+    ) -> pendulum.DateTime:
         if token_expiry_date_format:
             return pendulum.from_format(access_token_expires_in, token_expiry_date_format)
         else:
@@ -230,8 +277,12 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
             str: The current access_token, updated if it was previously expired.
         """
         if self.token_has_expired():
-            new_access_token, access_token_expires_in, new_refresh_token = self.refresh_access_token()
-            new_token_expiry_date = self.get_new_token_expiry_date(access_token_expires_in, self._token_expiry_date_format)
+            new_access_token, access_token_expires_in, new_refresh_token = (
+                self.refresh_access_token()
+            )
+            new_token_expiry_date: pendulum.DateTime = self.get_new_token_expiry_date(
+                access_token_expires_in, self._token_expiry_date_format
+            )
             self.access_token = new_access_token
             self.set_refresh_token(new_refresh_token)
             self.set_token_expiry_date(new_token_expiry_date)
@@ -239,12 +290,16 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
             #  Usually, a class shouldn't care about the implementation details but to keep backward compatibility where we print the
             #  message directly in the console, this is needed
             if not isinstance(self._message_repository, NoopMessageRepository):
-                self._message_repository.emit_message(create_connector_config_control_message(self._connector_config))
+                self._message_repository.emit_message(
+                    create_connector_config_control_message(self._connector_config)  # type: ignore [arg-type]
+                )
             else:
-                emit_configuration_as_airbyte_control_message(self._connector_config)
+                emit_configuration_as_airbyte_control_message(self._connector_config)  # type: ignore [arg-type]
         return self.access_token
 
-    def refresh_access_token(self) -> Tuple[str, str, str]:
+    def refresh_access_token(  # type: ignore[override]  # Signature doesn't match base class
+        self,
+    ) -> Tuple[str, str, str]:
         response_json = self._get_refresh_access_token_response()
         return (
             response_json[self.get_access_token_name()],

airbyte_cdk/sources/streams/http/requests_native_auth/token.py

@@ -6,7 +6,9 @@ import base64
 from itertools import cycle
 from typing import List
 
-from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_token import AbstractHeaderAuthenticator
+from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_token import (
+    AbstractHeaderAuthenticator,
+)
 
 
 class MultipleTokenAuthenticator(AbstractHeaderAuthenticator):
@@ -24,7 +26,9 @@ class MultipleTokenAuthenticator(AbstractHeaderAuthenticator):
     def token(self) -> str:
         return f"{self._auth_method} {next(self._tokens_iter)}"
 
-    def __init__(self, tokens: List[str], auth_method: str = "Bearer", auth_header: str = "Authorization"):
+    def __init__(
+        self, tokens: List[str], auth_method: str = "Bearer", auth_header: str = "Authorization"
+    ):
         self._auth_method = auth_method
         self._auth_header = auth_header
         self._tokens = tokens
@@ -65,7 +69,13 @@ class BasicHttpAuthenticator(AbstractHeaderAuthenticator):
     def token(self) -> str:
         return f"{self._auth_method} {self._token}"
 
-    def __init__(self, username: str, password: str = "", auth_method: str = "Basic", auth_header: str = "Authorization"):
+    def __init__(
+        self,
+        username: str,
+        password: str = "",
+        auth_method: str = "Basic",
+        auth_header: str = "Authorization",
+    ):
         auth_string = f"{username}:{password}".encode("utf8")
         b64_encoded = base64.b64encode(auth_string).decode("utf8")
         self._auth_header = auth_header

airbyte_cdk/sources/types.py

@@ -0,0 +1,154 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from __future__ import annotations
+
+from typing import Any, ItemsView, Iterator, KeysView, List, Mapping, Optional, ValuesView
+
+import orjson
+
+# A FieldPointer designates a path to a field inside a mapping. For example, retrieving ["k1", "k1.2"] in the object {"k1" :{"k1.2":
+# "hello"}] returns "hello"
+FieldPointer = List[str]
+Config = Mapping[str, Any]
+ConnectionDefinition = Mapping[str, Any]
+StreamState = Mapping[str, Any]
+
+
+class Record(Mapping[str, Any]):
+    def __init__(
+        self,
+        data: Mapping[str, Any],
+        stream_name: str,
+        associated_slice: Optional[StreamSlice] = None,
+        is_file_transfer_message: bool = False,
+    ):
+        self._data = data
+        self._associated_slice = associated_slice
+        self.stream_name = stream_name
+        self.is_file_transfer_message = is_file_transfer_message
+
+    @property
+    def data(self) -> Mapping[str, Any]:
+        return self._data
+
+    @property
+    def associated_slice(self) -> Optional[StreamSlice]:
+        return self._associated_slice
+
+    def __repr__(self) -> str:
+        return repr(self._data)
+
+    def __getitem__(self, key: str) -> Any:
+        return self._data[key]
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __iter__(self) -> Any:
+        return iter(self._data)
+
+    def __contains__(self, item: object) -> bool:
+        return item in self._data
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Record):
+            # noinspection PyProtectedMember
+            return self._data == other._data
+        return False
+
+    def __ne__(self, other: object) -> bool:
+        return not self.__eq__(other)
+
+
+class StreamSlice(Mapping[str, Any]):
+    def __init__(
+        self,
+        *,
+        partition: Mapping[str, Any],
+        cursor_slice: Mapping[str, Any],
+        extra_fields: Optional[Mapping[str, Any]] = None,
+    ) -> None:
+        """
+        :param partition: The partition keys representing a unique partition in the stream.
+        :param cursor_slice: The incremental cursor slice keys, such as dates or pagination tokens.
+        :param extra_fields: Additional fields that should not be part of the partition but passed along, such as metadata from the parent stream.
+        """
+        self._partition = partition
+        self._cursor_slice = cursor_slice
+        self._extra_fields = extra_fields or {}
+
+        # Ensure that partition keys do not overlap with cursor slice keys
+        if partition.keys() & cursor_slice.keys():
+            raise ValueError("Keys for partition and incremental sync cursor should not overlap")
+
+        self._stream_slice = dict(partition) | dict(cursor_slice)
+
+    @property
+    def partition(self) -> Mapping[str, Any]:
+        """Returns the partition portion of the stream slice."""
+        p = self._partition
+        while isinstance(p, StreamSlice):
+            p = p.partition
+        return p
+
+    @property
+    def cursor_slice(self) -> Mapping[str, Any]:
+        """Returns the cursor slice portion of the stream slice."""
+        c = self._cursor_slice
+        while isinstance(c, StreamSlice):
+            c = c.cursor_slice
+        return c
+
+    @property
+    def extra_fields(self) -> Mapping[str, Any]:
+        """Returns the extra fields that are not part of the partition."""
+        return self._extra_fields
+
+    def __repr__(self) -> str:
+        return repr(self._stream_slice)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        raise ValueError("StreamSlice is immutable")
+
+    def __getitem__(self, key: str) -> Any:
+        return self._stream_slice[key]
+
+    def __len__(self) -> int:
+        return len(self._stream_slice)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._stream_slice)
+
+    def __contains__(self, item: Any) -> bool:
+        return item in self._stream_slice
+
+    def keys(self) -> KeysView[str]:
+        return self._stream_slice.keys()
+
+    def items(self) -> ItemsView[str, Any]:
+        return self._stream_slice.items()
+
+    def values(self) -> ValuesView[Any]:
+        return self._stream_slice.values()
+
+    def get(self, key: str, default: Any = None) -> Optional[Any]:
+        return self._stream_slice.get(key, default)
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, dict):
+            return self._stream_slice == other
+        if isinstance(other, StreamSlice):
+            # noinspection PyProtectedMember
+            return self._partition == other._partition and self._cursor_slice == other._cursor_slice
+        return False
+
+    def __ne__(self, other: Any) -> bool:
+        return not self.__eq__(other)
+
+    def __json_serializable__(self) -> Any:
+        return self._stream_slice
+
+    def __hash__(self) -> int:
+        return hash(orjson.dumps(self._stream_slice, option=orjson.OPT_SORT_KEYS))

airbyte_cdk/sources/utils/record_helper.py

@@ -1,12 +1,18 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
+import time
+from collections.abc import Mapping as ABCMapping
+from typing import Any, Mapping, Optional
 
-import datetime
-from typing import Any, Mapping
-
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteRecordMessage, AirbyteTraceMessage
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteRecordMessage,
+    AirbyteTraceMessage,
+)
 from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 
@@ -15,24 +21,33 @@ def stream_data_to_airbyte_message(
     stream_name: str,
     data_or_message: StreamData,
     transformer: TypeTransformer = TypeTransformer(TransformConfig.NoTransform),
-    schema: Mapping[str, Any] = None,
+    schema: Optional[Mapping[str, Any]] = None,
+    is_file_transfer_message: bool = False,
 ) -> AirbyteMessage:
     if schema is None:
         schema = {}
 
-    if isinstance(data_or_message, Mapping):
-        data = dict(data_or_message)
-        now_millis = int(datetime.datetime.now().timestamp() * 1000)
-        # Transform object fields according to config. Most likely you will
-        # need it to normalize values against json schema. By default no action
-        # taken unless configured. See
-        # docs/connector-development/cdk-python/schemas.md for details.
-        transformer.transform(data, schema)  # type: ignore
-        message = AirbyteRecordMessage(stream=stream_name, data=data, emitted_at=now_millis)
-        return AirbyteMessage(type=MessageType.RECORD, record=message)
-    elif isinstance(data_or_message, AirbyteTraceMessage):
-        return AirbyteMessage(type=MessageType.TRACE, trace=data_or_message)
-    elif isinstance(data_or_message, AirbyteLogMessage):
-        return AirbyteMessage(type=MessageType.LOG, log=data_or_message)
-    else:
-        raise ValueError(f"Unexpected type for data_or_message: {type(data_or_message)}: {data_or_message}")
+    match data_or_message:
+        case ABCMapping():
+            data = dict(data_or_message)
+            now_millis = time.time_ns() // 1_000_000
+            # Transform object fields according to config. Most likely you will
+            # need it to normalize values against json schema. By default no action
+            # taken unless configured. See
+            # docs/connector-development/cdk-python/schemas.md for details.
+            transformer.transform(data, schema)
+            if is_file_transfer_message:
+                message = AirbyteFileTransferRecordMessage(
+                    stream=stream_name, file=data, emitted_at=now_millis, data={}
+                )
+            else:
+                message = AirbyteRecordMessage(stream=stream_name, data=data, emitted_at=now_millis)
+            return AirbyteMessage(type=MessageType.RECORD, record=message)
+        case AirbyteTraceMessage():
+            return AirbyteMessage(type=MessageType.TRACE, trace=data_or_message)
+        case AirbyteLogMessage():
+            return AirbyteMessage(type=MessageType.LOG, log=data_or_message)
+        case _:
+            raise ValueError(
+                f"Unexpected type for data_or_message: {type(data_or_message)}: {data_or_message}"
+            )

airbyte_cdk/sources/utils/schema_helpers.py

@@ -10,11 +10,12 @@ import pkgutil
 from typing import Any, ClassVar, Dict, List, Mapping, MutableMapping, Optional, Tuple
 
 import jsonref
-from airbyte_cdk.models import ConnectorSpecification, FailureType
-from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 from jsonschema import RefResolver, validate
 from jsonschema.exceptions import ValidationError
-from pydantic import BaseModel, Field
+from pydantic.v1 import BaseModel, Field
+
+from airbyte_cdk.models import ConnectorSpecification, FailureType
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
 class JsonFileLoader:
@@ -74,7 +75,9 @@ def _expand_refs(schema: Any, ref_resolver: Optional[RefResolver] = None) -> Non
         if "$ref" in schema:
             ref_url = schema.pop("$ref")
             _, definition = ref_resolver.resolve(ref_url)
-            _expand_refs(definition, ref_resolver=ref_resolver)  # expand refs in definitions as well
+            _expand_refs(
+                definition, ref_resolver=ref_resolver
+            )  # expand refs in definitions as well
             schema.update(definition)
         else:
             for key, value in schema.items():
@@ -152,7 +155,9 @@ class ResourceSchemaLoader:
             base = os.path.dirname(package.__file__) + "/"
         else:
             raise ValueError(f"Package {package} does not have a valid __file__ field")
-        resolved = jsonref.JsonRef.replace_refs(raw_schema, loader=JsonFileLoader(base, "schemas/shared"), base_uri=base)
+        resolved = jsonref.JsonRef.replace_refs(
+            raw_schema, loader=JsonFileLoader(base, "schemas/shared"), base_uri=base
+        )
         resolved = resolve_ref_links(resolved)
         if isinstance(resolved, dict):
             return resolved
@@ -160,7 +165,9 @@ class ResourceSchemaLoader:
             raise ValueError(f"Expected resolved to be a dict. Got {resolved}")
 
 
-def check_config_against_spec_or_exit(config: Mapping[str, Any], spec: ConnectorSpecification) -> None:
+def check_config_against_spec_or_exit(
+    config: Mapping[str, Any], spec: ConnectorSpecification
+) -> None:
     """
     Check config object against spec. In case of spec is invalid, throws
     an exception with validation error description.