airbyte-cdk 6.54.10__py3-none-any.whl → 6.54.11__py3-none-any.whl

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -9,7 +9,7 @@ import time
 from collections import OrderedDict
 from copy import deepcopy
 from datetime import timedelta
-from typing import Any, Callable, Iterable, Mapping, MutableMapping, Optional
+from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional
 
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
@@ -66,8 +66,8 @@ class ConcurrentPerPartitionCursor(Cursor):
     _GLOBAL_STATE_KEY = "state"
     _PERPARTITION_STATE_KEY = "states"
     _IS_PARTITION_DUPLICATION_LOGGED = False
-    _KEY = 0
-    _VALUE = 1
+    _PARENT_STATE = 0
+    _GENERATION_SEQUENCE = 1
 
     def __init__(
         self,
@@ -99,19 +99,29 @@ class ConcurrentPerPartitionCursor(Cursor):
         self._semaphore_per_partition: OrderedDict[str, threading.Semaphore] = OrderedDict()
 
         # Parent-state tracking: store each partition’s parent state in creation order
-        self._partition_parent_state_map: OrderedDict[str, Mapping[str, Any]] = OrderedDict()
+        self._partition_parent_state_map: OrderedDict[str, tuple[Mapping[str, Any], int]] = (
+            OrderedDict()
+        )
+        self._parent_state: Optional[StreamState] = None
+
+        # Tracks when the last slice for partition is emitted
+        self._partitions_done_generating_stream_slices: set[str] = set()
+        # Used to track the index of partitions that are not closed yet
+        self._processing_partitions_indexes: List[int] = list()
+        self._generated_partitions_count: int = 0
+        # Dictionary to map partition keys to their index
+        self._partition_key_to_index: dict[str, int] = {}
 
-        self._finished_partitions: set[str] = set()
         self._lock = threading.Lock()
-        self._timer = Timer()
-        self._new_global_cursor: Optional[StreamState] = None
         self._lookback_window: int = 0
-        self._parent_state: Optional[StreamState] = None
+        self._new_global_cursor: Optional[StreamState] = None
         self._number_of_partitions: int = 0
         self._use_global_cursor: bool = use_global_cursor
         self._partition_serializer = PerPartitionKeySerializer()
+
         # Track the last time a state message was emitted
         self._last_emission_time: float = 0.0
+        self._timer = Timer()
 
         self._set_initial_state(stream_state)
 
@@ -157,60 +167,37 @@ class ConcurrentPerPartitionCursor(Cursor):
                 self._cursor_per_partition[partition_key].close_partition(partition=partition)
                 cursor = self._cursor_per_partition[partition_key]
                 if (
-                    partition_key in self._finished_partitions
+                    partition_key in self._partitions_done_generating_stream_slices
                     and self._semaphore_per_partition[partition_key]._value == 0
                 ):
                     self._update_global_cursor(cursor.state[self.cursor_field.cursor_field_key])
 
+            # Clean up the partition if it is fully processed
+            self._cleanup_if_done(partition_key)
+
             self._check_and_update_parent_state()
 
             self._emit_state_message()
 
     def _check_and_update_parent_state(self) -> None:
-        """
-        Pop the leftmost partition state from _partition_parent_state_map only if
-        *all partitions* up to (and including) that partition key in _semaphore_per_partition
-        are fully finished (i.e. in _finished_partitions and semaphore._value == 0).
-        Additionally, delete finished semaphores with a value of 0 to free up memory,
-        as they are only needed to track errors and completion status.
-        """
         last_closed_state = None
 
         while self._partition_parent_state_map:
-            # Look at the earliest partition key in creation order
-            earliest_key = next(iter(self._partition_parent_state_map))
-
-            # Verify ALL partitions from the left up to earliest_key are finished
-            all_left_finished = True
-            for p_key, sem in list(
-                self._semaphore_per_partition.items()
-            ):  # Use list to allow modification during iteration
-                # If any earlier partition is still not finished, we must stop
-                if p_key not in self._finished_partitions or sem._value != 0:
-                    all_left_finished = False
-                    break
-                # Once we've reached earliest_key in the semaphore order, we can stop checking
-                if p_key == earliest_key:
-                    break
-
-            # If the partitions up to earliest_key are not all finished, break the while-loop
-            if not all_left_finished:
-                break
+            earliest_key, (candidate_state, candidate_seq) = next(
+                iter(self._partition_parent_state_map.items())
+            )
 
-            # Pop the leftmost entry from parent-state map
-            _, closed_parent_state = self._partition_parent_state_map.popitem(last=False)
-            last_closed_state = closed_parent_state
+            # if any partition that started <= candidate_seq is still open, we must wait
+            if (
+                self._processing_partitions_indexes
+                and self._processing_partitions_indexes[0] <= candidate_seq
+            ):
+                break
 
-            # Clean up finished semaphores with value 0 up to and including earliest_key
-            for p_key in list(self._semaphore_per_partition.keys()):
-                sem = self._semaphore_per_partition[p_key]
-                if p_key in self._finished_partitions and sem._value == 0:
-                    del self._semaphore_per_partition[p_key]
-                    logger.debug(f"Deleted finished semaphore for partition {p_key} with value 0")
-                if p_key == earliest_key:
-                    break
+            # safe to pop
+            self._partition_parent_state_map.popitem(last=False)
+            last_closed_state = candidate_state
 
-        # Update _parent_state if we popped at least one partition
         if last_closed_state is not None:
             self._parent_state = last_closed_state
 
@@ -289,18 +276,24 @@ class ConcurrentPerPartitionCursor(Cursor):
             if not self._IS_PARTITION_DUPLICATION_LOGGED:
                 logger.warning(f"Partition duplication detected for stream {self._stream_name}")
                 self._IS_PARTITION_DUPLICATION_LOGGED = True
+            return
         else:
             self._semaphore_per_partition[partition_key] = threading.Semaphore(0)
 
         with self._lock:
+            seq = self._generated_partitions_count
+            self._generated_partitions_count += 1
+            self._processing_partitions_indexes.append(seq)
+            self._partition_key_to_index[partition_key] = seq
+
             if (
                 len(self._partition_parent_state_map) == 0
                 or self._partition_parent_state_map[
                     next(reversed(self._partition_parent_state_map))
-                ]
+                ][self._PARENT_STATE]
                 != parent_state
             ):
-                self._partition_parent_state_map[partition_key] = deepcopy(parent_state)
+                self._partition_parent_state_map[partition_key] = (deepcopy(parent_state), seq)
 
         for cursor_slice, is_last_slice, _ in iterate_with_last_flag_and_state(
             cursor.stream_slices(),
@@ -308,7 +301,7 @@ class ConcurrentPerPartitionCursor(Cursor):
         ):
             self._semaphore_per_partition[partition_key].release()
             if is_last_slice:
-                self._finished_partitions.add(partition_key)
+                self._partitions_done_generating_stream_slices.add(partition_key)
             yield StreamSlice(
                 partition=partition, cursor_slice=cursor_slice, extra_fields=partition.extra_fields
             )
@@ -338,14 +331,11 @@ class ConcurrentPerPartitionCursor(Cursor):
             while len(self._cursor_per_partition) > self.DEFAULT_MAX_PARTITIONS_NUMBER - 1:
                 # Try removing finished partitions first
                 for partition_key in list(self._cursor_per_partition.keys()):
-                    if partition_key in self._finished_partitions and (
-                        partition_key not in self._semaphore_per_partition
-                        or self._semaphore_per_partition[partition_key]._value == 0
-                    ):
+                    if partition_key not in self._partition_key_to_index:
                         oldest_partition = self._cursor_per_partition.pop(
                             partition_key
                         )  # Remove the oldest partition
-                        logger.warning(
+                        logger.debug(
                             f"The maximum number of partitions has been reached. Dropping the oldest finished partition: {oldest_partition}. Over limit: {self._number_of_partitions - self.DEFAULT_MAX_PARTITIONS_NUMBER}."
                         )
                         break
@@ -474,6 +464,25 @@ class ConcurrentPerPartitionCursor(Cursor):
         ):
             self._new_global_cursor = {self.cursor_field.cursor_field_key: copy.deepcopy(value)}
 
+    def _cleanup_if_done(self, partition_key: str) -> None:
+        """
+        Free every in-memory structure that belonged to a completed partition:
+        cursor, semaphore, flag inside `_finished_partitions`
+        """
+        if not (
+            partition_key in self._partitions_done_generating_stream_slices
+            and self._semaphore_per_partition[partition_key]._value == 0
+        ):
+            return
+
+        self._semaphore_per_partition.pop(partition_key, None)
+        self._partitions_done_generating_stream_slices.discard(partition_key)
+
+        seq = self._partition_key_to_index.pop(partition_key)
+        self._processing_partitions_indexes.remove(seq)
+
+        logger.debug(f"Partition {partition_key} fully processed and cleaned up.")
+
     def _to_partition_key(self, partition: Mapping[str, Any]) -> str:
         return self._partition_serializer.to_partition_key(partition)
 

airbyte_cdk/sources/declarative/interpolation/jinja.py

@@ -147,16 +147,18 @@ class JinjaInterpolation(Interpolation):
             # It can be returned as is
             return s
 
+    @staticmethod
     @cache
-    def _find_undeclared_variables(self, s: Optional[str]) -> Set[str]:
+    def _find_undeclared_variables(s: Optional[str]) -> Set[str]:
         """
         Find undeclared variables and cache them
         """
         ast = _ENVIRONMENT.parse(s)  # type: ignore # parse is able to handle None
         return meta.find_undeclared_variables(ast)
 
+    @staticmethod
     @cache
-    def _compile(self, s: str) -> Template:
+    def _compile(s: str) -> Template:
         """
         We must cache the Jinja Template ourselves because we're using `from_string` instead of a template loader
         """

{airbyte_cdk-6.54.10.dist-info → airbyte_cdk-6.54.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 6.54.10
+Version: 6.54.11
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://airbyte.com
 License: MIT

{airbyte_cdk-6.54.10.dist-info → airbyte_cdk-6.54.11.dist-info}/RECORD

@@ -111,7 +111,7 @@ airbyte_cdk/sources/declarative/extractors/record_selector.py,sha256=vCpwX1PVRFP
 airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py,sha256=WJyA2OYIEgFpVP5Y3o0tIj69AV6IKkn9B16MeXaEItI,6513
 airbyte_cdk/sources/declarative/extractors/type_transformer.py,sha256=d6Y2Rfg8pMVEEnHllfVksWZdNVOU55yk34O03dP9muY,1626
 airbyte_cdk/sources/declarative/incremental/__init__.py,sha256=U1oZKtBaEC6IACmvziY9Wzg7Z8EgF4ZuR7NwvjlB_Sk,1255
-airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py,sha256=8zxfM1aXSwdfTbpL9PHwgg43ZRbNUQRcX038D7YTI_0,22768
+airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py,sha256=araWk039M89c6lQHEUltfM1VI_xGw9gZIDXRWWF6SkM,22591
 airbyte_cdk/sources/declarative/incremental/datetime_based_cursor.py,sha256=Rbe6lJLTtZ5en33MwZiB9-H9-AwDMNHgwBZs8EqhYqk,22172
 airbyte_cdk/sources/declarative/incremental/declarative_cursor.py,sha256=5Bhw9VRPyIuCaD0wmmq_L3DZsa-rJgtKSEUzSd8YYD0,536
 airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py,sha256=2tsE6FgXzemf4fZZ4uGtd8QpRBl9GJ2CRqSNJE5p0EI,16077
@@ -125,7 +125,7 @@ airbyte_cdk/sources/declarative/interpolation/interpolated_mapping.py,sha256=h36
 airbyte_cdk/sources/declarative/interpolation/interpolated_nested_mapping.py,sha256=myVaNtFqxOAwrbp93rgd1dhkqyuvXvET9rsimQ89ktc,1873
 airbyte_cdk/sources/declarative/interpolation/interpolated_string.py,sha256=CQkHqGlfa87G6VYMtBAQWin7ECKpfMdrDcg0JO5_rhc,3212
 airbyte_cdk/sources/declarative/interpolation/interpolation.py,sha256=9IoeuWam3L6GyN10L6U8xNWXmkt9cnahSDNkez1OmFY,982
-airbyte_cdk/sources/declarative/interpolation/jinja.py,sha256=wm_W3_usYU_3eUdSpzrj21qsk8-vcK8TtWEHQsH7SY4,7245
+airbyte_cdk/sources/declarative/interpolation/jinja.py,sha256=oFGKs3oX0xO6DOL4E9x8rhxwbEoRcgx4HJVIL1RQ9c4,7269
 airbyte_cdk/sources/declarative/interpolation/macros.py,sha256=xRcmjape4_WGmKMJpmBsKY0k4OHJDM46Hv3V-dlSz3w,5640
 airbyte_cdk/sources/declarative/manifest_declarative_source.py,sha256=ciXtM7Qhus170ZwP8B9Ac4VScX2FPBYvlbZRv_r376U,24692
 airbyte_cdk/sources/declarative/migrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -421,9 +421,9 @@ airbyte_cdk/utils/slice_hasher.py,sha256=EDxgROHDbfG-QKQb59m7h_7crN1tRiawdf5uU7G
 airbyte_cdk/utils/spec_schema_transformations.py,sha256=-5HTuNsnDBAhj-oLeQXwpTGA0HdcjFOf2zTEMUTTg_Y,816
 airbyte_cdk/utils/stream_status_utils.py,sha256=ZmBoiy5HVbUEHAMrUONxZvxnvfV9CesmQJLDTAIWnWw,1171
 airbyte_cdk/utils/traced_exception.py,sha256=C8uIBuCL_E4WnBAOPSxBicD06JAldoN9fGsQDp463OY,6292
-airbyte_cdk-6.54.10.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
-airbyte_cdk-6.54.10.dist-info/LICENSE_SHORT,sha256=aqF6D1NcESmpn-cqsxBtszTEnHKnlsp8L4x9wAh3Nxg,55
-airbyte_cdk-6.54.10.dist-info/METADATA,sha256=D26T_w15_WHbsqtSJzrHNySmgZg9FnltqLhbESGxdDg,6344
-airbyte_cdk-6.54.10.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-airbyte_cdk-6.54.10.dist-info/entry_points.txt,sha256=AKWbEkHfpzzk9nF9tqBUaw1MbvTM4mGtEzmZQm0ZWvM,139
-airbyte_cdk-6.54.10.dist-info/RECORD,,
+airbyte_cdk-6.54.11.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
+airbyte_cdk-6.54.11.dist-info/LICENSE_SHORT,sha256=aqF6D1NcESmpn-cqsxBtszTEnHKnlsp8L4x9wAh3Nxg,55
+airbyte_cdk-6.54.11.dist-info/METADATA,sha256=GAc-NcvDVzTVP4sUrNphva1hP0kyZt-ANEM5Qz5CgPY,6344
+airbyte_cdk-6.54.11.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+airbyte_cdk-6.54.11.dist-info/entry_points.txt,sha256=AKWbEkHfpzzk9nF9tqBUaw1MbvTM4mGtEzmZQm0ZWvM,139
+airbyte_cdk-6.54.11.dist-info/RECORD,,