Merge branch 'lazebnyi/add-lazy-read-to-simple-retriver' of github.com:airbytehq/airbyte-python-cdk into lazebnyi/add-lazy-read-to-simple-retriver

Author: lazebnyi

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -116,6 +116,19 @@ definitions:
         type: array
         items:
           "$ref": "#/definitions/AddedFieldDefinition"
+      condition:
+        description: Fields will be added if expression is evaluated to True.
+        type: string
+        default: ""
+        interpolation_context:
+          - config
+          - property
+          - parameters
+        examples:
+          - "{{ property|string == '' }}"
+          - "{{ property is integer }}"
+          - "{{ property|length > 5 }}"
+          - "{{ property == 'some_string_to_match' }}"
       $parameters:
         type: object
         additionalProperties: true
@@ -2265,6 +2278,10 @@ definitions:
         title: Delete Origin Value
         description: Whether to delete the origin value or keep it. Default is False.
         type: boolean
+      replace_record:
+        title: Replace Origin Record
+        description: Whether to replace the origin record or not. Default is False.
+        type: boolean
       $parameters:
         type: object
         additionalProperties: true
@@ -3000,6 +3017,10 @@ definitions:
           - "$ref": "#/definitions/SchemaNormalization"
           - "$ref": "#/definitions/CustomSchemaNormalization"
         default: None
+      transform_before_filtering:
+        description: If true, transformation will be applied before record filtering.
+        type: boolean
+        default: false
       $parameters:
         type: object
         additionalProperties: true

airbyte_cdk/sources/declarative/declarative_stream.py

@@ -14,6 +14,7 @@
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
 from airbyte_cdk.sources.declarative.retrievers import SimpleRetriever
+from airbyte_cdk.sources.declarative.retrievers.async_retriever import AsyncRetriever
 from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.declarative.schema import DefaultSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
@@ -76,11 +77,17 @@ def primary_key(self, value: str) -> None:
 
     @property
     def exit_on_rate_limit(self) -> bool:
+        if isinstance(self.retriever, AsyncRetriever):
+            return self.retriever.exit_on_rate_limit
+
         return self.retriever.requester.exit_on_rate_limit  # type: ignore # abstract Retriever class has not requester attribute
 
     @exit_on_rate_limit.setter
     def exit_on_rate_limit(self, value: bool) -> None:
-        self.retriever.requester.exit_on_rate_limit = value  # type: ignore[attr-defined]
+        if isinstance(self.retriever, AsyncRetriever):
+            self.retriever.exit_on_rate_limit = value
+        else:
+            self.retriever.requester.exit_on_rate_limit = value  # type: ignore[attr-defined]
 
     @property  # type: ignore
     def name(self) -> str:

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -877,6 +877,11 @@ class DpathFlattenFields(BaseModel):
         description="Whether to delete the origin value or keep it. Default is False.",
         title="Delete Origin Value",
     )
+    replace_record: Optional[bool] = Field(
+        None,
+        description="Whether to replace the origin record or not. Default is False.",
+        title="Replace Origin Record",
+    )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
@@ -1460,6 +1465,16 @@ class AddFields(BaseModel):
         description="List of transformations (path and corresponding value) that will be added to the record.",
         title="Fields",
     )
+    condition: Optional[str] = Field(
+        "",
+        description="Fields will be added if expression is evaluated to True.,",
+        examples=[
+            "{{ property|string == '' }}",
+            "{{ property is integer }}",
+            "{{ property|length > 5 }}",
+            "{{ property == 'some_string_to_match' }}",
+        ],
+    )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
@@ -1771,6 +1786,10 @@ class RecordSelector(BaseModel):
         description="Responsible for normalization according to the schema.",
         title="Schema Normalization",
     )
+    transform_before_filtering: Optional[bool] = Field(
+        False,
+        description="If true, transformation will be applied before record filtering.",
+    )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -713,7 +713,11 @@ def create_add_fields(self, model: AddFieldsModel, config: Config, **kwargs: Any
             )
             for added_field_definition_model in model.fields
         ]
-        return AddFields(fields=added_field_definitions, parameters=model.parameters or {})
+        return AddFields(
+            fields=added_field_definitions,
+            condition=model.condition or "",
+            parameters=model.parameters or {},
+        )
 
     def create_keys_to_lower_transformation(
         self, model: KeysToLowerModel, config: Config, **kwargs: Any
@@ -749,6 +753,7 @@ def create_dpath_flatten_fields(
             delete_origin_value=model.delete_origin_value
             if model.delete_origin_value is not None
             else False,
+            replace_record=model.replace_record if model.replace_record is not None else False,
             parameters=model.parameters or {},
         )
 
@@ -1903,6 +1908,10 @@ def _merge_stream_slicers(
     ) -> Optional[StreamSlicer]:
         retriever_model = model.retriever
 
+        stream_slicer = self._build_stream_slicer_from_partition_router(
+            retriever_model, config, stream_name=model.name
+        )
+
         if retriever_model.type == "AsyncRetriever":
             is_not_datetime_cursor = (
                 model.incremental_sync.type != "DatetimeBasedCursor"
@@ -1922,13 +1931,11 @@ def _merge_stream_slicers(
                     "AsyncRetriever with cursor other than DatetimeBasedCursor is not supported yet."
                 )
 
-            if is_partition_router:
+            if is_partition_router and not stream_slicer:
                 # Note that this development is also done in parallel to the per partition development which once merged
                 # we could support here by calling create_concurrent_cursor_from_perpartition_cursor
                 raise ValueError("Per partition state is not supported yet for AsyncRetriever.")
 
-        stream_slicer = self._build_stream_slicer_from_partition_router(retriever_model, config, stream_name=model.name)
-
         if model.incremental_sync:
             return self._build_incremental_cursor(model, stream_slicer, config)
 
@@ -2607,7 +2614,9 @@ def create_record_selector(
             else None
         )
 
-        transform_before_filtering = False
+        assert model.transform_before_filtering is not None  # for mypy
+
+        transform_before_filtering = model.transform_before_filtering
         if client_side_incremental_sync:
             record_filter = ClientSideIncrementalRecordFilterDecorator(
                 config=config,

airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py

@@ -320,20 +320,15 @@ def set_initial_state(self, stream_state: StreamState) -> None:
 
         parent_state = stream_state.get("parent_state", {})
 
-        # If `parent_state` doesn't exist and at least one parent stream has an incremental dependency,
-        # copy the child state to parent streams with incremental dependencies.
-        incremental_dependency = any(
-            [parent_config.incremental_dependency for parent_config in self.parent_stream_configs]
-        )
-        if not parent_state and not incremental_dependency:
-            return
-
-        if not parent_state and incremental_dependency:
-            # 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 (
+                not parent_state.get(parent_config.stream.name, {})
+                and parent_config.incremental_dependency
+            ):
+                # Migrate child state to parent state format
+                parent_state = self._migrate_child_state_to_parent_state(stream_state)
+
             if parent_config.incremental_dependency:
                 parent_config.stream.state = parent_state.get(parent_config.stream.name, {})
 

airbyte_cdk/sources/declarative/retrievers/async_retriever.py

@@ -36,6 +36,36 @@ class AsyncRetriever(Retriever):
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
 
+    @property
+    def exit_on_rate_limit(self) -> bool:
+        """
+        Whether to exit on rate limit. This is a property of the job repository
+        and not the stream slicer. The stream slicer is responsible for creating
+        the jobs, but the job repository is responsible for managing the rate
+        limits and other job-related properties.
+
+        Note:
+         - If the `creation_requester` cannot place / create the job - it might be the case of the RateLimits
+         - If the `creation_requester` can place / create the job - it means all other requesters should successfully manage
+           to complete the results.
+        """
+        job_orchestrator = self.stream_slicer._job_orchestrator
+        if job_orchestrator is None:
+            # Default value when orchestrator is not available
+            return False
+        return job_orchestrator._job_repository.creation_requester.exit_on_rate_limit  # type: ignore
+
+    @exit_on_rate_limit.setter
+    def exit_on_rate_limit(self, value: bool) -> None:
+        """
+        Sets the `exit_on_rate_limit` property of the job repository > creation_requester,
+        meaning that the Job cannot be placed / created if the rate limit is reached.
+        Thus no further work on managing jobs is expected to be done.
+        """
+        job_orchestrator = self.stream_slicer._job_orchestrator
+        if job_orchestrator is not None:
+            job_orchestrator._job_repository.creation_requester.exit_on_rate_limit = value  # type: ignore[attr-defined, assignment]
+
     @property
     def state(self) -> StreamState:
         """

airbyte_cdk/sources/declarative/transformations/add_fields.py

@@ -1,12 +1,13 @@
 #
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
 #
 
 from dataclasses import InitVar, dataclass, field
 from typing import Any, Dict, List, Mapping, Optional, Type, Union
 
 import dpath
 
+from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.types import Config, FieldPointer, StreamSlice, StreamState
@@ -86,11 +87,16 @@ class AddFields(RecordTransformation):
 
     fields: List[AddedFieldDefinition]
     parameters: InitVar[Mapping[str, Any]]
+    condition: str = ""
     _parsed_fields: List[ParsedAddFieldDefinition] = field(
         init=False, repr=False, default_factory=list
     )
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        self._filter_interpolator = InterpolatedBoolean(
+            condition=self.condition, parameters=parameters
+        )
+
         for add_field in self.fields:
             if len(add_field.path) < 1:
                 raise ValueError(
@@ -132,7 +138,9 @@ def transform(
         for parsed_field in self._parsed_fields:
             valid_types = (parsed_field.value_type,) if parsed_field.value_type else None
             value = parsed_field.value.eval(config, valid_types=valid_types, **kwargs)
-            dpath.new(record, parsed_field.path, value)
+            is_empty_condition = not self.condition
+            if is_empty_condition or self._filter_interpolator.eval(config, value=value, **kwargs):
+                dpath.new(record, parsed_field.path, value)
 
     def __eq__(self, other: Any) -> bool:
         return bool(self.__dict__ == other.__dict__)

airbyte_cdk/sources/declarative/transformations/dpath_flatten_fields.py

@@ -15,13 +15,15 @@ class DpathFlattenFields(RecordTransformation):
 
     field_path: List[Union[InterpolatedString, str]] path to the field to flatten.
     delete_origin_value: bool = False whether to delete origin field or keep it. Default is False.
+    replace_record: bool = False whether to replace origin record or not. Default is False.
 
     """
 
     config: Config
     field_path: List[Union[InterpolatedString, str]]
     parameters: InitVar[Mapping[str, Any]]
     delete_origin_value: bool = False
+    replace_record: bool = False
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._field_path = [
@@ -48,8 +50,12 @@ def transform(
             extracted = dpath.get(record, path, default=[])
 
         if isinstance(extracted, dict):
-            conflicts = set(extracted.keys()) & set(record.keys())
-            if not conflicts:
-                if self.delete_origin_value:
-                    dpath.delete(record, path)
+            if self.replace_record and extracted:
+                dpath.delete(record, "**")
                 record.update(extracted)
+            else:
+                conflicts = set(extracted.keys()) & set(record.keys())
+                if not conflicts:
+                    if self.delete_origin_value:
+                        dpath.delete(record, path)
+                    record.update(extracted)