airbyte-cdk 6.31.2.dev0__py3-none-any.whl → 6.33.0__py3-none-any.whl

airbyte_cdk/sources/declarative/partition_routers/async_job_partition_router.py

@@ -4,9 +4,9 @@ from dataclasses import InitVar, dataclass, field
 from typing import Any, Callable, Iterable, Mapping, Optional
 
 from airbyte_cdk.models import FailureType
+from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
 from airbyte_cdk.sources.declarative.async_job.job_orchestrator import (
     AsyncJobOrchestrator,
-    AsyncPartition,
 )
 from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
     SinglePartitionRouter,
@@ -42,12 +42,12 @@ class AsyncJobPartitionRouter(StreamSlicer):
 
         for completed_partition in self._job_orchestrator.create_and_get_completed_partitions():
             yield StreamSlice(
-                partition=dict(completed_partition.stream_slice.partition)
-                | {"partition": completed_partition},
+                partition=dict(completed_partition.stream_slice.partition),
                 cursor_slice=completed_partition.stream_slice.cursor_slice,
+                extra_fields={"jobs": list(completed_partition.jobs)},
             )
 
-    def fetch_records(self, partition: AsyncPartition) -> Iterable[Mapping[str, Any]]:
+    def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
         """
         This method of fetching records extends beyond what a PartitionRouter/StreamSlicer should
         be responsible for. However, this was added in because the JobOrchestrator is required to
@@ -62,4 +62,4 @@ class AsyncJobPartitionRouter(StreamSlicer):
                 failure_type=FailureType.system_error,
             )
 
-        return self._job_orchestrator.fetch_records(partition=partition)
+        return self._job_orchestrator.fetch_records(async_jobs=async_jobs)

airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py

@@ -289,24 +289,58 @@ class SubstreamPartitionRouter(PartitionRouter):
             return
 
         if not parent_state and incremental_dependency:
-            # Attempt to retrieve child state
-            substream_state = list(stream_state.values())
-            substream_state = substream_state[0] if substream_state else {}  # type: ignore [assignment]  # Incorrect type for assignment
-            parent_state = {}
-
-            # Copy child state to parent streams with incremental dependencies
-            if substream_state:
-                for parent_config in self.parent_stream_configs:
-                    if parent_config.incremental_dependency:
-                        parent_state[parent_config.stream.name] = {
-                            parent_config.stream.cursor_field: substream_state
-                        }
+            # Migrate child state to parent state format
+            parent_state = self._migrate_child_state_to_parent_state(stream_state)
 
         # Set state for each parent stream with an incremental dependency
         for parent_config in self.parent_stream_configs:
             if parent_config.incremental_dependency:
                 parent_config.stream.state = parent_state.get(parent_config.stream.name, {})
 
+    def _migrate_child_state_to_parent_state(self, stream_state: StreamState) -> StreamState:
+        """
+        Migrate the child stream state to the parent stream's state format.
+
+        This method converts the global or child state into a format compatible with parent
+        streams. The migration occurs only for parent streams with incremental dependencies.
+        The method filters out per-partition states and retains only the global state in the
+        format `{cursor_field: cursor_value}`.
+
+        Args:
+            stream_state (StreamState): The state to migrate. Expected formats include:
+                - {"updated_at": "2023-05-27T00:00:00Z"}
+                - {"states": [...] } (ignored during migration)
+
+        Returns:
+            StreamState: A migrated state for parent streams in the format:
+                {
+                    "parent_stream_name": {"parent_stream_cursor": "2023-05-27T00:00:00Z"}
+                }
+
+        Example:
+            Input: {"updated_at": "2023-05-27T00:00:00Z"}
+            Output: {
+                "parent_stream_name": {"parent_stream_cursor": "2023-05-27T00:00:00Z"}
+            }
+        """
+        substream_state_values = list(stream_state.values())
+        substream_state = substream_state_values[0] if substream_state_values else {}
+
+        # Ignore per-partition states or invalid formats
+        if isinstance(substream_state, (list, dict)) or len(substream_state_values) != 1:
+            return {}
+
+        # Copy child state to parent streams with incremental dependencies
+        parent_state = {}
+        if substream_state:
+            for parent_config in self.parent_stream_configs:
+                if parent_config.incremental_dependency:
+                    parent_state[parent_config.stream.name] = {
+                        parent_config.stream.cursor_field: substream_state
+                    }
+
+        return parent_state
+
     def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
         """
         Get the state of the parent streams.

airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py

@@ -8,6 +8,7 @@ from typing import Any, List, Mapping, Optional, Union
 import requests
 
 from airbyte_cdk.sources.streams.http.error_handlers import ErrorHandler
+from airbyte_cdk.sources.streams.http.error_handlers.backoff_strategy import BackoffStrategy
 from airbyte_cdk.sources.streams.http.error_handlers.response_models import (
     ErrorResolution,
     ResponseAction,
@@ -77,3 +78,24 @@ class CompositeErrorHandler(ErrorHandler):
             return matched_error_resolution
 
         return create_fallback_error_resolution(response_or_exception)
+
+    @property
+    def backoff_strategies(self) -> Optional[List[BackoffStrategy]]:
+        """
+        Combines backoff strategies from all child error handlers into a single flattened list.
+
+        When used with HttpRequester, note the following behavior:
+        - In HttpRequester.__post_init__, the entire list of backoff strategies is assigned to the error handler
+        - However, the error handler's backoff_time() method only ever uses the first non-None strategy in the list
+        - This means that if any backoff strategies are present, the first non-None strategy becomes the default
+        - This applies to both user-defined response filters and errors from DEFAULT_ERROR_MAPPING
+        - The list structure is not used to map different strategies to different error conditions
+        - Therefore, subsequent strategies in the list will not be used
+
+        Returns None if no handlers have strategies defined, which will result in HttpRequester using its default backoff strategy.
+        """
+        all_strategies = []
+        for handler in self.error_handlers:
+            if hasattr(handler, "backoff_strategies") and handler.backoff_strategies:
+                all_strategies.extend(handler.backoff_strategies)
+        return all_strategies if all_strategies else None

airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py

@@ -151,16 +151,16 @@ class HttpResponseFilter:
         :param response: The HTTP response which can be used during interpolation
         :return: The evaluated error message string to be emitted
         """
-        return self.error_message.eval(  # type: ignore [no-any-return, union-attr]
+        return self.error_message.eval(  # type: ignore[no-any-return, union-attr]
             self.config, response=self._safe_response_json(response), headers=response.headers
         )
 
     def _response_matches_predicate(self, response: requests.Response) -> bool:
         return (
             bool(
-                self.predicate.condition  # type: ignore [union-attr]
-                and self.predicate.eval(  # type: ignore [union-attr]
-                    None,  # type: ignore [arg-type]
+                self.predicate.condition  # type:ignore[union-attr]
+                and self.predicate.eval(  # type:ignore[union-attr]
+                    None,  # type: ignore[arg-type]
                     response=self._safe_response_json(response),
                     headers=response.headers,
                 )

airbyte_cdk/sources/declarative/retrievers/async_retriever.py

@@ -6,7 +6,7 @@ from typing import Any, Iterable, Mapping, Optional
 
 from typing_extensions import deprecated
 
-from airbyte_cdk.models import FailureType
+from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
 from airbyte_cdk.sources.declarative.async_job.job_orchestrator import AsyncPartition
 from airbyte_cdk.sources.declarative.extractors.record_selector import RecordSelector
 from airbyte_cdk.sources.declarative.partition_routers.async_job_partition_router import (
@@ -16,7 +16,6 @@ from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.source import ExperimentalClassWarning
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
-from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
 @deprecated(
@@ -57,9 +56,9 @@ class AsyncRetriever(Retriever):
 
         return self.state
 
-    def _validate_and_get_stream_slice_partition(
+    def _validate_and_get_stream_slice_jobs(
         self, stream_slice: Optional[StreamSlice] = None
-    ) -> AsyncPartition:
+    ) -> Iterable[AsyncJob]:
         """
         Validates the stream_slice argument and returns the partition from it.
 
@@ -73,12 +72,7 @@ class AsyncRetriever(Retriever):
             AirbyteTracedException: If the stream_slice is not an instance of StreamSlice or if the partition is not present in the stream_slice.
 
         """
-        if not isinstance(stream_slice, StreamSlice) or "partition" not in stream_slice.partition:
-            raise AirbyteTracedException(
-                message="Invalid arguments to AsyncJobRetriever.read_records: stream_slice is no optional. Please contact Airbyte Support",
-                failure_type=FailureType.system_error,
-            )
-        return stream_slice["partition"]  # type: ignore  # stream_slice["partition"] has been added as an AsyncPartition as part of stream_slices
+        return stream_slice.extra_fields.get("jobs", []) if stream_slice else []
 
     def stream_slices(self) -> Iterable[Optional[StreamSlice]]:
         return self.stream_slicer.stream_slices()
@@ -89,8 +83,8 @@ class AsyncRetriever(Retriever):
         stream_slice: Optional[StreamSlice] = None,
     ) -> Iterable[StreamData]:
         stream_state: StreamState = self._get_stream_state()
-        partition: AsyncPartition = self._validate_and_get_stream_slice_partition(stream_slice)
-        records: Iterable[Mapping[str, Any]] = self.stream_slicer.fetch_records(partition)
+        jobs: Iterable[AsyncJob] = self._validate_and_get_stream_slice_jobs(stream_slice)
+        records: Iterable[Mapping[str, Any]] = self.stream_slicer.fetch_records(jobs)
 
         yield from self.record_selector.filter_and_transform(
             all_data=records,

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -163,7 +163,7 @@ class SimpleRetriever(Retriever):
             stream_slice,
             next_page_token,
             self._paginator.get_request_headers,
-            self.stream_slicer.get_request_headers,
+            self.request_option_provider.get_request_headers,
         )
         if isinstance(headers, str):
             raise ValueError("Request headers cannot be a string")

airbyte_cdk/sources/declarative/schema/__init__.py

@@ -4,6 +4,7 @@
 
 from airbyte_cdk.sources.declarative.schema.default_schema_loader import DefaultSchemaLoader
 from airbyte_cdk.sources.declarative.schema.dynamic_schema_loader import (
+    ComplexFieldType,
     DynamicSchemaLoader,
     SchemaTypeIdentifier,
     TypesMap,
@@ -18,6 +19,7 @@ __all__ = [
     "SchemaLoader",
     "InlineSchemaLoader",
     "DynamicSchemaLoader",
+    "ComplexFieldType",
     "TypesMap",
     "SchemaTypeIdentifier",
 ]

airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py

@@ -18,7 +18,7 @@ from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.source import ExperimentalClassWarning
 from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
 
-AIRBYTE_DATA_TYPES: Mapping[str, Mapping[str, Any]] = {
+AIRBYTE_DATA_TYPES: Mapping[str, MutableMapping[str, Any]] = {
     "string": {"type": ["null", "string"]},
     "boolean": {"type": ["null", "boolean"]},
     "date": {"type": ["null", "string"], "format": "date"},
@@ -45,6 +45,25 @@ AIRBYTE_DATA_TYPES: Mapping[str, Mapping[str, Any]] = {
 }
 
 
+@deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
+@dataclass(frozen=True)
+class ComplexFieldType:
+    """
+    Identifies complex field type
+    """
+
+    field_type: str
+    items: Optional[Union[str, "ComplexFieldType"]] = None
+
+    def __post_init__(self) -> None:
+        """
+        Enforces that `items` is only used when `field_type` is a array
+        """
+        # `items_type` is valid only for array target types
+        if self.items and self.field_type != "array":
+            raise ValueError("'items' can only be used when 'field_type' is an array.")
+
+
 @deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
 @dataclass(frozen=True)
 class TypesMap:
@@ -52,7 +71,7 @@ class TypesMap:
     Represents a mapping between a current type and its corresponding target type.
     """
 
-    target_type: Union[List[str], str]
+    target_type: Union[List[str], str, ComplexFieldType]
     current_type: Union[List[str], str]
     condition: Optional[str]
 
@@ -135,8 +154,9 @@ class DynamicSchemaLoader(SchemaLoader):
         transformed_properties = self._transform(properties, {})
 
         return {
-            "$schema": "http://json-schema.org/draft-07/schema#",
+            "$schema": "https://json-schema.org/draft-07/schema#",
             "type": "object",
+            "additionalProperties": True,
             "properties": transformed_properties,
         }
 
@@ -188,18 +208,37 @@ class DynamicSchemaLoader(SchemaLoader):
             first_type = self._get_airbyte_type(mapped_field_type[0])
             second_type = self._get_airbyte_type(mapped_field_type[1])
             return {"oneOf": [first_type, second_type]}
+
         elif isinstance(mapped_field_type, str):
             return self._get_airbyte_type(mapped_field_type)
+
+        elif isinstance(mapped_field_type, ComplexFieldType):
+            return self._resolve_complex_type(mapped_field_type)
+
         else:
             raise ValueError(
                 f"Invalid data type. Available string or two items list of string. Got {mapped_field_type}."
             )
 
+    def _resolve_complex_type(self, complex_type: ComplexFieldType) -> Mapping[str, Any]:
+        if not complex_type.items:
+            return self._get_airbyte_type(complex_type.field_type)
+
+        field_type = self._get_airbyte_type(complex_type.field_type)
+
+        field_type["items"] = (
+            self._get_airbyte_type(complex_type.items)
+            if isinstance(complex_type.items, str)
+            else self._resolve_complex_type(complex_type.items)
+        )
+
+        return field_type
+
     def _replace_type_if_not_valid(
         self,
         field_type: Union[List[str], str],
         raw_schema: MutableMapping[str, Any],
-    ) -> Union[List[str], str]:
+    ) -> Union[List[str], str, ComplexFieldType]:
         """
         Replaces a field type if it matches a type mapping in `types_map`.
         """
@@ -216,7 +255,7 @@ class DynamicSchemaLoader(SchemaLoader):
         return field_type
 
     @staticmethod
-    def _get_airbyte_type(field_type: str) -> Mapping[str, Any]:
+    def _get_airbyte_type(field_type: str) -> MutableMapping[str, Any]:
         """
         Maps a field type to its corresponding Airbyte type definition.
         """

airbyte_cdk/sources/http_logger.py

@@ -45,7 +45,7 @@ def format_http_message(
         log_message["http"]["is_auxiliary"] = is_auxiliary  # type: ignore [index]
     if stream_name:
         log_message["airbyte_cdk"] = {"stream": {"name": stream_name}}
-    return log_message  # type: ignore [return-value]  # got "dict[str, object]", expected "dict[str, JsonType]"
+    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/streams/concurrent/clamping.py

@@ -0,0 +1,99 @@
+from abc import ABC
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Callable
+
+from airbyte_cdk.sources.streams.concurrent.cursor_types import CursorValueType
+
+
+class ClampingStrategy(ABC):
+    def clamp(self, value: CursorValueType) -> CursorValueType:
+        raise NotImplementedError()
+
+
+class NoClamping(ClampingStrategy):
+    def clamp(self, value: CursorValueType) -> CursorValueType:
+        return value
+
+
+class ClampingEndProvider:
+    def __init__(
+        self,
+        clamping_strategy: ClampingStrategy,
+        end_provider: Callable[[], CursorValueType],
+        granularity: timedelta,
+    ) -> None:
+        self._clamping_strategy = clamping_strategy
+        self._end_provider = end_provider
+        self._granularity = granularity
+
+    def __call__(self) -> CursorValueType:
+        return self._clamping_strategy.clamp(self._end_provider()) - self._granularity
+
+
+class DayClampingStrategy(ClampingStrategy):
+    def __init__(self, is_ceiling: bool = True) -> None:
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        return_value = value.replace(hour=0, minute=0, second=0, microsecond=0)
+        if self._is_ceiling:
+            return return_value + timedelta(days=1)
+        return return_value
+
+
+class MonthClampingStrategy(ClampingStrategy):
+    def __init__(self, is_ceiling: bool = True) -> None:
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        return_value = value.replace(hour=0, minute=0, second=0, microsecond=0)
+        needs_to_round = value.day != 1
+        if not needs_to_round:
+            return return_value
+
+        return self._ceil(return_value) if self._is_ceiling else return_value.replace(day=1)
+
+    def _ceil(self, value: datetime) -> datetime:
+        return value.replace(
+            year=value.year + 1 if value.month == 12 else value.year,
+            month=(value.month % 12) + 1,
+            day=1,
+            hour=0,
+            minute=0,
+            second=0,
+            microsecond=0,
+        )
+
+
+class Weekday(Enum):
+    """
+    These integer values map to the same ones used by the Datetime.date.weekday() implementation
+    """
+
+    MONDAY = 0
+    TUESDAY = 1
+    WEDNESDAY = 2
+    THURSDAY = 3
+    FRIDAY = 4
+    SATURDAY = 5
+    SUNDAY = 6
+
+
+class WeekClampingStrategy(ClampingStrategy):
+    def __init__(self, day_of_week: Weekday, is_ceiling: bool = True) -> None:
+        self._day_of_week = day_of_week.value
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        days_diff_to_ceiling = (
+            7 - (value.weekday() - self._day_of_week)
+            if value.weekday() > self._day_of_week
+            else abs(value.weekday() - self._day_of_week)
+        )
+        delta = (
+            timedelta(days_diff_to_ceiling)
+            if self._is_ceiling
+            else timedelta(days_diff_to_ceiling - 7)
+        )
+        return value.replace(hour=0, minute=0, second=0, microsecond=0) + delta