Refactor ConcurrentPerPartitionCursor to not use ConcurrentCursor with _use_global_cursor

Author: tolik0

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -40,12 +40,6 @@ properties:
     "$ref": "#/definitions/Spec"
   concurrency_level:
     "$ref": "#/definitions/ConcurrencyLevel"
-  api_budget:
-    title: API Budget
-    description: Defines how many requests can be made to the API in a given time frame. This field accepts either a generic APIBudget or an HTTP-specific configuration (HTTPAPIBudget) to be applied across all streams.
-    anyOf:
-      - "$ref": "#/definitions/APIBudget"
-      - "$ref": "#/definitions/HTTPAPIBudget"
   metadata:
     type: object
     description: For internal Airbyte use only - DO NOT modify manually. Used by consumers of declarative manifests for storing related metadata.
@@ -1371,207 +1365,6 @@ definitions:
       $parameters:
         type: object
         additional_properties: true
-  APIBudget:
-    title: API Budget
-    description: >
-      A generic API budget configuration that defines the policies (rate limiting rules)
-      and the maximum number of attempts to acquire a call credit. This budget does not automatically
-      update itself based on HTTP response headers.
-    type: object
-    required:
-      - type
-      - policies
-    properties:
-      type:
-        type: string
-        enum: [APIBudget]
-      policies:
-        title: Policies
-        description: List of call rate policies that define how many calls are allowed.
-        type: array
-        items:
-          anyOf:
-            - "$ref": "#/definitions/FixedWindowCallRatePolicy"
-            - "$ref": "#/definitions/MovingWindowCallRatePolicy"
-            - "$ref": "#/definitions/UnlimitedCallRatePolicy"
-      maximum_attempts_to_acquire:
-        title: Maximum Attempts to Acquire
-        description: The maximum number of attempts to acquire a call before giving up.
-        type: integer
-        default: 100000
-    additionalProperties: true
-  HTTPAPIBudget:
-    title: HTTP API Budget
-    description: >
-      An HTTP-specific API budget that extends APIBudget by updating rate limiting information based
-      on HTTP response headers. It extracts available calls and the next reset timestamp from the HTTP responses.
-    type: object
-    required:
-      - type
-      - policies
-    properties:
-      type:
-        type: string
-        enum: [HTTPAPIBudget]
-      policies:
-        title: Policies
-        description: List of call rate policies that define how many calls are allowed.
-        type: array
-        items:
-          anyOf:
-            - "$ref": "#/definitions/FixedWindowCallRatePolicy"
-            - "$ref": "#/definitions/MovingWindowCallRatePolicy"
-            - "$ref": "#/definitions/UnlimitedCallRatePolicy"
-      ratelimit_reset_header:
-        title: Rate Limit Reset Header
-        description: The HTTP response header name that indicates when the rate limit resets.
-        type: string
-        default: "ratelimit-reset"
-      ratelimit_remaining_header:
-        title: Rate Limit Remaining Header
-        description: The HTTP response header name that indicates the number of remaining allowed calls.
-        type: string
-        default: "ratelimit-remaining"
-      status_codes_for_ratelimit_hit:
-        title: Status Codes for Rate Limit Hit
-        description: List of HTTP status codes that indicate a rate limit has been hit.
-        type: array
-        items:
-          type: integer
-        default: [429]
-      maximum_attempts_to_acquire:
-        title: Maximum Attempts to Acquire
-        description: The maximum number of attempts to acquire a call before giving up.
-        type: integer
-        default: 100000
-    additionalProperties: true
-  FixedWindowCallRatePolicy:
-    title: Fixed Window Call Rate Policy
-    description: A policy that allows a fixed number of calls within a specific time window.
-    type: object
-    required:
-      - type
-      - next_reset_ts
-      - period
-      - call_limit
-      - matchers
-    properties:
-      type:
-        type: string
-        enum: [FixedWindowCallRatePolicy]
-      next_reset_ts:
-        title: Next Reset Timestamp
-        description: The timestamp when the rate limit will reset.
-        type: string
-        format: date-time
-      period:
-        title: Period
-        description: The time interval for the rate limit window.
-        type: string
-        format: duration
-      call_limit:
-        title: Call Limit
-        description: The maximum number of calls allowed within the period.
-        type: integer
-      matchers:
-        title: Matchers
-        description: List of matchers that define which requests this policy applies to.
-        type: array
-        items:
-          "$ref": "#/definitions/HttpRequestRegexMatcher"
-    additionalProperties: true
-  MovingWindowCallRatePolicy:
-    title: Moving Window Call Rate Policy
-    description: A policy that allows a fixed number of calls within a moving time window.
-    type: object
-    required:
-      - type
-      - rates
-      - matchers
-    properties:
-      type:
-        type: string
-        enum: [MovingWindowCallRatePolicy]
-      rates:
-        title: Rates
-        description: List of rates that define the call limits for different time intervals.
-        type: array
-        items:
-          "$ref": "#/definitions/Rate"
-      matchers:
-        title: Matchers
-        description: List of matchers that define which requests this policy applies to.
-        type: array
-        items:
-          "$ref": "#/definitions/HttpRequestRegexMatcher"
-    additionalProperties: true
-  UnlimitedCallRatePolicy:
-    title: Unlimited Call Rate Policy
-    description: A policy that allows unlimited calls for specific requests.
-    type: object
-    required:
-      - type
-      - matchers
-    properties:
-      type:
-        type: string
-        enum: [UnlimitedCallRatePolicy]
-      matchers:
-        title: Matchers
-        description: List of matchers that define which requests this policy applies to.
-        type: array
-        items:
-          "$ref": "#/definitions/HttpRequestRegexMatcher"
-    additionalProperties: true
-  Rate:
-    title: Rate
-    description: Defines a rate limit with a specific number of calls allowed within a time interval.
-    type: object
-    required:
-      - limit
-      - interval
-    properties:
-      limit:
-        title: Limit
-        description: The maximum number of calls allowed within the interval.
-        type: integer
-      interval:
-        title: Interval
-        description: The time interval for the rate limit.
-        type: string
-        format: duration
-    additionalProperties: true
-  HttpRequestRegexMatcher:
-    title: HTTP Request Matcher
-    description: >
-      Matches HTTP requests based on method, base URL, URL path pattern, query parameters, and headers.
-      Use `url_base` to specify the scheme and host (without trailing slash) and
-      `url_path_pattern` to apply a regex to the request path.
-    type: object
-    properties:
-      method:
-        title: Method
-        description: The HTTP method to match (e.g., GET, POST).
-        type: string
-      url_base:
-        title: URL Base
-        description: The base URL (scheme and host, e.g. "https://api.example.com") to match.
-        type: string
-      url_path_pattern:
-        title: URL Path Pattern
-        description: A regular expression pattern to match the URL path.
-        type: string
-      params:
-        title: Parameters
-        description: The query parameters to match.
-        type: object
-        additionalProperties: true
-      headers:
-        title: Headers
-        description: The headers to match.
-        type: object
-        additionalProperties: true
-    additionalProperties: true
   DefaultErrorHandler:
     title: Default Error Handler
     description: Component defining how to handle errors. Default behavior includes only retrying server errors (HTTP 5XX) and too many requests (HTTP 429) with an exponential backoff.

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -142,20 +142,16 @@ def close_partition(self, partition: Partition) -> None:
             raise ValueError("stream_slice cannot be None")
 
         partition_key = self._to_partition_key(stream_slice.partition)
-        self._cursor_per_partition[partition_key].close_partition(partition=partition)
+        if not self._use_global_cursor:
+            self._cursor_per_partition[partition_key].close_partition(partition=partition)
         with self._lock:
             self._semaphore_per_partition[partition_key].acquire()
             cursor = self._cursor_per_partition[partition_key]
             if (
                 partition_key in self._finished_partitions
                 and self._semaphore_per_partition[partition_key]._value == 0
             ):
-                if (
-                    self._new_global_cursor is None
-                    or self._new_global_cursor[self.cursor_field.cursor_field_key]
-                    < cursor.state[self.cursor_field.cursor_field_key]
-                ):
-                    self._new_global_cursor = copy.deepcopy(cursor.state)
+                self._update_global_cursor(cursor.state[self.cursor_field.cursor_field_key])
             if not self._use_global_cursor:
                 self._emit_state_message()
 
@@ -366,9 +362,22 @@ def observe(self, record: Record) -> None:
             raise ValueError(
                 "Invalid state as stream slices that are emitted should refer to an existing cursor"
             )
-        self._cursor_per_partition[
-            self._to_partition_key(record.associated_slice.partition)
-        ].observe(record)
+
+        record_cursor = self._connector_state_converter.parse_value(
+            self._cursor_field.extract_value(record)
+        )
+        self._update_global_cursor(record_cursor)
+        if not self._use_global_cursor:
+            self._cursor_per_partition[
+                self._to_partition_key(record.associated_slice.partition)
+            ].observe(record)
+
+    def _update_global_cursor(self, value: Mapping[str, Any]) -> None:
+        if (
+            self._new_global_cursor is None
+            or self._new_global_cursor[self.cursor_field.cursor_field_key] < value
+        ):
+            self._new_global_cursor = {self.cursor_field.cursor_field_key: copy.deepcopy(value)}
 
     def _to_partition_key(self, partition: Mapping[str, Any]) -> str:
         return self._partition_serializer.to_partition_key(partition)

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -137,10 +137,6 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
             self._source_config, config
         )
 
-        api_budget_model = self._source_config.get("api_budget")
-        if api_budget_model:
-            self._constructor.set_api_budget(api_budget_model, config)
-
         source_streams = [
             self._constructor.create_component(
                 DeclarativeStreamModel,