airbyte-cdk 0.72.0__py3-none-any.whl → 6.13.1.dev4106__py3-none-any.whl

airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_concurrent_cursor.py

@@ -11,13 +11,15 @@ from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level, Type
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
-from airbyte_cdk.sources.file_based.stream.concurrent.cursor.abstract_concurrent_file_based_cursor import AbstractConcurrentFileBasedCursor
+from airbyte_cdk.sources.file_based.stream.concurrent.cursor.abstract_concurrent_file_based_cursor import (
+    AbstractConcurrentFileBasedCursor,
+)
 from airbyte_cdk.sources.file_based.stream.cursor import DefaultFileBasedCursor
 from airbyte_cdk.sources.file_based.types import StreamState
 from airbyte_cdk.sources.message.repository import MessageRepository
 from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.types import Record
 
 if TYPE_CHECKING:
     from airbyte_cdk.sources.file_based.stream.concurrent.adapters import FileBasedStreamPartition
@@ -27,7 +29,9 @@ _NULL_FILE = ""
 
 class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
     CURSOR_FIELD = "_ab_source_file_last_modified"
-    DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL = DefaultFileBasedCursor.DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL
+    DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL = (
+        DefaultFileBasedCursor.DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL
+    )
     DEFAULT_MAX_HISTORY_SIZE = 10_000
     DATE_TIME_FORMAT = DefaultFileBasedCursor.DATE_TIME_FORMAT
     zero_value = datetime.min
@@ -51,7 +55,8 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
         self._connector_state_manager = connector_state_manager
         self._cursor_field = cursor_field
         self._time_window_if_history_is_full = timedelta(
-            days=stream_config.days_to_sync_if_history_is_full or self.DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL
+            days=stream_config.days_to_sync_if_history_is_full
+            or self.DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL
         )
         self._state_lock = RLock()
         self._pending_files_lock = RLock()
@@ -70,7 +75,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
     def close_partition(self, partition: Partition) -> None:
         with self._pending_files_lock:
             if self._pending_files is None:
-                raise RuntimeError("Expected pending partitions to be set but it was not. This is unexpected. Please contact Support.")
+                raise RuntimeError(
+                    "Expected pending partitions to be set but it was not. This is unexpected. Please contact Support."
+                )
 
     def set_pending_partitions(self, partitions: List["FileBasedStreamPartition"]) -> None:
         with self._pending_files_lock:
@@ -81,7 +88,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
                     continue
                 for file in _slice["files"]:
                     if file.uri in self._pending_files.keys():
-                        raise RuntimeError(f"Already found file {_slice} in pending files. This is unexpected. Please contact Support.")
+                        raise RuntimeError(
+                            f"Already found file {_slice} in pending files. This is unexpected. Please contact Support."
+                        )
                 self._pending_files.update({file.uri: file})
 
     def _compute_prev_sync_cursor(self, value: Optional[StreamState]) -> Tuple[datetime, str]:
@@ -96,7 +105,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
         # represents the start time that the file was uploaded, we can usually expect that all previous
         # files have already been uploaded. If that's the case, they'll be in history and we'll skip
         # re-uploading them.
-        earliest_file_cursor_value = self._get_cursor_key_from_file(self._compute_earliest_file_in_history())
+        earliest_file_cursor_value = self._get_cursor_key_from_file(
+            self._compute_earliest_file_in_history()
+        )
         cursor_str = min(prev_cursor_str, earliest_file_cursor_value)
         cursor_dt, cursor_uri = cursor_str.split("_", 1)
         return datetime.strptime(cursor_dt, self.DATE_TIME_FORMAT), cursor_uri
@@ -109,8 +120,13 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
     def _compute_earliest_file_in_history(self) -> Optional[RemoteFile]:
         with self._state_lock:
             if self._file_to_datetime_history:
-                filename, last_modified = min(self._file_to_datetime_history.items(), key=lambda f: (f[1], f[0]))
-                return RemoteFile(uri=filename, last_modified=datetime.strptime(last_modified, self.DATE_TIME_FORMAT))
+                filename, last_modified = min(
+                    self._file_to_datetime_history.items(), key=lambda f: (f[1], f[0])
+                )
+                return RemoteFile(
+                    uri=filename,
+                    last_modified=datetime.strptime(last_modified, self.DATE_TIME_FORMAT),
+                )
             else:
                 return None
 
@@ -120,7 +136,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
         :param file: The file to add
         """
         if self._pending_files is None:
-            raise RuntimeError("Expected pending partitions to be set but it was not. This is unexpected. Please contact Support.")
+            raise RuntimeError(
+                "Expected pending partitions to be set but it was not. This is unexpected. Please contact Support."
+            )
         with self._pending_files_lock:
             with self._state_lock:
                 if file.uri not in self._pending_files:
@@ -135,7 +153,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
                     )
                 else:
                     self._pending_files.pop(file.uri)
-                self._file_to_datetime_history[file.uri] = file.last_modified.strftime(self.DATE_TIME_FORMAT)
+                self._file_to_datetime_history[file.uri] = file.last_modified.strftime(
+                    self.DATE_TIME_FORMAT
+                )
                 if len(self._file_to_datetime_history) > self.DEFAULT_MAX_HISTORY_SIZE:
                     # Get the earliest file based on its last modified date and its uri
                     oldest_file = self._compute_earliest_file_in_history()
@@ -155,7 +175,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
                 self._stream_namespace,
                 new_state,
             )
-            state_message = self._connector_state_manager.create_state_message(self._stream_name, self._stream_namespace)
+            state_message = self._connector_state_manager.create_state_message(
+                self._stream_name, self._stream_namespace
+            )
             self._message_repository.emit_message(state_message)
 
     def _get_new_cursor_value(self) -> str:
@@ -183,12 +205,19 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
     def _compute_latest_file_in_history(self) -> Optional[RemoteFile]:
         with self._state_lock:
             if self._file_to_datetime_history:
-                filename, last_modified = max(self._file_to_datetime_history.items(), key=lambda f: (f[1], f[0]))
-                return RemoteFile(uri=filename, last_modified=datetime.strptime(last_modified, self.DATE_TIME_FORMAT))
+                filename, last_modified = max(
+                    self._file_to_datetime_history.items(), key=lambda f: (f[1], f[0])
+                )
+                return RemoteFile(
+                    uri=filename,
+                    last_modified=datetime.strptime(last_modified, self.DATE_TIME_FORMAT),
+                )
             else:
                 return None
 
-    def get_files_to_sync(self, all_files: Iterable[RemoteFile], logger: logging.Logger) -> Iterable[RemoteFile]:
+    def get_files_to_sync(
+        self, all_files: Iterable[RemoteFile], logger: logging.Logger
+    ) -> Iterable[RemoteFile]:
         """
         Given the list of files in the source, return the files that should be synced.
         :param all_files: All files in the source
@@ -210,7 +239,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
         with self._state_lock:
             if file.uri in self._file_to_datetime_history:
                 # If the file's uri is in the history, we should sync the file if it has been modified since it was synced
-                updated_at_from_history = datetime.strptime(self._file_to_datetime_history[file.uri], self.DATE_TIME_FORMAT)
+                updated_at_from_history = datetime.strptime(
+                    self._file_to_datetime_history[file.uri], self.DATE_TIME_FORMAT
+                )
                 if file.last_modified < updated_at_from_history:
                     self._message_repository.emit_message(
                         AirbyteMessage(
@@ -246,7 +277,9 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
         """
         with self._state_lock:
             if self._file_to_datetime_history is None:
-                raise RuntimeError("The history object has not been set. This is unexpected. Please contact Support.")
+                raise RuntimeError(
+                    "The history object has not been set. This is unexpected. Please contact Support."
+                )
             return len(self._file_to_datetime_history) >= self.DEFAULT_MAX_HISTORY_SIZE
 
     def _compute_start_time(self) -> datetime:
@@ -268,7 +301,10 @@ class FileBasedConcurrentCursor(AbstractConcurrentFileBasedCursor):
         Get the state of the cursor.
         """
         with self._state_lock:
-            return {"history": self._file_to_datetime_history, self._cursor_field.cursor_field_key: self._get_new_cursor_value()}
+            return {
+                "history": self._file_to_datetime_history,
+                self._cursor_field.cursor_field_key: self._get_new_cursor_value(),
+            }
 
     def set_initial_state(self, value: StreamState) -> None:
         pass

airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_final_state_cursor.py

@@ -9,12 +9,14 @@ from typing import TYPE_CHECKING, Any, Iterable, List, MutableMapping, Optional
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
-from airbyte_cdk.sources.file_based.stream.concurrent.cursor.abstract_concurrent_file_based_cursor import AbstractConcurrentFileBasedCursor
+from airbyte_cdk.sources.file_based.stream.concurrent.cursor.abstract_concurrent_file_based_cursor import (
+    AbstractConcurrentFileBasedCursor,
+)
 from airbyte_cdk.sources.file_based.types import StreamState
 from airbyte_cdk.sources.message import MessageRepository
-from airbyte_cdk.sources.streams import FULL_REFRESH_SENTINEL_STATE_KEY
+from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.types import Record
 
 if TYPE_CHECKING:
     from airbyte_cdk.sources.file_based.stream.concurrent.adapters import FileBasedStreamPartition
@@ -24,7 +26,11 @@ class FileBasedFinalStateCursor(AbstractConcurrentFileBasedCursor):
     """Cursor that is used to guarantee at least one state message is emitted for a concurrent file-based stream."""
 
     def __init__(
-        self, stream_config: FileBasedStreamConfig, message_repository: MessageRepository, stream_namespace: Optional[str], **kwargs: Any
+        self,
+        stream_config: FileBasedStreamConfig,
+        message_repository: MessageRepository,
+        stream_namespace: Optional[str],
+        **kwargs: Any,
     ):
         self._stream_name = stream_config.name
         self._stream_namespace = stream_namespace
@@ -32,11 +38,11 @@ class FileBasedFinalStateCursor(AbstractConcurrentFileBasedCursor):
         # Normally the connector state manager operates at the source-level. However, we only need it to write the sentinel
         # state message rather than manage overall source state. This is also only temporary as we move to the resumable
         # full refresh world where every stream uses a FileBasedConcurrentCursor with incremental state.
-        self._connector_state_manager = ConnectorStateManager(stream_instance_map={})
+        self._connector_state_manager = ConnectorStateManager()
 
     @property
     def state(self) -> MutableMapping[str, Any]:
-        return {FULL_REFRESH_SENTINEL_STATE_KEY: True}
+        return {NO_CURSOR_STATE_KEY: True}
 
     def observe(self, record: Record) -> None:
         pass
@@ -50,7 +56,9 @@ class FileBasedFinalStateCursor(AbstractConcurrentFileBasedCursor):
     def add_file(self, file: RemoteFile) -> None:
         pass
 
-    def get_files_to_sync(self, all_files: Iterable[RemoteFile], logger: logging.Logger) -> Iterable[RemoteFile]:
+    def get_files_to_sync(
+        self, all_files: Iterable[RemoteFile], logger: logging.Logger
+    ) -> Iterable[RemoteFile]:
         return all_files
 
     def get_state(self) -> MutableMapping[str, Any]:
@@ -66,6 +74,10 @@ class FileBasedFinalStateCursor(AbstractConcurrentFileBasedCursor):
         pass
 
     def ensure_at_least_one_state_emitted(self) -> None:
-        self._connector_state_manager.update_state_for_stream(self._stream_name, self._stream_namespace, self.state)
-        state_message = self._connector_state_manager.create_state_message(self._stream_name, self._stream_namespace)
+        self._connector_state_manager.update_state_for_stream(
+            self._stream_name, self._stream_namespace, self.state
+        )
+        state_message = self._connector_state_manager.create_state_message(
+            self._stream_name, self._stream_namespace
+        )
         self._message_repository.emit_message(state_message)

airbyte_cdk/sources/file_based/stream/cursor/abstract_file_based_cursor.py

@@ -54,7 +54,9 @@ class AbstractFileBasedCursor(ABC):
         ...
 
     @abstractmethod
-    def get_files_to_sync(self, all_files: Iterable[RemoteFile], logger: logging.Logger) -> Iterable[RemoteFile]:
+    def get_files_to_sync(
+        self, all_files: Iterable[RemoteFile], logger: logging.Logger
+    ) -> Iterable[RemoteFile]:
         """
         Given the list of files in the source, return the files that should be synced.
         :param all_files: All files in the source

airbyte_cdk/sources/file_based/stream/cursor/default_file_based_cursor.py

@@ -8,7 +8,9 @@ from typing import Any, Iterable, MutableMapping, Optional
 
 from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
-from airbyte_cdk.sources.file_based.stream.cursor.abstract_file_based_cursor import AbstractFileBasedCursor
+from airbyte_cdk.sources.file_based.stream.cursor.abstract_file_based_cursor import (
+    AbstractFileBasedCursor,
+)
 from airbyte_cdk.sources.file_based.types import StreamState
 
 
@@ -19,14 +21,17 @@ class DefaultFileBasedCursor(AbstractFileBasedCursor):
     CURSOR_FIELD = "_ab_source_file_last_modified"
 
     def __init__(self, stream_config: FileBasedStreamConfig, **_: Any):
-        super().__init__(stream_config)
+        super().__init__(stream_config)  # type: ignore [safe-super]
         self._file_to_datetime_history: MutableMapping[str, str] = {}
         self._time_window_if_history_is_full = timedelta(
-            days=stream_config.days_to_sync_if_history_is_full or self.DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL
+            days=stream_config.days_to_sync_if_history_is_full
+            or self.DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL
         )
 
         if self._time_window_if_history_is_full <= timedelta():
-            raise ValueError(f"days_to_sync_if_history_is_full must be a positive timedelta, got {self._time_window_if_history_is_full}")
+            raise ValueError(
+                f"days_to_sync_if_history_is_full must be a positive timedelta, got {self._time_window_if_history_is_full}"
+            )
 
         self._start_time = self._compute_start_time()
         self._initial_earliest_file_in_history: Optional[RemoteFile] = None
@@ -37,7 +42,9 @@ class DefaultFileBasedCursor(AbstractFileBasedCursor):
         self._initial_earliest_file_in_history = self._compute_earliest_file_in_history()
 
     def add_file(self, file: RemoteFile) -> None:
-        self._file_to_datetime_history[file.uri] = file.last_modified.strftime(self.DATE_TIME_FORMAT)
+        self._file_to_datetime_history[file.uri] = file.last_modified.strftime(
+            self.DATE_TIME_FORMAT
+        )
         if len(self._file_to_datetime_history) > self.DEFAULT_MAX_HISTORY_SIZE:
             # Get the earliest file based on its last modified date and its uri
             oldest_file = self._compute_earliest_file_in_history()
@@ -60,7 +67,9 @@ class DefaultFileBasedCursor(AbstractFileBasedCursor):
         a string joining the last-modified timestamp of the last synced file and the name of the file.
         """
         if self._file_to_datetime_history.items():
-            filename, timestamp = max(self._file_to_datetime_history.items(), key=lambda x: (x[1], x[0]))
+            filename, timestamp = max(
+                self._file_to_datetime_history.items(), key=lambda x: (x[1], x[0])
+            )
             return f"{timestamp}_{filename}"
         return None
 
@@ -73,7 +82,9 @@ class DefaultFileBasedCursor(AbstractFileBasedCursor):
     def _should_sync_file(self, file: RemoteFile, logger: logging.Logger) -> bool:
         if file.uri in self._file_to_datetime_history:
             # If the file's uri is in the history, we should sync the file if it has been modified since it was synced
-            updated_at_from_history = datetime.strptime(self._file_to_datetime_history[file.uri], self.DATE_TIME_FORMAT)
+            updated_at_from_history = datetime.strptime(
+                self._file_to_datetime_history[file.uri], self.DATE_TIME_FORMAT
+            )
             if file.last_modified < updated_at_from_history:
                 logger.warning(
                     f"The file {file.uri}'s last modified date is older than the last time it was synced. This is unexpected. Skipping the file."
@@ -99,7 +110,9 @@ class DefaultFileBasedCursor(AbstractFileBasedCursor):
             # The file is not in the history and the history is complete. We know we need to sync the file
             return True
 
-    def get_files_to_sync(self, all_files: Iterable[RemoteFile], logger: logging.Logger) -> Iterable[RemoteFile]:
+    def get_files_to_sync(
+        self, all_files: Iterable[RemoteFile], logger: logging.Logger
+    ) -> Iterable[RemoteFile]:
         if self._is_history_full():
             logger.warning(
                 f"The state history is full. "
@@ -115,8 +128,12 @@ class DefaultFileBasedCursor(AbstractFileBasedCursor):
 
     def _compute_earliest_file_in_history(self) -> Optional[RemoteFile]:
         if self._file_to_datetime_history:
-            filename, last_modified = min(self._file_to_datetime_history.items(), key=lambda f: (f[1], f[0]))
-            return RemoteFile(uri=filename, last_modified=datetime.strptime(last_modified, self.DATE_TIME_FORMAT))
+            filename, last_modified = min(
+                self._file_to_datetime_history.items(), key=lambda f: (f[1], f[0])
+            )
+            return RemoteFile(
+                uri=filename, last_modified=datetime.strptime(last_modified, self.DATE_TIME_FORMAT)
+            )
         else:
             return None