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

airbyte_cdk/sources/http_logger.py

@@ -5,11 +5,16 @@
 from typing import Optional, Union
 
 import requests
+
 from airbyte_cdk.sources.message import LogMessage
 
 
 def format_http_message(
-    response: requests.Response, title: str, description: str, stream_name: Optional[str], is_auxiliary: bool = None
+    response: requests.Response,
+    title: str,
+    description: str,
+    stream_name: Optional[str],
+    is_auxiliary: bool | None = None,
 ) -> LogMessage:
     request = response.request
     log_message = {
@@ -37,10 +42,10 @@ def format_http_message(
         "url": {"full": request.url},
     }
     if is_auxiliary is not None:
-        log_message["http"]["is_auxiliary"] = is_auxiliary
+        log_message["http"]["is_auxiliary"] = is_auxiliary  # type: ignore [index]
     if stream_name:
         log_message["airbyte_cdk"] = {"stream": {"name": stream_name}}
-    return log_message
+    return log_message  # type: ignore [return-value]  # got "dict[str, object]", expected "dict[str, JsonType]"
 
 
 def _normalize_body_string(body_str: Optional[Union[str, bytes]]) -> Optional[str]:

airbyte_cdk/sources/message/__init__.py

@@ -10,4 +10,10 @@ from .repository import (
     NoopMessageRepository,
 )
 
-__all__ = ["InMemoryMessageRepository", "LogAppenderMessageRepositoryDecorator", "LogMessage", "MessageRepository", "NoopMessageRepository"]
+__all__ = [
+    "InMemoryMessageRepository",
+    "LogAppenderMessageRepositoryDecorator",
+    "LogMessage",
+    "MessageRepository",
+    "NoopMessageRepository",
+]

airbyte_cdk/sources/message/repository.py

@@ -28,7 +28,9 @@ _SEVERITY_BY_LOG_LEVEL = {
 
 def _is_severe_enough(threshold: Level, level: Level) -> bool:
     if threshold not in _SEVERITY_BY_LOG_LEVEL:
-        _LOGGER.warning(f"Log level {threshold} for threshold is not supported. This is probably a CDK bug. Please contact Airbyte.")
+        _LOGGER.warning(
+            f"Log level {threshold} for threshold is not supported. This is probably a CDK bug. Please contact Airbyte."
+        )
         return True
 
     if level not in _SEVERITY_BY_LOG_LEVEL:
@@ -80,7 +82,12 @@ class InMemoryMessageRepository(MessageRepository):
     def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
         if _is_severe_enough(self._log_level, level):
             self.emit_message(
-                AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=level, message=filter_secrets(json.dumps(message_provider()))))
+                AirbyteMessage(
+                    type=Type.LOG,
+                    log=AirbyteLogMessage(
+                        level=level, message=filter_secrets(json.dumps(message_provider()))
+                    ),
+                )
             )
 
     def consume_queue(self) -> Iterable[AirbyteMessage]:
@@ -89,7 +96,12 @@ class InMemoryMessageRepository(MessageRepository):
 
 
 class LogAppenderMessageRepositoryDecorator(MessageRepository):
-    def __init__(self, dict_to_append: LogMessage, decorated: MessageRepository, log_level: Level = Level.INFO):
+    def __init__(
+        self,
+        dict_to_append: LogMessage,
+        decorated: MessageRepository,
+        log_level: Level = Level.INFO,
+    ):
         self._dict_to_append = dict_to_append
         self._decorated = decorated
         self._log_level = log_level
@@ -106,7 +118,9 @@ class LogAppenderMessageRepositoryDecorator(MessageRepository):
     def consume_queue(self) -> Iterable[AirbyteMessage]:
         return self._decorated.consume_queue()
 
-    def _append_second_to_first(self, first: LogMessage, second: LogMessage, path: Optional[List[str]] = None) -> LogMessage:
+    def _append_second_to_first(
+        self, first: LogMessage, second: LogMessage, path: Optional[List[str]] = None
+    ) -> LogMessage:
         if path is None:
             path = []
 

airbyte_cdk/sources/source.py

@@ -5,27 +5,41 @@
 
 import logging
 from abc import ABC, abstractmethod
-from collections import defaultdict
-from typing import Any, Dict, Generic, Iterable, List, Mapping, MutableMapping, Optional, TypeVar, Union
+from typing import Any, Generic, Iterable, List, Mapping, Optional, TypeVar
 
 from airbyte_cdk.connector import BaseConnector, DefaultConnectorMixin, TConfig
-from airbyte_cdk.models import AirbyteCatalog, AirbyteMessage, AirbyteStateMessage, AirbyteStateType, ConfiguredAirbyteCatalog
+from airbyte_cdk.models import (
+    AirbyteCatalog,
+    AirbyteMessage,
+    AirbyteStateMessage,
+    AirbyteStateMessageSerializer,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
+)
 
 TState = TypeVar("TState")
 TCatalog = TypeVar("TCatalog")
 
 
+class ExperimentalClassWarning(DeprecationWarning):
+    pass
+
+
 class BaseSource(BaseConnector[TConfig], ABC, Generic[TConfig, TState, TCatalog]):
     @abstractmethod
-    def read_state(self, state_path: str) -> TState:
-        ...
+    def read_state(self, state_path: str) -> TState: ...
 
     @abstractmethod
-    def read_catalog(self, catalog_path: str) -> TCatalog:
-        ...
+    def read_catalog(self, catalog_path: str) -> TCatalog: ...
 
     @abstractmethod
-    def read(self, logger: logging.Logger, config: TConfig, catalog: TCatalog, state: Optional[TState] = None) -> Iterable[AirbyteMessage]:
+    def read(
+        self,
+        logger: logging.Logger,
+        config: TConfig,
+        catalog: TCatalog,
+        state: Optional[TState] = None,
+    ) -> Iterable[AirbyteMessage]:
         """
         Returns a generator of the AirbyteMessages generated by reading the source with the given configuration, catalog, and state.
         """
@@ -40,12 +54,12 @@ class BaseSource(BaseConnector[TConfig], ABC, Generic[TConfig, TState, TCatalog]
 
 class Source(
     DefaultConnectorMixin,
-    BaseSource[Mapping[str, Any], Union[List[AirbyteStateMessage], MutableMapping[str, Any]], ConfiguredAirbyteCatalog],
+    BaseSource[Mapping[str, Any], List[AirbyteStateMessage], ConfiguredAirbyteCatalog],
     ABC,
 ):
     # can be overridden to change an input state.
     @classmethod
-    def read_state(cls, state_path: str) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
+    def read_state(cls, state_path: str) -> List[AirbyteStateMessage]:
         """
         Retrieves the input state of a sync by reading from the specified JSON file. Incoming state can be deserialized into either
         a JSON object for legacy state input or as a list of AirbyteStateMessages for the per-stream state format. Regardless of the
@@ -53,39 +67,29 @@ class Source(
         :param state_path: The filepath to where the stream states are located
         :return: The complete stream state based on the connector's previous sync
         """
+        parsed_state_messages = []
         if state_path:
             state_obj = BaseConnector._read_json_file(state_path)
-            if not state_obj:
-                return cls._emit_legacy_state_format({})
-            if isinstance(state_obj, List):
-                parsed_state_messages = []
+            if state_obj:
                 for state in state_obj:  # type: ignore  # `isinstance(state_obj, List)` ensures that this is a list
-                    parsed_message = AirbyteStateMessage.parse_obj(state)
-                    if not parsed_message.stream and not parsed_message.data and not parsed_message.global_:
-                        raise ValueError("AirbyteStateMessage should contain either a stream, global, or state field")
+                    parsed_message = AirbyteStateMessageSerializer.load(state)
+                    if (
+                        not parsed_message.stream
+                        and not parsed_message.data
+                        and not parsed_message.global_
+                    ):
+                        raise ValueError(
+                            "AirbyteStateMessage should contain either a stream, global, or state field"
+                        )
                     parsed_state_messages.append(parsed_message)
-                return parsed_state_messages
-            else:
-                return cls._emit_legacy_state_format(state_obj)  # type: ignore  # assuming it is a dict
-        return cls._emit_legacy_state_format({})
-
-    @classmethod
-    def _emit_legacy_state_format(cls, state_obj: Dict[str, Any]) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
-        """
-        Existing connectors that override read() might not be able to interpret the new state format. We temporarily
-        send state in the old format for these connectors, but once all have been upgraded, this method can be removed,
-        and we can then emit state in the list format.
-        """
-        # vars(self.__class__) checks if the current class directly overrides the read() function
-        if "read" in vars(cls):
-            return defaultdict(dict, state_obj)
-        else:
-            if state_obj:
-                return [AirbyteStateMessage(type=AirbyteStateType.LEGACY, data=state_obj)]
-            else:
-                return []
+        return parsed_state_messages
 
     # can be overridden to change an input catalog
     @classmethod
     def read_catalog(cls, catalog_path: str) -> ConfiguredAirbyteCatalog:
-        return ConfiguredAirbyteCatalog.parse_obj(cls._read_json_file(catalog_path))
+        return ConfiguredAirbyteCatalogSerializer.load(cls._read_json_file(catalog_path))
+
+    @property
+    def name(self) -> str:
+        """Source name"""
+        return self.__class__.__name__

airbyte_cdk/sources/streams/__init__.py

@@ -3,6 +3,6 @@
 #
 
 # Initialize Streams Package
-from .core import FULL_REFRESH_SENTINEL_STATE_KEY, IncrementalMixin, Stream
+from .core import NO_CURSOR_STATE_KEY, CheckpointMixin, IncrementalMixin, Stream
 
-__all__ = ["FULL_REFRESH_SENTINEL_STATE_KEY", "IncrementalMixin", "Stream"]
+__all__ = ["NO_CURSOR_STATE_KEY", "IncrementalMixin", "CheckpointMixin", "Stream"]

airbyte_cdk/sources/streams/availability_strategy.py

@@ -5,9 +5,10 @@
 import logging
 import typing
 from abc import ABC, abstractmethod
-from typing import Optional, Tuple
+from typing import Any, Mapping, Optional, Tuple
 
-from airbyte_cdk.sources.streams import Stream
+from airbyte_cdk.models import SyncMode
+from airbyte_cdk.sources.streams.core import Stream, StreamData
 
 if typing.TYPE_CHECKING:
     from airbyte_cdk.sources import Source
@@ -19,7 +20,9 @@ class AvailabilityStrategy(ABC):
     """
 
     @abstractmethod
-    def check_availability(self, stream: Stream, logger: logging.Logger, source: Optional["Source"]) -> Tuple[bool, Optional[str]]:
+    def check_availability(
+        self, stream: Stream, logger: logging.Logger, source: Optional["Source"] = None
+    ) -> Tuple[bool, Optional[str]]:
         """
         Checks stream availability.
 
@@ -31,3 +34,51 @@ class AvailabilityStrategy(ABC):
           for some reason and the str should describe what went wrong and how to
           resolve the unavailability, if possible.
         """
+
+    @staticmethod
+    def get_first_stream_slice(stream: Stream) -> Optional[Mapping[str, Any]]:
+        """
+        Gets the first stream_slice from a given stream's stream_slices.
+        :param stream: stream
+        :raises StopIteration: if there is no first slice to return (the stream_slices generator is empty)
+        :return: first stream slice from 'stream_slices' generator (`None` is a valid stream slice)
+        """
+        # We wrap the return output of stream_slices() because some implementations return types that are iterable,
+        # but not iterators such as lists or tuples
+        slices = iter(
+            stream.stream_slices(
+                cursor_field=stream.cursor_field,  # type: ignore[arg-type]
+                sync_mode=SyncMode.full_refresh,
+            )
+        )
+        return next(slices)
+
+    @staticmethod
+    def get_first_record_for_slice(
+        stream: Stream, stream_slice: Optional[Mapping[str, Any]]
+    ) -> StreamData:
+        """
+        Gets the first record for a stream_slice of a stream.
+
+        :param stream: stream instance from which to read records
+        :param stream_slice: stream_slice parameters for slicing the stream
+        :raises StopIteration: if there is no first record to return (the read_records generator is empty)
+        :return: StreamData containing the first record in the slice
+        """
+        # Store the original value of exit_on_rate_limit
+        original_exit_on_rate_limit = stream.exit_on_rate_limit
+
+        try:
+            # Ensure exit_on_rate_limit is safely set to True if possible
+            stream.exit_on_rate_limit = True
+
+            # We wrap the return output of read_records() because some implementations return types that are iterable,
+            # but not iterators such as lists or tuples
+            records_for_slice = iter(
+                stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice)
+            )
+
+            return next(records_for_slice)
+        finally:
+            # Restore the original exit_on_rate_limit value
+            stream.exit_on_rate_limit = original_exit_on_rate_limit

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
@@ -76,7 +75,9 @@ class AbstractCallRatePolicy(abc.ABC):
         """
 
     @abc.abstractmethod
-    def update(self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]) -> None:
+    def update(
+        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
+    ) -> None:
         """Update call rate counting with current values
 
         :param available_calls:
@@ -202,12 +203,20 @@ class UnlimitedCallRatePolicy(BaseCallRatePolicy):
     def try_acquire(self, request: Any, weight: int) -> None:
         """Do nothing"""
 
-    def update(self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]) -> None:
+    def update(
+        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
+    ) -> None:
         """Do nothing"""
 
 
 class FixedWindowCallRatePolicy(BaseCallRatePolicy):
-    def __init__(self, next_reset_ts: datetime.datetime, period: timedelta, call_limit: int, matchers: list[RequestMatcher]):
+    def __init__(
+        self,
+        next_reset_ts: datetime.datetime,
+        period: timedelta,
+        call_limit: int,
+        matchers: list[RequestMatcher],
+    ):
         """A policy that allows {call_limit} calls within a {period} time interval
 
         :param next_reset_ts: next call rate reset time point
@@ -235,7 +244,8 @@ class FixedWindowCallRatePolicy(BaseCallRatePolicy):
             if self._calls_num + weight > self._call_limit:
                 reset_in = self._next_reset_ts - datetime.datetime.now()
                 error_message = (
-                    f"reached maximum number of allowed calls {self._call_limit} " f"per {self._offset} interval, next reset in {reset_in}."
+                    f"reached maximum number of allowed calls {self._call_limit} "
+                    f"per {self._offset} interval, next reset in {reset_in}."
                 )
                 raise CallRateLimitHit(
                     error=error_message,
@@ -247,7 +257,9 @@ class FixedWindowCallRatePolicy(BaseCallRatePolicy):
 
             self._calls_num += weight
 
-    def update(self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]) -> None:
+    def update(
+        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
+    ) -> None:
         """Update call rate counters, by default, only reacts to decreasing updates of available_calls and changes to call_reset_ts.
         We ignore updates with available_calls > current_available_calls to support call rate limits that are lower than API limits.
 
@@ -260,12 +272,18 @@ class FixedWindowCallRatePolicy(BaseCallRatePolicy):
 
             if available_calls is not None and current_available_calls > available_calls:
                 logger.debug(
-                    "got rate limit update from api, adjusting available calls from %s to %s", current_available_calls, available_calls
+                    "got rate limit update from api, adjusting available calls from %s to %s",
+                    current_available_calls,
+                    available_calls,
                 )
                 self._calls_num = self._call_limit - available_calls
 
             if call_reset_ts is not None and call_reset_ts != self._next_reset_ts:
-                logger.debug("got rate limit update from api, adjusting reset time from %s to %s", self._next_reset_ts, call_reset_ts)
+                logger.debug(
+                    "got rate limit update from api, adjusting reset time from %s to %s",
+                    self._next_reset_ts,
+                    call_reset_ts,
+                )
                 self._next_reset_ts = call_reset_ts
 
     def _update_current_window(self) -> None:
@@ -292,7 +310,10 @@ class MovingWindowCallRatePolicy(BaseCallRatePolicy):
         """
         if not rates:
             raise ValueError("The list of rates can not be empty")
-        pyrate_rates = [PyRateRate(limit=rate.limit, interval=int(rate.interval.total_seconds() * 1000)) for rate in rates]
+        pyrate_rates = [
+            PyRateRate(limit=rate.limit, interval=int(rate.interval.total_seconds() * 1000))
+            for rate in rates
+        ]
         self._bucket = InMemoryBucket(pyrate_rates)
         # Limiter will create the background task that clears old requests in the bucket
         self._limiter = Limiter(self._bucket)
@@ -320,14 +341,18 @@ class MovingWindowCallRatePolicy(BaseCallRatePolicy):
                     time_to_wait=timedelta(milliseconds=time_to_wait),
                 )
 
-    def update(self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]) -> None:
+    def update(
+        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
+    ) -> None:
         """Adjust call bucket to reflect the state of the API server
 
         :param available_calls:
         :param call_reset_ts:
         :return:
         """
-        if available_calls is not None and call_reset_ts is None:  # we do our best to sync buckets with API
+        if (
+            available_calls is not None and call_reset_ts is None
+        ):  # we do our best to sync buckets with API
             if available_calls == 0:
                 with self._limiter.lock:
                     items_to_add = self._bucket.count() < self._bucket.rates[0].limit
@@ -350,7 +375,9 @@ class AbstractAPIBudget(abc.ABC):
     """
 
     @abc.abstractmethod
-    def acquire_call(self, request: Any, block: bool = True, timeout: Optional[float] = None) -> None:
+    def acquire_call(
+        self, request: Any, block: bool = True, timeout: Optional[float] = None
+    ) -> None:
         """Try to get a call from budget, will block by default
 
         :param request:
@@ -375,7 +402,9 @@ class AbstractAPIBudget(abc.ABC):
 class APIBudget(AbstractAPIBudget):
     """Default APIBudget implementation"""
 
-    def __init__(self, policies: list[AbstractCallRatePolicy], maximum_attempts_to_acquire: int = 100000) -> None:
+    def __init__(
+        self, policies: list[AbstractCallRatePolicy], maximum_attempts_to_acquire: int = 100000
+    ) -> None:
         """Constructor
 
         :param policies: list of policies in this budget
@@ -392,7 +421,9 @@ class APIBudget(AbstractAPIBudget):
                 return policy
         return None
 
-    def acquire_call(self, request: Any, block: bool = True, timeout: Optional[float] = None) -> None:
+    def acquire_call(
+        self, request: Any, block: bool = True, timeout: Optional[float] = None
+    ) -> None:
         """Try to get a call from budget, will block by default.
         Matchers will be called sequentially in the same order they were added.
         The first matcher that returns True will
@@ -417,7 +448,9 @@ class APIBudget(AbstractAPIBudget):
         """
         pass
 
-    def _do_acquire(self, request: Any, policy: AbstractCallRatePolicy, block: bool, timeout: Optional[float]) -> None:
+    def _do_acquire(
+        self, request: Any, policy: AbstractCallRatePolicy, block: bool, timeout: Optional[float]
+    ) -> None:
         """Internal method to try to acquire a call credit
 
         :param request:
@@ -439,14 +472,20 @@ class APIBudget(AbstractAPIBudget):
                     else:
                         time_to_wait = exc.time_to_wait
 
-                    time_to_wait = max(timedelta(0), time_to_wait)  # sometimes we get negative duration
-                    logger.info("reached call limit %s. going to sleep for %s", exc.rate, time_to_wait)
+                    time_to_wait = max(
+                        timedelta(0), time_to_wait
+                    )  # sometimes we get negative duration
+                    logger.info(
+                        "reached call limit %s. going to sleep for %s", exc.rate, time_to_wait
+                    )
                     time.sleep(time_to_wait.total_seconds())
                 else:
                     raise
 
         if last_exception:
-            logger.info("we used all %s attempts to acquire and failed", self._maximum_attempts_to_acquire)
+            logger.info(
+                "we used all %s attempts to acquire and failed", self._maximum_attempts_to_acquire
+            )
             raise last_exception
 
 
@@ -481,9 +520,13 @@ class HttpAPIBudget(APIBudget):
             reset_ts = self.get_reset_ts_from_response(response)
             policy.update(available_calls=available_calls, call_reset_ts=reset_ts)
 
-    def get_reset_ts_from_response(self, response: requests.Response) -> Optional[datetime.datetime]:
+    def get_reset_ts_from_response(
+        self, response: requests.Response
+    ) -> Optional[datetime.datetime]:
         if response.headers.get(self._ratelimit_reset_header):
-            return datetime.datetime.fromtimestamp(int(response.headers[self._ratelimit_reset_header]))
+            return datetime.datetime.fromtimestamp(
+                int(response.headers[self._ratelimit_reset_header])
+            )
         return None
 
     def get_calls_left_from_response(self, response: requests.Response) -> Optional[int]:

airbyte_cdk/sources/streams/checkpoint/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+
+from .checkpoint_reader import (
+    CheckpointMode,
+    CheckpointReader,
+    CursorBasedCheckpointReader,
+    FullRefreshCheckpointReader,
+    IncrementalCheckpointReader,
+    LegacyCursorBasedCheckpointReader,
+    ResumableFullRefreshCheckpointReader,
+)
+from .cursor import Cursor
+from .resumable_full_refresh_cursor import ResumableFullRefreshCursor
+
+__all__ = [
+    "CheckpointMode",
+    "CheckpointReader",
+    "Cursor",
+    "CursorBasedCheckpointReader",
+    "FullRefreshCheckpointReader",
+    "IncrementalCheckpointReader",
+    "LegacyCursorBasedCheckpointReader",
+    "ResumableFullRefreshCheckpointReader",
+    "ResumableFullRefreshCursor",
+]