airbyte-cdk 6.7.1rc4__py3-none-any.whl → 6.7.2__py3-none-any.whl

airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_concurrent_cursor.py

@@ -19,7 +19,7 @@ from airbyte_cdk.sources.file_based.types import StreamState
 from airbyte_cdk.sources.message.repository import MessageRepository
 from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.types import Record
 
 if TYPE_CHECKING:
     from airbyte_cdk.sources.file_based.stream.concurrent.adapters import FileBasedStreamPartition

airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_final_state_cursor.py

@@ -16,7 +16,7 @@ from airbyte_cdk.sources.file_based.types import StreamState
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.types import Record
 
 if TYPE_CHECKING:
     from airbyte_cdk.sources.file_based.stream.concurrent.adapters import FileBasedStreamPartition

airbyte_cdk/sources/http_logger.py

@@ -5,6 +5,7 @@
 from typing import Optional, Union
 
 import requests
+
 from airbyte_cdk.sources.message import LogMessage
 
 

airbyte_cdk/sources/streams/call_rate.py

@@ -14,9 +14,8 @@ from urllib import parse
 
 import requests
 import requests_cache
-from pyrate_limiter import InMemoryBucket, Limiter
+from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock
 from pyrate_limiter import Rate as PyRateRate
-from pyrate_limiter import RateItem, TimeClock
 from pyrate_limiter.exceptions import BucketFullException
 
 # prevents mypy from complaining about missing session attributes in LimiterMixin

airbyte_cdk/sources/streams/concurrent/abstract_stream.py

@@ -5,12 +5,13 @@
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, Mapping, Optional
 
+from deprecated.classic import deprecated
+
 from airbyte_cdk.models import AirbyteStream
 from airbyte_cdk.sources.source import ExperimentalClassWarning
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import StreamAvailability
 from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from deprecated.classic import deprecated
 
 
 @deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)

airbyte_cdk/sources/streams/concurrent/adapters.py

@@ -8,6 +8,8 @@ import logging
 from functools import lru_cache
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
 
+from deprecated.classic import deprecated
+
 from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
@@ -37,12 +39,10 @@ from airbyte_cdk.sources.streams.concurrent.helpers import (
 )
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_cdk.sources.streams.core import StreamData
+from airbyte_cdk.sources.types import Record
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig
 from airbyte_cdk.sources.utils.slice_logger import SliceLogger
-from deprecated.classic import deprecated
-
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
 """
@@ -294,7 +294,11 @@ class StreamPartition(Partition):
                     self._stream.transformer.transform(
                         data_to_return, self._stream.get_json_schema()
                     )
-                    yield Record(data_to_return, self)
+                    yield Record(
+                        data=data_to_return,
+                        stream_name=self.stream_name(),
+                        associated_slice=self._slice,
+                    )
                 else:
                     self._message_repository.emit_message(record_data)
         except Exception as e:

airbyte_cdk/sources/streams/concurrent/availability_strategy.py

@@ -6,9 +6,10 @@ import logging
 from abc import ABC, abstractmethod
 from typing import Optional
 
-from airbyte_cdk.sources.source import ExperimentalClassWarning
 from deprecated.classic import deprecated
 
+from airbyte_cdk.sources.source import ExperimentalClassWarning
+
 
 class StreamAvailability(ABC):
     @abstractmethod

airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -3,6 +3,7 @@
 #
 
 import functools
+import logging
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional, Protocol, Tuple
 
@@ -10,12 +11,13 @@ from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_cdk.sources.streams.concurrent.partitions.stream_slicer import StreamSlicer
 from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import (
     AbstractStreamStateConverter,
 )
-from airbyte_cdk.sources.types import StreamSlice
+from airbyte_cdk.sources.types import Record, StreamSlice
+
+LOGGER = logging.getLogger("airbyte")
 
 
 def _extract_value(mapping: Mapping[str, Any], path: List[str]) -> Any:
@@ -173,9 +175,11 @@ class ConcurrentCursor(Cursor):
         self.start, self._concurrent_state = self._get_concurrent_state(stream_state)
         self._lookback_window = lookback_window
         self._slice_range = slice_range
-        self._most_recent_cursor_value_per_partition: MutableMapping[Partition, Any] = {}
+        self._most_recent_cursor_value_per_partition: MutableMapping[StreamSlice, Any] = {}
         self._has_closed_at_least_one_slice = False
         self._cursor_granularity = cursor_granularity
+        # Flag to track if the logger has been triggered (per stream)
+        self._should_be_synced_logger_triggered = False
 
     @property
     def state(self) -> MutableMapping[str, Any]:
@@ -210,12 +214,12 @@ class ConcurrentCursor(Cursor):
 
     def observe(self, record: Record) -> None:
         most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
-            record.partition
+            record.associated_slice
         )
         cursor_value = self._extract_cursor_value(record)
 
         if most_recent_cursor_value is None or most_recent_cursor_value < cursor_value:
-            self._most_recent_cursor_value_per_partition[record.partition] = cursor_value
+            self._most_recent_cursor_value_per_partition[record.associated_slice] = cursor_value
 
     def _extract_cursor_value(self, record: Record) -> Any:
         return self._connector_state_converter.parse_value(self._cursor_field.extract_value(record))
@@ -231,7 +235,9 @@ class ConcurrentCursor(Cursor):
         self._has_closed_at_least_one_slice = True
 
     def _add_slice_to_state(self, partition: Partition) -> None:
-        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(partition)
+        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
+            partition.to_slice()
+        )
 
         if self._slice_boundary_fields:
             if "slices" not in self.state:
@@ -442,3 +448,21 @@ class ConcurrentCursor(Cursor):
             return lower + step
         except OverflowError:
             return self._end_provider()
+
+    def should_be_synced(self, record: Record) -> bool:
+        """
+        Determines if a record should be synced based on its cursor value.
+        :param record: The record to evaluate
+
+        :return: True if the record's cursor value falls within the sync boundaries
+        """
+        try:
+            record_cursor_value: CursorValueType = self._extract_cursor_value(record)  # type: ignore  # cursor_field is converted to an InterpolatedString in __post_init__
+        except ValueError:
+            if not self._should_be_synced_logger_triggered:
+                LOGGER.warning(
+                    f"Could not find cursor field `{self.cursor_field.cursor_field_key}` in record. The incremental sync will assume it needs to be synced"
+                )
+                self._should_be_synced_logger_triggered = True
+            return True
+        return self.start <= record_cursor_value <= self._end_provider()

airbyte_cdk/sources/streams/concurrent/partitions/partition.py

@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, Mapping, Optional
 
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.types import Record
 
 
 class Partition(ABC):

airbyte_cdk/sources/streams/concurrent/partitions/types.py

@@ -8,7 +8,7 @@ from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentin
     PartitionGenerationCompletedSentinel,
 )
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.types import Record
 
 
 class PartitionCompleteSentinel:

airbyte_cdk/sources/streams/concurrent/state_converters/datetime_stream_state_converter.py

@@ -7,6 +7,7 @@ from datetime import datetime, timedelta, timezone
 from typing import Any, Callable, List, MutableMapping, Optional, Tuple
 
 import pendulum
+from pendulum.datetime import DateTime
 
 # FIXME We would eventually like the Concurrent package do be agnostic of the declarative package. However, this is a breaking change and
 #  the goal in the short term is only to fix the issue we are seeing for source-declarative-manifest.
@@ -16,7 +17,6 @@ from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_sta
     AbstractStreamStateConverter,
     ConcurrencyCompatibleStateType,
 )
-from pendulum.datetime import DateTime
 
 
 class DateTimeStreamStateConverter(AbstractStreamStateConverter):

airbyte_cdk/sources/streams/core.py

@@ -10,6 +10,8 @@ from dataclasses import dataclass
 from functools import cached_property, lru_cache
 from typing import Any, Dict, Iterable, Iterator, List, Mapping, MutableMapping, Optional, Union
 
+from deprecated import deprecated
+
 import airbyte_cdk.sources.utils.casing as casing
 from airbyte_cdk.models import (
     AirbyteMessage,
@@ -35,7 +37,6 @@ from airbyte_cdk.sources.types import StreamSlice
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig, ResourceSchemaLoader
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger, SliceLogger
 from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
-from deprecated import deprecated
 
 # A stream's read method can return one of the following types:
 # Mapping[str, Any]: The content of an AirbyteRecordMessage

airbyte_cdk/sources/streams/http/error_handlers/default_error_mapping.py

@@ -4,12 +4,13 @@
 
 from typing import Mapping, Type, Union
 
+from requests.exceptions import InvalidSchema, InvalidURL, RequestException
+
 from airbyte_cdk.models import FailureType
 from airbyte_cdk.sources.streams.http.error_handlers.response_models import (
     ErrorResolution,
     ResponseAction,
 )
-from requests.exceptions import InvalidSchema, InvalidURL, RequestException
 
 DEFAULT_ERROR_MAPPING: Mapping[Union[int, str, Type[Exception]], ErrorResolution] = {
     InvalidSchema: ErrorResolution(

airbyte_cdk/sources/streams/http/error_handlers/http_status_error_handler.py

@@ -7,6 +7,7 @@ from datetime import timedelta
 from typing import Mapping, Optional, Union
 
 import requests
+
 from airbyte_cdk.models import FailureType
 from airbyte_cdk.sources.streams.http.error_handlers.default_error_mapping import (
     DEFAULT_ERROR_MAPPING,

airbyte_cdk/sources/streams/http/error_handlers/json_error_message_parser.py

@@ -5,6 +5,7 @@
 from typing import Optional
 
 import requests
+
 from airbyte_cdk.sources.streams.http.error_handlers import ErrorMessageParser
 from airbyte_cdk.sources.utils.types import JsonType
 

airbyte_cdk/sources/streams/http/error_handlers/response_models.py

@@ -5,9 +5,10 @@ from enum import Enum
 from typing import Optional, Union
 
 import requests
+from requests import HTTPError
+
 from airbyte_cdk.models import FailureType
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
-from requests import HTTPError
 
 
 class ResponseAction(Enum):

airbyte_cdk/sources/streams/http/http.py

@@ -9,6 +9,9 @@ from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optio
 from urllib.parse import urljoin
 
 import requests
+from deprecated import deprecated
+from requests.auth import AuthBase
+
 from airbyte_cdk.models import AirbyteMessage, FailureType, SyncMode
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.message.repository import InMemoryMessageRepository
@@ -33,8 +36,6 @@ from airbyte_cdk.sources.streams.http.error_handlers.response_models import (
 from airbyte_cdk.sources.streams.http.http_client import HttpClient
 from airbyte_cdk.sources.types import Record, StreamSlice
 from airbyte_cdk.sources.utils.types import JsonType
-from deprecated import deprecated
-from requests.auth import AuthBase
 
 # list of all possible HTTP methods which can be used for sending of request bodies
 BODY_REQUEST_METHODS = ("GET", "POST", "PUT", "PATCH")

airbyte_cdk/sources/streams/http/http_client.py

@@ -6,11 +6,13 @@ import logging
 import os
 import urllib
 from pathlib import Path
-from typing import Any, Callable, List, Mapping, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
 
 import orjson
 import requests
 import requests_cache
+from requests.auth import AuthBase
+
 from airbyte_cdk.models import (
     AirbyteMessageSerializer,
     AirbyteStreamStatus,
@@ -43,13 +45,13 @@ from airbyte_cdk.sources.streams.http.rate_limiting import (
     rate_limit_default_backoff_handler,
     user_defined_backoff_handler,
 )
+from airbyte_cdk.sources.utils.types import JsonType
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH
 from airbyte_cdk.utils.stream_status_utils import (
     as_airbyte_message as stream_status_as_airbyte_message,
 )
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
-from requests.auth import AuthBase
 
 BODY_REQUEST_METHODS = ("GET", "POST", "PUT", "PATCH")
 
@@ -74,6 +76,7 @@ class MessageRepresentationAirbyteTracedErrors(AirbyteTracedException):
 class HttpClient:
     _DEFAULT_MAX_RETRY: int = 5
     _DEFAULT_MAX_TIME: int = 60 * 10
+    _ACTIONS_TO_RETRY_ON = {ResponseAction.RETRY, ResponseAction.RATE_LIMITED}
 
     def __init__(
         self,
@@ -91,11 +94,9 @@ class HttpClient:
     ):
         self._name = name
         self._api_budget: APIBudget = api_budget or APIBudget(policies=[])
-        self._is_session_owner = False
         if session:
             self._session = session
         else:
-            self._is_session_owner = True
             self._use_cache = use_cache
             self._session = self._request_session()
             self._session.mount(
@@ -116,13 +117,10 @@ class HttpClient:
         else:
             self._backoff_strategies = [DefaultBackoffStrategy()]
         self._error_message_parser = error_message_parser or JsonErrorMessageParser()
+        self._request_attempt_count: Dict[requests.PreparedRequest, int] = {}
         self._disable_retries = disable_retries
         self._message_repository = message_repository
 
-    def __del__(self):
-        if self._is_session_owner:
-            self._session.close()
-
     @property
     def cache_filename(self) -> str:
         """
@@ -275,6 +273,13 @@ class HttpClient:
         log_formatter: Optional[Callable[[requests.Response], Any]] = None,
         exit_on_rate_limit: Optional[bool] = False,
     ) -> requests.Response:
+        if request not in self._request_attempt_count:
+            self._request_attempt_count[request] = 1
+        else:
+            self._request_attempt_count[request] += 1
+            if hasattr(self._session, "auth") and isinstance(self._session.auth, AuthBase):
+                self._session.auth(request)
+
         self._logger.debug(
             "Making outbound API request",
             extra={"headers": request.headers, "url": request.url, "request_body": request.body},
@@ -332,6 +337,40 @@ class HttpClient:
 
         return response  # type: ignore # will either return a valid response of type requests.Response or raise an exception
 
+    def _get_response_body(self, response: requests.Response) -> Optional[JsonType]:
+        """
+        Extracts and returns the body of an HTTP response.
+
+        This method attempts to parse the response body as JSON. If the response
+        body is not valid JSON, it falls back to decoding the response content
+        as a UTF-8 string. If both attempts fail, it returns None.
+
+        Args:
+            response (requests.Response): The HTTP response object.
+
+        Returns:
+            Optional[JsonType]: The parsed JSON object as a string, the decoded
+            response content as a string, or None if both parsing attempts fail.
+        """
+        try:
+            return str(response.json())
+        except requests.exceptions.JSONDecodeError:
+            try:
+                return response.content.decode("utf-8")
+            except Exception:
+                return "The Content of the Response couldn't be decoded."
+
+    def _evict_key(self, prepared_request: requests.PreparedRequest) -> None:
+        """
+        Addresses high memory consumption when enabling concurrency in https://github.com/airbytehq/oncall/issues/6821.
+
+        The `_request_attempt_count` attribute keeps growing as multiple requests are made using the same `http_client`.
+        To mitigate this issue, we evict keys for completed requests once we confirm that no further retries are needed.
+        This helps manage memory usage more efficiently while maintaining the necessary logic for retry attempts.
+        """
+        if prepared_request in self._request_attempt_count:
+            del self._request_attempt_count[prepared_request]
+
     def _handle_error_resolution(
         self,
         response: Optional[requests.Response],
@@ -340,6 +379,9 @@ class HttpClient:
         error_resolution: ErrorResolution,
         exit_on_rate_limit: Optional[bool] = False,
     ) -> None:
+        if error_resolution.response_action not in self._ACTIONS_TO_RETRY_ON:
+            self._evict_key(request)
+
         # Emit stream status RUNNING with the reason RATE_LIMITED to log that the rate limit has been reached
         if error_resolution.response_action == ResponseAction.RATE_LIMITED:
             # TODO: Update to handle with message repository when concurrent message repository is ready
@@ -360,13 +402,18 @@ class HttpClient:
 
         if error_resolution.response_action == ResponseAction.FAIL:
             if response is not None:
-                error_message = f"'{request.method}' request to '{request.url}' failed with status code '{response.status_code}' and error message '{response.content}'"
+                filtered_response_message = filter_secrets(
+                    f"Request (body): '{str(request.body)}'. Response (body): '{self._get_response_body(response)}'. Response (headers): '{response.headers}'."
+                )
+                error_message = f"'{request.method}' request to '{request.url}' failed with status code '{response.status_code}' and error message: '{self._error_message_parser.parse_response_error_message(response)}'. {filtered_response_message}"
             else:
                 error_message = (
                     f"'{request.method}' request to '{request.url}' failed with exception: '{exc}'"
                 )
 
-            self._logger.warning(filter_secrets(error_message))
+            # ensure the exception message is emitted before raised
+            self._logger.error(error_message)
+
             raise MessageRepresentationAirbyteTracedErrors(
                 internal_message=error_message,
                 message=error_resolution.error_message or error_message,
@@ -390,7 +437,7 @@ class HttpClient:
             for backoff_strategy in self._backoff_strategies:
                 backoff_time = backoff_strategy.backoff_time(
                     response_or_exception=response if response is not None else exc,
-                    attempt_count=0,
+                    attempt_count=self._request_attempt_count[request],
                 )
                 if backoff_time:
                     user_defined_backoff_time = backoff_time

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
 

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

@@ -6,6 +6,7 @@ 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,

airbyte_cdk/sources/types.py

@@ -6,6 +6,8 @@ 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]
@@ -15,9 +17,17 @@ StreamState = Mapping[str, Any]
 
 
 class Record(Mapping[str, Any]):
-    def __init__(self, data: Mapping[str, Any], associated_slice: Optional[StreamSlice]):
+    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]:
@@ -139,3 +149,6 @@ class StreamSlice(Mapping[str, Any]):
 
     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/schema_helpers.py

@@ -10,12 +10,13 @@ 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.v1 import BaseModel, Field
 
+from airbyte_cdk.models import ConnectorSpecification, FailureType
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
 
 class JsonFileLoader:
     """

airbyte_cdk/sql/secrets.py

@@ -6,9 +6,10 @@ from __future__ import annotations
 import json
 from typing import TYPE_CHECKING, Any
 
-from airbyte_cdk.sql import exceptions as exc
 from pydantic_core import CoreSchema, core_schema
 
+from airbyte_cdk.sql import exceptions as exc
+
 if TYPE_CHECKING:
     from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler, ValidationInfo
     from pydantic.json_schema import JsonSchemaValue

airbyte_cdk/sql/shared/sql_processor.py

@@ -13,6 +13,12 @@ from typing import TYPE_CHECKING, Any, final
 import pandas as pd
 import sqlalchemy
 import ulid
+from airbyte_protocol_dataclasses.models import AirbyteStateMessage
+from pandas import Index
+from pydantic import BaseModel, Field
+from sqlalchemy import Column, Table, and_, create_engine, insert, null, select, text, update
+from sqlalchemy.exc import ProgrammingError, SQLAlchemyError
+
 from airbyte_cdk.sql import exceptions as exc
 from airbyte_cdk.sql._util.hashing import one_way_hash
 from airbyte_cdk.sql._util.name_normalizers import LowerCaseNormalizer
@@ -24,16 +30,10 @@ from airbyte_cdk.sql.constants import (
 )
 from airbyte_cdk.sql.secrets import SecretString
 from airbyte_cdk.sql.types import SQLTypeConverter
-from airbyte_protocol_dataclasses.models import AirbyteStateMessage
-from pandas import Index
-from pydantic import BaseModel, Field
-from sqlalchemy import Column, Table, and_, create_engine, insert, null, select, text, update
-from sqlalchemy.exc import ProgrammingError, SQLAlchemyError
 
 if TYPE_CHECKING:
     from collections.abc import Generator
 
-    from airbyte_cdk.sql.shared.catalog_providers import CatalogProvider
     from sqlalchemy.engine import Connection, Engine
     from sqlalchemy.engine.cursor import CursorResult
     from sqlalchemy.engine.reflection import Inspector
@@ -41,6 +41,8 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.elements import TextClause
     from sqlalchemy.sql.type_api import TypeEngine
 
+    from airbyte_cdk.sql.shared.catalog_providers import CatalogProvider
+
 
 class SQLRuntimeError(Exception):
     """Raised when an SQL operation fails."""

airbyte_cdk/test/entrypoint_wrapper.py

@@ -23,6 +23,10 @@ from io import StringIO
 from pathlib import Path
 from typing import Any, List, Mapping, Optional, Union
 
+from orjson import orjson
+from pydantic import ValidationError as V2ValidationError
+from serpyco_rs import SchemaValidationError
+
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
 from airbyte_cdk.exception_handler import assemble_uncaught_exception
 from airbyte_cdk.logger import AirbyteLogFormatter
@@ -40,9 +44,6 @@ from airbyte_cdk.models import (
     Type,
 )
 from airbyte_cdk.sources import Source
-from orjson import orjson
-from pydantic import ValidationError as V2ValidationError
-from serpyco_rs import SchemaValidationError
 
 
 class EntrypointOutput:

airbyte_cdk/test/mock_http/mocker.py

@@ -7,6 +7,7 @@ from types import TracebackType
 from typing import Callable, List, Optional, Union
 
 import requests_mock
+
 from airbyte_cdk.test.mock_http import HttpRequest, HttpRequestMatcher, HttpResponse
 
 

airbyte_cdk/utils/schema_inferrer.py

@@ -5,11 +5,12 @@
 from collections import defaultdict
 from typing import Any, Dict, List, Mapping, Optional
 
-from airbyte_cdk.models import AirbyteRecordMessage
 from genson import SchemaBuilder, SchemaNode
 from genson.schema.strategies.object import Object
 from genson.schema.strategies.scalar import Number
 
+from airbyte_cdk.models import AirbyteRecordMessage
+
 # schema keywords
 _TYPE = "type"
 _NULL_TYPE = "null"

airbyte_cdk/utils/slice_hasher.py

@@ -1,6 +1,6 @@
 import hashlib
 import json
-from typing import Any, Mapping, Optional, Final
+from typing import Any, Final, Mapping, Optional
 
 
 class SliceEncoder(json.JSONEncoder):