removing lock from concurrent cursor

Author: maxi297

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -96,13 +96,12 @@ def on_partition(self, partition: Partition) -> None:
         """
         stream_name = partition.stream_name()
         self._streams_to_running_partitions[stream_name].add(partition)
-        cursor = self._stream_name_to_instance[stream_name].cursor
         if self._slice_logger.should_log_slice_message(self._logger):
             self._message_repository.emit_message(
                 self._slice_logger.create_slice_log_message(partition.to_slice())
             )
         self._thread_pool_manager.submit(
-            self._partition_reader.process_partition, partition, cursor
+            self._partition_reader.process_partition, partition
         )
 
     def on_partition_complete_sentinel(
@@ -116,16 +115,26 @@ def on_partition_complete_sentinel(
         """
         partition = sentinel.partition
 
-        partitions_running = self._streams_to_running_partitions[partition.stream_name()]
-        if partition in partitions_running:
-            partitions_running.remove(partition)
-            # If all partitions were generated and this was the last one, the stream is done
-            if (
-                partition.stream_name() not in self._streams_currently_generating_partitions
-                and len(partitions_running) == 0
-            ):
-                yield from self._on_stream_is_done(partition.stream_name())
-        yield from self._message_repository.consume_queue()
+        try:
+            if sentinel.is_successful:
+                stream = self._stream_name_to_instance[partition.stream_name()]
+                stream.cursor.close_partition(partition)
+        except Exception as exception:
+            self._flag_exception(partition.stream_name(), exception)
+            yield AirbyteTracedException.from_exception(
+                exception, stream_descriptor=StreamDescriptor(name=partition.stream_name())
+            ).as_sanitized_airbyte_message()
+        finally:
+            partitions_running = self._streams_to_running_partitions[partition.stream_name()]
+            if partition in partitions_running:
+                partitions_running.remove(partition)
+                # If all partitions were generated and this was the last one, the stream is done
+                if (
+                    partition.stream_name() not in self._streams_currently_generating_partitions
+                    and len(partitions_running) == 0
+                ):
+                    yield from self._on_stream_is_done(partition.stream_name())
+            yield from self._message_repository.consume_queue()
 
     def on_record(self, record: Record) -> Iterable[AirbyteMessage]:
         """
@@ -154,6 +163,7 @@ def on_record(self, record: Record) -> Iterable[AirbyteMessage]:
                     stream.as_airbyte_stream(), AirbyteStreamStatus.RUNNING
                 )
             self._record_counter[stream.name] += 1
+            stream.cursor.observe(record)
         yield message
         yield from self._message_repository.consume_queue()
 

airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -198,7 +198,7 @@ def __init__(
         # not thread safe. When multiple partitions are being closed by the cursor at the same time, it is
         # possible for one partition to update concurrent_state after a second partition has already read
         # the previous state. This can lead to the second partition overwriting the previous one's state.
-        self._lock = threading.Lock()
+        # self._lock = threading.Lock()
 
     @property
     def state(self) -> MutableMapping[str, Any]:
@@ -273,14 +273,14 @@ def _extract_cursor_value(self, record: Record) -> Any:
         return self._connector_state_converter.parse_value(self._cursor_field.extract_value(record))
 
     def close_partition(self, partition: Partition) -> None:
-        with self._lock:
-            slice_count_before = len(self._concurrent_state.get("slices", []))
-            self._add_slice_to_state(partition)
-            if slice_count_before < len(
-                self._concurrent_state["slices"]
-            ):  # only emit if at least one slice has been processed
-                self._merge_partitions()
-                self._emit_state_message()
+        #with self._lock:
+        slice_count_before = len(self._concurrent_state.get("slices", []))
+        self._add_slice_to_state(partition)
+        if slice_count_before < len(
+            self._concurrent_state["slices"]
+        ):  # only emit if at least one slice has been processed
+            self._merge_partitions()
+            self._emit_state_message()
         self._has_closed_at_least_one_slice = True
 
     def _add_slice_to_state(self, partition: Partition) -> None:

airbyte_cdk/sources/streams/concurrent/partition_reader.py

@@ -60,7 +60,7 @@ def __init__(
         self._queue = queue
         self._partition_logger = partition_logger
 
-    def process_partition(self, partition: Partition, cursor: Cursor) -> None:
+    def process_partition(self, partition: Partition) -> None:
         """
         Process a partition and put the records in the output queue.
         When all the partitions are added to the queue, a sentinel is added to the queue to indicate that all the partitions have been generated.
@@ -78,8 +78,6 @@ def process_partition(self, partition: Partition, cursor: Cursor) -> None:
 
             for record in partition.read():
                 self._queue.put(record)
-                cursor.observe(record)
-            cursor.close_partition(partition)
             self._queue.put(PartitionCompleteSentinel(partition, self._IS_SUCCESSFUL))
         except Exception as e:
             self._queue.put(StreamThreadException(e, partition.stream_name()))