Merge branch 'main' into devin/1774478445-memory-failfast

Author: pnilan

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -76,12 +76,19 @@
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
+from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
+    GroupingPartitionRouter,
+)
+from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
+    SubstreamPartitionRouter,
+)
 from airbyte_cdk.sources.declarative.resolvers import COMPONENTS_RESOLVER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.spec.spec import Spec
 from airbyte_cdk.sources.declarative.types import Config, ConnectionDefinition
 from airbyte_cdk.sources.message.concurrent_repository import ConcurrentMessageRepository
 from airbyte_cdk.sources.message.repository import InMemoryMessageRepository
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
+from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from airbyte_cdk.sources.streams.concurrent.partitions.types import QueueItem
 from airbyte_cdk.sources.utils.slice_logger import (
     AlwaysLogSliceLogger,
@@ -405,6 +412,8 @@ def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]:  # type: i
         if api_budget_model:
             self._constructor.set_api_budget(api_budget_model, self._config)
 
+        prepared_configs = self._initialize_cache_for_parent_streams(deepcopy(stream_configs))
+
         source_streams = [
             self._constructor.create_component(
                 (
@@ -416,10 +425,70 @@ def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]:  # type: i
                 self._config,
                 emit_connector_builder_messages=self._emit_connector_builder_messages,
             )
-            for stream_config in self._initialize_cache_for_parent_streams(deepcopy(stream_configs))
+            for stream_config in prepared_configs
         ]
+
+        self._apply_stream_groups(source_streams)
+
         return source_streams
 
+    def _apply_stream_groups(self, streams: List[AbstractStream]) -> None:
+        """Set block_simultaneous_read on streams based on the manifest's stream_groups config.
+
+        Iterates over the resolved manifest's stream_groups and matches group membership
+        against actual created stream instances by name. Validates that no stream shares a
+        group with any of its parent streams, which would cause a deadlock.
+        """
+        stream_groups = self._source_config.get("stream_groups", {})
+        if not stream_groups:
+            return
+
+        # Build stream_name -> group_name mapping from the resolved manifest
+        stream_name_to_group: Dict[str, str] = {}
+        for group_name, group_config in stream_groups.items():
+            for stream_ref in group_config.get("streams", []):
+                if isinstance(stream_ref, dict):
+                    stream_name = stream_ref.get("name", "")
+                    if stream_name:
+                        stream_name_to_group[stream_name] = group_name
+
+        # Validate no stream shares a group with any of its ancestor streams
+        stream_name_to_instance: Dict[str, AbstractStream] = {s.name: s for s in streams}
+
+        def _collect_all_ancestor_names(stream_name: str) -> Set[str]:
+            """Recursively collect all ancestor stream names."""
+            ancestors: Set[str] = set()
+            inst = stream_name_to_instance.get(stream_name)
+            if not isinstance(inst, DefaultStream):
+                return ancestors
+            router = inst.get_partition_router()
+            if isinstance(router, GroupingPartitionRouter):
+                router = router.underlying_partition_router
+            if not isinstance(router, SubstreamPartitionRouter):
+                return ancestors
+            for parent_config in router.parent_stream_configs:
+                parent_name = parent_config.stream.name
+                ancestors.add(parent_name)
+                ancestors.update(_collect_all_ancestor_names(parent_name))
+            return ancestors
+
+        for stream in streams:
+            if not isinstance(stream, DefaultStream) or stream.name not in stream_name_to_group:
+                continue
+            group_name = stream_name_to_group[stream.name]
+            for ancestor_name in _collect_all_ancestor_names(stream.name):
+                if stream_name_to_group.get(ancestor_name) == group_name:
+                    raise ValueError(
+                        f"Stream '{stream.name}' and its parent stream '{ancestor_name}' "
+                        f"are both in group '{group_name}'. "
+                        f"A child stream must not share a group with its parent to avoid deadlock."
+                    )
+
+        # Apply group to matching stream instances
+        for stream in streams:
+            if isinstance(stream, DefaultStream) and stream.name in stream_name_to_group:
+                stream.block_simultaneous_read = stream_name_to_group[stream.name]
+
     @staticmethod
     def _initialize_cache_for_parent_streams(
         stream_configs: List[Dict[str, Any]],

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -45,6 +45,15 @@ properties:
     "$ref": "#/definitions/ConcurrencyLevel"
   api_budget:
     "$ref": "#/definitions/HTTPAPIBudget"
+  stream_groups:
+    title: Stream Groups
+    description: >
+      Groups of streams that share a common resource and should not be read simultaneously.
+      Each group defines a set of stream references and an action that controls how concurrent
+      reads are managed. Only applies to ConcurrentDeclarativeSource.
+    type: object
+    additionalProperties:
+      "$ref": "#/definitions/StreamGroup"
   max_concurrent_async_job_count:
     title: Maximum Concurrent Asynchronous Jobs
     description: Maximum number of concurrent asynchronous jobs to run. This property is only relevant for sources/streams that support asynchronous job execution through the AsyncRetriever (e.g. a report-based stream that initiates a job, polls the job status, and then fetches the job results). This is often set by the API's maximum number of concurrent jobs on the account level. Refer to the API's documentation for this information.
@@ -1832,6 +1841,15 @@ definitions:
         description: The headers to match.
         type: object
         additionalProperties: true
+      weight:
+        title: Weight
+        description: >
+          The weight of a request matching this matcher when acquiring a call from the rate limiter.
+          Different endpoints can consume different amounts from a shared budget by specifying
+          different weights. If not set, each request counts as 1.
+        anyOf:
+          - type: integer
+          - type: string
     additionalProperties: true
   DefaultErrorHandler:
     title: Default Error Handler
@@ -4191,6 +4209,43 @@ definitions:
             - "$ref": "#/definitions/ConfigRemoveFields"
             - "$ref": "#/definitions/CustomConfigTransformation"
         default: []
+  StreamGroup:
+    title: Stream Group
+    description: >
+      A group of streams that share a common resource and should not be read simultaneously.
+      Streams in the same group will be blocked from concurrent reads based on the specified action.
+    type: object
+    required:
+      - streams
+      - action
+    properties:
+      streams:
+        title: Streams
+        description: >
+          List of references to streams that belong to this group.
+        type: array
+        items:
+          anyOf:
+            - "$ref": "#/definitions/DeclarativeStream"
+      action:
+        title: Action
+        description: The action to apply to streams in this group.
+        "$ref": "#/definitions/BlockSimultaneousSyncsAction"
+  BlockSimultaneousSyncsAction:
+    title: Block Simultaneous Syncs Action
+    description: >
+      Action that prevents streams in the same group from being read concurrently.
+      When applied to a stream group, streams with this action will be deferred if
+      another stream in the same group is currently active.
+      This is useful for APIs that don't allow concurrent access to the same
+      endpoint or session. Only applies to ConcurrentDeclarativeSource.
+    type: object
+    required:
+      - type
+    properties:
+      type:
+        type: string
+        enum: [BlockSimultaneousSyncsAction]
   SubstreamPartitionRouter:
     title: Substream Partition Router
     description: Partition router that is used to retrieve records that have been partitioned according to records from the specified parent streams. An example of a parent stream is automobile brands and the substream would be the various car models associated with each branch.

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -486,6 +486,11 @@ class Config:
     headers: Optional[Dict[str, Any]] = Field(
         None, description="The headers to match.", title="Headers"
     )
+    weight: Optional[Union[int, str]] = Field(
+        None,
+        description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.",
+        title="Weight",
+    )
 
 
 class DpathExtractor(BaseModel):
@@ -2394,6 +2399,11 @@ class Config:
     spec: Optional[Spec] = None
     concurrency_level: Optional[ConcurrencyLevel] = None
     api_budget: Optional[HTTPAPIBudget] = None
+    stream_groups: Optional[Dict[str, StreamGroup]] = Field(
+        None,
+        description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.",
+        title="Stream Groups",
+    )
     max_concurrent_async_job_count: Optional[Union[int, str]] = Field(
         None,
         description="Maximum number of concurrent asynchronous jobs to run. This property is only relevant for sources/streams that support asynchronous job execution through the AsyncRetriever (e.g. a report-based stream that initiates a job, polls the job status, and then fetches the job results). This is often set by the API's maximum number of concurrent jobs on the account level. Refer to the API's documentation for this information.",
@@ -2429,6 +2439,11 @@ class Config:
     spec: Optional[Spec] = None
     concurrency_level: Optional[ConcurrencyLevel] = None
     api_budget: Optional[HTTPAPIBudget] = None
+    stream_groups: Optional[Dict[str, StreamGroup]] = Field(
+        None,
+        description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.",
+        title="Stream Groups",
+    )
     max_concurrent_async_job_count: Optional[Union[int, str]] = Field(
         None,
         description="Maximum number of concurrent asynchronous jobs to run. This property is only relevant for sources/streams that support asynchronous job execution through the AsyncRetriever (e.g. a report-based stream that initiates a job, polls the job status, and then fetches the job results). This is often set by the API's maximum number of concurrent jobs on the account level. Refer to the API's documentation for this information.",
@@ -3096,6 +3111,23 @@ class AsyncRetriever(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class BlockSimultaneousSyncsAction(BaseModel):
+    type: Literal["BlockSimultaneousSyncsAction"]
+
+
+class StreamGroup(BaseModel):
+    streams: List[str] = Field(
+        ...,
+        description='List of references to streams that belong to this group. Use JSON references to stream definitions (e.g., "#/definitions/my_stream").',
+        title="Streams",
+    )
+    action: BlockSimultaneousSyncsAction = Field(
+        ...,
+        description="The action to apply to streams in this group.",
+        title="Action",
+    )
+
+
 class SubstreamPartitionRouter(BaseModel):
     type: Literal["SubstreamPartitionRouter"]
     parent_stream_configs: List[ParentStreamConfig] = Field(

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -4387,12 +4387,21 @@ def create_rate(self, model: RateModel, config: Config, **kwargs: Any) -> Rate:
     def create_http_request_matcher(
         self, model: HttpRequestRegexMatcherModel, config: Config, **kwargs: Any
     ) -> HttpRequestRegexMatcher:
+        weight = model.weight
+        if weight is not None:
+            if isinstance(weight, str):
+                weight = int(InterpolatedString.create(weight, parameters={}).eval(config))
+            else:
+                weight = int(weight)
+            if weight < 1:
+                raise ValueError(f"weight must be >= 1, got {weight}")
         return HttpRequestRegexMatcher(
             method=model.method,
             url_base=model.url_base,
             url_path_pattern=model.url_path_pattern,
             params=model.params,
             headers=model.headers,
+            weight=weight,
         )
 
     def set_api_budget(self, component_definition: ComponentDefinition, config: Config) -> None:

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -166,14 +166,22 @@ def __init__(
         url_path_pattern: Optional[str] = None,
         params: Optional[Mapping[str, Any]] = None,
         headers: Optional[Mapping[str, Any]] = None,
+        weight: Optional[int] = None,
     ):
         """
         :param method: HTTP method (e.g. "GET", "POST"); compared case-insensitively.
         :param url_base: Base URL (scheme://host) that must match.
         :param url_path_pattern: A regex pattern that will be applied to the path portion of the URL.
         :param params: Dictionary of query parameters that must be present in the request.
         :param headers: Dictionary of headers that must be present (header keys are compared case-insensitively).
+        :param weight: The weight of a request matching this matcher. If set, this value is used
+            when acquiring a call from the rate limiter, enabling cost-based rate limiting
+            where different endpoints consume different amounts from a shared budget.
+            If not set, each request counts as 1.
         """
+        if weight is not None and weight < 1:
+            raise ValueError(f"weight must be >= 1, got {weight}")
+        self._weight = weight
         self._method = method.upper() if method else None
 
         # Normalize the url_base if provided: remove trailing slash.
@@ -242,11 +250,16 @@ def __call__(self, request: Any) -> bool:
 
         return True
 
+    @property
+    def weight(self) -> Optional[int]:
+        """The weight of a request matching this matcher, or None if not set."""
+        return self._weight
+
     def __str__(self) -> str:
         regex = self._url_path_pattern.pattern if self._url_path_pattern else None
         return (
             f"HttpRequestRegexMatcher(method={self._method}, url_base={self._url_base}, "
-            f"url_path_pattern={regex}, params={self._params}, headers={self._headers})"
+            f"url_path_pattern={regex}, params={self._params}, headers={self._headers}, weight={self._weight})"
         )
 
 
@@ -265,6 +278,22 @@ def matches(self, request: Any) -> bool:
             return True
         return any(matcher(request) for matcher in self._matchers)
 
+    def get_weight(self, request: Any) -> int:
+        """Get the weight for a request based on the first matching matcher.
+
+        If a matcher has a weight configured, that weight is used.
+        Otherwise, defaults to 1.
+
+        :param request: a request object
+        :return: the weight for this request
+        """
+        for matcher in self._matchers:
+            if matcher(request):
+                if isinstance(matcher, HttpRequestRegexMatcher) and matcher.weight is not None:
+                    return matcher.weight
+                return 1
+        return 1
+
 
 class UnlimitedCallRatePolicy(BaseCallRatePolicy):
     """
@@ -420,6 +449,11 @@ def __init__(self, rates: list[Rate], matchers: list[RequestMatcher]):
     def try_acquire(self, request: Any, weight: int) -> None:
         if not self.matches(request):
             raise ValueError("Request does not match the policy")
+        lowest_limit = min(rate.limit for rate in self._bucket.rates)
+        if weight > lowest_limit:
+            raise ValueError(
+                f"Weight can not exceed the lowest configured rate limit ({lowest_limit})"
+            )
 
         try:
             self._limiter.try_acquire(request, weight=weight)
@@ -596,7 +630,8 @@ def _do_acquire(
         # sometimes we spend all budget before a second attempt, so we have a few more attempts
         for attempt in range(1, self._maximum_attempts_to_acquire):
             try:
-                policy.try_acquire(request, weight=1)
+                weight = policy.get_weight(request) if isinstance(policy, BaseCallRatePolicy) else 1
+                policy.try_acquire(request, weight=weight)
                 return
             except CallRateLimitHit as exc:
                 last_exception = exc

airbyte_cdk/sources/streams/call_rate.py

@@ -85,6 +85,21 @@ def cursor(self) -> Cursor:
         :return: The cursor associated with this stream.
         """
 
+    @property
+    def block_simultaneous_read(self) -> str:
+        """
+        Override to return a non-empty group name if this stream should block simultaneous reads.
+        When a non-empty string is returned, prevents starting partition generation for this stream if:
+        - Another stream with the same group name is already active
+        - Any of its parent streams are in an active group
+
+        This allows grouping multiple streams that share the same resource (e.g., API endpoint or session)
+        to prevent them from running concurrently, even if they don't have a parent-child relationship.
+
+        :return: Group name for blocking (non-empty string), or "" to allow concurrent reading
+        """
+        return ""  # Default: allow concurrent reading
+
     @abstractmethod
     def check_availability(self) -> StreamAvailability:
         """

airbyte_cdk/sources/streams/concurrent/abstract_stream.py

@@ -196,6 +196,11 @@ def cursor_field(self) -> Union[str, List[str]]:
     def cursor(self) -> Optional[Cursor]:  # type: ignore[override] # StreamFaced expects to use only airbyte_cdk.sources.streams.concurrent.cursor.Cursor
         return self._cursor
 
+    @property
+    def block_simultaneous_read(self) -> str:
+        """Returns the blocking group name from the underlying stream"""
+        return self._abstract_stream.block_simultaneous_read
+
     # FIXME the lru_cache seems to be mostly there because of typing issue
     @lru_cache(maxsize=None)
     def get_json_schema(self) -> Mapping[str, Any]: