Merge remote-tracking branch 'origin/main' into aj/feat/add-standard-tests-cli

Author: aaronsteers

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -239,8 +239,8 @@ def get_token_expiry_date(self) -> AirbyteDateTime:
     def _has_access_token_been_initialized(self) -> bool:
         return self._access_token is not None
 
-    def set_token_expiry_date(self, value: Union[str, int]) -> None:
-        self._token_expiry_date = self._parse_token_expiration_date(value)
+    def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
+        self._token_expiry_date = value
 
     def get_assertion_name(self) -> str:
         return self.assertion_name

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -19,7 +19,10 @@
 from airbyte_cdk.sources.declarative.extractors.record_filter import (
     ClientSideIncrementalRecordFilterDecorator,
 )
-from airbyte_cdk.sources.declarative.incremental import ConcurrentPerPartitionCursor
+from airbyte_cdk.sources.declarative.incremental import (
+    ConcurrentPerPartitionCursor,
+    GlobalSubstreamCursor,
+)
 from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
 from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import (
     PerPartitionWithGlobalCursor,
@@ -361,7 +364,8 @@ def _group_streams(
                     == DatetimeBasedCursorModel.__name__
                     and hasattr(declarative_stream.retriever, "stream_slicer")
                     and isinstance(
-                        declarative_stream.retriever.stream_slicer, PerPartitionWithGlobalCursor
+                        declarative_stream.retriever.stream_slicer,
+                        (GlobalSubstreamCursor, PerPartitionWithGlobalCursor),
                     )
                 ):
                     stream_state = self._connector_state_manager.get_stream_state(

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -2307,6 +2307,27 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
+  KeyTransformation:
+    title: Transformation to apply for extracted object keys by Dpath Flatten Fields
+    type: object
+    required:
+      - type
+    properties:
+      type:
+        type: string
+        enum: [ KeyTransformation ]
+      prefix:
+        title: Key Prefix
+        description: Prefix to add for object keys. If not provided original keys remain unchanged.
+        type: string
+        examples:
+          - flattened_
+      suffix:
+        title: Key Suffix
+        description: Suffix to add for object keys. If not provided original keys remain unchanged.
+        type: string
+        examples:
+          - _flattened
   DpathFlattenFields:
     title: Dpath Flatten Fields
     description: A transformation that flatten field values to the to top of the record.
@@ -2335,6 +2356,11 @@ definitions:
         title: Replace Origin Record
         description: Whether to replace the origin record or not. Default is False.
         type: boolean
+      key_transformation:
+        title: Key transformation
+        description: Transformation for object keys. If not provided, original key will be used.
+        type: object
+        "$ref": "#/definitions/KeyTransformation"
       $parameters:
         type: object
         additionalProperties: true

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -879,6 +879,25 @@ class FlattenFields(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class KeyTransformation(BaseModel):
+    prefix: Optional[Union[str, None]] = Field(
+        None,
+        description="Prefix to add for object keys. If not provided original keys remain unchanged.",
+        examples=[
+            "flattened_",
+        ],
+        title="Key Prefix",
+    )
+    suffix: Optional[Union[str, None]] = Field(
+        None,
+        description="Suffix to add for object keys. If not provided original keys remain unchanged.",
+        examples=[
+            "_flattened",
+        ],
+        title="Key Suffix",
+    )
+
+
 class DpathFlattenFields(BaseModel):
     type: Literal["DpathFlattenFields"]
     field_path: List[str] = Field(
@@ -897,6 +916,11 @@ class DpathFlattenFields(BaseModel):
         description="Whether to replace the origin record or not. Default is False.",
         title="Replace Origin Record",
     )
+    key_transformation: Optional[Union[KeyTransformation, None]] = Field(
+        None,
+        description="Transformation for object keys. If not provided, original key will be used.",
+        title="Key transformation",
+    )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -498,6 +498,7 @@
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
 from airbyte_cdk.sources.declarative.transformations.dpath_flatten_fields import (
     DpathFlattenFields,
+    KeyTransformation,
 )
 from airbyte_cdk.sources.declarative.transformations.flatten_fields import (
     FlattenFields,
@@ -790,13 +791,24 @@ def create_dpath_flatten_fields(
         self, model: DpathFlattenFieldsModel, config: Config, **kwargs: Any
     ) -> DpathFlattenFields:
         model_field_path: List[Union[InterpolatedString, str]] = [x for x in model.field_path]
+        key_transformation = (
+            KeyTransformation(
+                config=config,
+                prefix=model.key_transformation.prefix,
+                suffix=model.key_transformation.suffix,
+                parameters=model.parameters or {},
+            )
+            if model.key_transformation is not None
+            else None
+        )
         return DpathFlattenFields(
             config=config,
             field_path=model_field_path,
             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,
+            key_transformation=key_transformation,
             parameters=model.parameters or {},
         )
 
@@ -1427,7 +1439,9 @@ def create_concurrent_cursor_from_perpartition_cursor(
         stream_state = self.apply_stream_state_migrations(stream_state_migrations, stream_state)
 
         # Per-partition state doesn't make sense for GroupingPartitionRouter, so force the global state
-        use_global_cursor = isinstance(partition_router, GroupingPartitionRouter)
+        use_global_cursor = isinstance(
+            partition_router, GroupingPartitionRouter
+        ) or component_definition.get("global_substream_cursor", False)
 
         # Return the concurrent cursor and state converter
         return ConcurrentPerPartitionCursor(
@@ -2054,6 +2068,7 @@ def create_default_paginator(
         config: Config,
         *,
         url_base: str,
+        extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None,
         decoder: Optional[Decoder] = None,
         cursor_used_for_stop_condition: Optional[DeclarativeCursor] = None,
     ) -> Union[DefaultPaginator, PaginatorTestReadDecorator]:
@@ -2075,7 +2090,10 @@ def create_default_paginator(
             else None
         )
         pagination_strategy = self._create_component_from_model(
-            model=model.pagination_strategy, config=config, decoder=decoder_to_use
+            model=model.pagination_strategy,
+            config=config,
+            decoder=decoder_to_use,
+            extractor_model=extractor_model,
         )
         if cursor_used_for_stop_condition:
             pagination_strategy = StopConditionPaginationStrategyDecorator(
@@ -2572,7 +2590,12 @@ def create_oauth_authenticator(
         )
 
     def create_offset_increment(
-        self, model: OffsetIncrementModel, config: Config, decoder: Decoder, **kwargs: Any
+        self,
+        model: OffsetIncrementModel,
+        config: Config,
+        decoder: Decoder,
+        extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None,
+        **kwargs: Any,
     ) -> OffsetIncrement:
         if isinstance(decoder, PaginationDecoderDecorator):
             inner_decoder = decoder.decoder
@@ -2587,10 +2610,24 @@ def create_offset_increment(
                 self._UNSUPPORTED_DECODER_ERROR.format(decoder_type=type(inner_decoder))
             )
 
+        # Ideally we would instantiate the runtime extractor from highest most level (in this case the SimpleRetriever)
+        # so that it can be shared by OffSetIncrement and RecordSelector. However, due to how we instantiate the
+        # decoder with various decorators here, but not in create_record_selector, it is simpler to retain existing
+        # behavior by having two separate extractors with identical behavior since they use the same extractor model.
+        # When we have more time to investigate we can look into reusing the same component.
+        extractor = (
+            self._create_component_from_model(
+                model=extractor_model, config=config, decoder=decoder_to_use
+            )
+            if extractor_model
+            else None
+        )
+
         return OffsetIncrement(
             page_size=model.page_size,
             config=config,
             decoder=decoder_to_use,
+            extractor=extractor,
             inject_on_first_request=model.inject_on_first_request or False,
             parameters=model.parameters or {},
         )
@@ -2954,6 +2991,7 @@ def create_simple_retriever(
                 model=model.paginator,
                 config=config,
                 url_base=url_base,
+                extractor_model=model.record_selector.extractor,
                 decoder=decoder,
                 cursor_used_for_stop_condition=cursor_used_for_stop_condition,
             )

airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py

@@ -12,6 +12,7 @@
     JsonDecoder,
     PaginationDecoderDecorator,
 )
+from airbyte_cdk.sources.declarative.extractors.dpath_extractor import RecordExtractor
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
     PaginationStrategy,
@@ -46,6 +47,7 @@ class OffsetIncrement(PaginationStrategy):
     config: Config
     page_size: Optional[Union[str, int]]
     parameters: InitVar[Mapping[str, Any]]
+    extractor: Optional[RecordExtractor]
     decoder: Decoder = field(
         default_factory=lambda: PaginationDecoderDecorator(decoder=JsonDecoder(parameters={}))
     )
@@ -75,6 +77,14 @@ def next_page_token(
     ) -> Optional[Any]:
         decoded_response = next(self.decoder.decode(response))
 
+        if self.extractor:
+            page_size_from_response = len(list(self.extractor.extract_records(response=response)))
+            # The extractor could return 0 records which is valid, but evaluates to False. Our fallback in other
+            # cases as the best effort option is to use the incoming last_page_size
+            last_page_size = (
+                page_size_from_response if page_size_from_response is not None else last_page_size
+            )
+
         # Stop paginating when there are fewer records than the page size or the current page has no records
         if (
             self._page_size

airbyte_cdk/sources/declarative/requesters/query_properties/property_chunking.py

@@ -52,8 +52,10 @@ def get_request_property_chunks(
         chunk_size = 0
         for property_field in property_fields:
             # If property_limit_type is not defined, we default to property_count which is just an incrementing count
+            # todo: Add ability to specify parameter delimiter representation and take into account in property_field_size
             property_field_size = (
                 len(property_field)
+                + 1  # The +1 represents the extra character for the delimiter in between properties
                 if self.property_limit_type == PropertyLimitType.characters
                 else 1
             )