Add exception if stream_slices was executed two times

airbytehq · Jan 16, 2025 · daa6873 · daa6873
1 parent df0993e
commit daa6873
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py b/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py
@@ -1,9 +1,9 @@
-import copy
-import logging
-
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
+
+import copy
+import logging
 import threading
 from collections import OrderedDict
 from copy import deepcopy
@@ -171,6 +171,9 @@ def _emit_state_message(self) -> None:
         self._message_repository.emit_message(state_message)
 
     def stream_slices(self) -> Iterable[StreamSlice]:
+        if self._timer.is_running():
+            raise RuntimeError("stream_slices has been executed more than once.")
+
         slices = self._partition_router.stream_slices()
         self._timer.start()
         for partition in slices:

diff --git a/airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py b/airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py
@@ -64,6 +64,9 @@ def finish(self) -> int:
         else:
             raise RuntimeError("Global substream cursor timer not started")
 
+    def is_running(self) -> bool:
+        return self._start is not None
+
 
 class GlobalSubstreamCursor(DeclarativeCursor):
     """