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

airbyte_cdk/sources/file_based/stream/default_file_based_stream.py

@@ -5,24 +5,32 @@
 import asyncio
 import itertools
 import traceback
+from collections import defaultdict
 from copy import deepcopy
 from functools import cache
-from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Set, Union
+from os import path
+from typing import Any, Dict, Iterable, List, Mapping, MutableMapping, Optional, Set, Tuple, Union
 
 from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, FailureType, Level
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.file_based.config.file_based_stream_config import PrimaryKeyType
 from airbyte_cdk.sources.file_based.exceptions import (
+    DuplicatedFilesError,
     FileBasedSourceError,
     InvalidSchemaError,
     MissingSchemaError,
-    NoFilesMatchingError,
     RecordParseError,
     SchemaInferenceError,
     StopSyncPerValidationPolicy,
 )
+from airbyte_cdk.sources.file_based.file_types import FileTransfer
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
-from airbyte_cdk.sources.file_based.schema_helpers import SchemaType, merge_schemas, schemaless_schema
+from airbyte_cdk.sources.file_based.schema_helpers import (
+    SchemaType,
+    file_transfer_schema,
+    merge_schemas,
+    schemaless_schema,
+)
 from airbyte_cdk.sources.file_based.stream import AbstractFileBasedStream
 from airbyte_cdk.sources.file_based.stream.cursor import AbstractFileBasedCursor
 from airbyte_cdk.sources.file_based.types import StreamSlice
@@ -33,17 +41,29 @@ from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
 class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
-
     """
     The default file-based stream.
     """
 
+    FILE_TRANSFER_KW = "use_file_transfer"
+    PRESERVE_DIRECTORY_STRUCTURE_KW = "preserve_directory_structure"
+    FILES_KEY = "files"
     DATE_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
     ab_last_mod_col = "_ab_source_file_last_modified"
     ab_file_name_col = "_ab_source_file_url"
+    modified = "modified"
+    source_file_url = "source_file_url"
     airbyte_columns = [ab_last_mod_col, ab_file_name_col]
+    use_file_transfer = False
+    preserve_directory_structure = True
 
     def __init__(self, **kwargs: Any):
+        if self.FILE_TRANSFER_KW in kwargs:
+            self.use_file_transfer = kwargs.pop(self.FILE_TRANSFER_KW, False)
+        if self.PRESERVE_DIRECTORY_STRUCTURE_KW in kwargs:
+            self.preserve_directory_structure = kwargs.pop(
+                self.PRESERVE_DIRECTORY_STRUCTURE_KW, True
+            )
         super().__init__(**kwargs)
 
     @property
@@ -55,28 +75,84 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
         """State setter, accept state serialized by state getter."""
         self._cursor.set_initial_state(value)
 
-    @property
+    @property  # type: ignore # mypy complains wrong type, but AbstractFileBasedCursor is parent of file-based cursors
     def cursor(self) -> Optional[AbstractFileBasedCursor]:
         return self._cursor
 
     @cursor.setter
     def cursor(self, value: AbstractFileBasedCursor) -> None:
         if self._cursor is not None:
-            raise RuntimeError(f"Cursor for stream {self.name} is already set. This is unexpected. Please contact Support.")
+            raise RuntimeError(
+                f"Cursor for stream {self.name} is already set. This is unexpected. Please contact Support."
+            )
         self._cursor = value
 
     @property
     def primary_key(self) -> PrimaryKeyType:
-        return self.config.primary_key or self.get_parser().get_parser_defined_primary_key(self.config)
+        return self.config.primary_key or self.get_parser().get_parser_defined_primary_key(
+            self.config
+        )
+
+    def _filter_schema_invalid_properties(
+        self, configured_catalog_json_schema: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        if self.use_file_transfer:
+            return {
+                "type": "object",
+                "properties": {
+                    "file_path": {"type": "string"},
+                    "file_size": {"type": "string"},
+                    self.ab_file_name_col: {"type": "string"},
+                },
+            }
+        else:
+            return super()._filter_schema_invalid_properties(configured_catalog_json_schema)
+
+    def _duplicated_files_names(
+        self, slices: List[dict[str, List[RemoteFile]]]
+    ) -> List[dict[str, List[str]]]:
+        seen_file_names: Dict[str, List[str]] = defaultdict(list)
+        for file_slice in slices:
+            for file_found in file_slice[self.FILES_KEY]:
+                file_name = path.basename(file_found.uri)
+                seen_file_names[file_name].append(file_found.uri)
+        return [
+            {file_name: paths} for file_name, paths in seen_file_names.items() if len(paths) > 1
+        ]
 
     def compute_slices(self) -> Iterable[Optional[Mapping[str, Any]]]:
         # Sort files by last_modified, uri and return them grouped by last_modified
         all_files = self.list_files()
         files_to_read = self._cursor.get_files_to_sync(all_files, self.logger)
         sorted_files_to_read = sorted(files_to_read, key=lambda f: (f.last_modified, f.uri))
-        slices = [{"files": list(group[1])} for group in itertools.groupby(sorted_files_to_read, lambda f: f.last_modified)]
+        slices = [
+            {self.FILES_KEY: list(group[1])}
+            for group in itertools.groupby(sorted_files_to_read, lambda f: f.last_modified)
+        ]
+        if slices and not self.preserve_directory_structure:
+            duplicated_files_names = self._duplicated_files_names(slices)
+            if duplicated_files_names:
+                raise DuplicatedFilesError(
+                    stream=self.name, duplicated_files_names=duplicated_files_names
+                )
         return slices
 
+    def transform_record(
+        self, record: dict[str, Any], file: RemoteFile, last_updated: str
+    ) -> dict[str, Any]:
+        # adds _ab_source_file_last_modified and _ab_source_file_url to the record
+        record[self.ab_last_mod_col] = last_updated
+        record[self.ab_file_name_col] = file.uri
+        return record
+
+    def transform_record_for_file_transfer(
+        self, record: dict[str, Any], file: RemoteFile
+    ) -> dict[str, Any]:
+        # timstamp() returns a float representing the number of seconds since the unix epoch
+        record[self.modified] = int(file.last_modified.timestamp()) * 1000
+        record[self.source_file_url] = file.uri
+        return record
+
     def read_records_from_slice(self, stream_slice: StreamSlice) -> Iterable[AirbyteMessage]:
         """
         Yield all records from all remote files in `list_files_for_this_sync`.
@@ -96,16 +172,33 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             n_skipped = line_no = 0
 
             try:
-                for record in parser.parse_records(self.config, file, self.stream_reader, self.logger, schema):
-                    line_no += 1
-                    if self.config.schemaless:
-                        record = {"data": record}
-                    elif not self.record_passes_validation_policy(record):
-                        n_skipped += 1
-                        continue
-                    record[self.ab_last_mod_col] = file_datetime_string
-                    record[self.ab_file_name_col] = file.uri
-                    yield stream_data_to_airbyte_message(self.name, record)
+                if self.use_file_transfer:
+                    self.logger.info(f"{self.name}: {file} file-based syncing")
+                    # todo: complete here the code to not rely on local parser
+                    file_transfer = FileTransfer()
+                    for record in file_transfer.get_file(
+                        self.config, file, self.stream_reader, self.logger
+                    ):
+                        line_no += 1
+                        if not self.record_passes_validation_policy(record):
+                            n_skipped += 1
+                            continue
+                        record = self.transform_record_for_file_transfer(record, file)
+                        yield stream_data_to_airbyte_message(
+                            self.name, record, is_file_transfer_message=True
+                        )
+                else:
+                    for record in parser.parse_records(
+                        self.config, file, self.stream_reader, self.logger, schema
+                    ):
+                        line_no += 1
+                        if self.config.schemaless:
+                            record = {"data": record}
+                        elif not self.record_passes_validation_policy(record):
+                            n_skipped += 1
+                            continue
+                        record = self.transform_record(record, file, file_datetime_string)
+                        yield stream_data_to_airbyte_message(self.name, record)
                 self._cursor.add_file(file)
 
             except StopSyncPerValidationPolicy:
@@ -172,49 +265,71 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
         }
         try:
             schema = self._get_raw_json_schema()
-        except (InvalidSchemaError, NoFilesMatchingError) as config_exception:
-            self.logger.exception(FileBasedSourceError.SCHEMA_INFERENCE_ERROR.value, exc_info=config_exception)
+        except InvalidSchemaError as config_exception:
             raise AirbyteTracedException(
                 internal_message="Please check the logged errors for more information.",
                 message=FileBasedSourceError.SCHEMA_INFERENCE_ERROR.value,
                 exception=AirbyteTracedException(exception=config_exception),
                 failure_type=FailureType.config_error,
             )
+        except AirbyteTracedException as ate:
+            raise ate
         except Exception as exc:
-            raise SchemaInferenceError(FileBasedSourceError.SCHEMA_INFERENCE_ERROR, stream=self.name) from exc
+            raise SchemaInferenceError(
+                FileBasedSourceError.SCHEMA_INFERENCE_ERROR, stream=self.name
+            ) from exc
         else:
             return {"type": "object", "properties": {**extra_fields, **schema["properties"]}}
 
     def _get_raw_json_schema(self) -> JsonSchema:
-        if self.config.input_schema:
+        if self.use_file_transfer:
+            return file_transfer_schema
+        elif self.config.input_schema:
             return self.config.get_input_schema()  # type: ignore
         elif self.config.schemaless:
             return schemaless_schema
         else:
             files = self.list_files()
-            total_n_files = len(files)
-
-            if total_n_files == 0:
-                raise NoFilesMatchingError(FileBasedSourceError.EMPTY_STREAM, stream=self.name)
+            first_n_files = len(files)
 
-            max_n_files_for_schema_inference = self._discovery_policy.get_max_n_files_for_schema_inference(self.get_parser())
-            if total_n_files > max_n_files_for_schema_inference:
-                # Use the most recent files for schema inference, so we pick up schema changes during discovery.
-                files = sorted(files, key=lambda x: x.last_modified, reverse=True)[:max_n_files_for_schema_inference]
-                self.logger.warn(
-                    msg=f"Refusing to infer schema for all {total_n_files} files; using {max_n_files_for_schema_inference} files."
+            if self.config.recent_n_files_to_read_for_schema_discovery:
+                self.logger.info(
+                    msg=(
+                        f"Only first {self.config.recent_n_files_to_read_for_schema_discovery} files will be used to infer schema "
+                        f"for stream {self.name} due to limitation in config."
+                    )
                 )
+                first_n_files = self.config.recent_n_files_to_read_for_schema_discovery
 
-            inferred_schema = self.infer_schema(files)
+        if first_n_files == 0:
+            self.logger.warning(
+                msg=f"No files were identified in the stream {self.name}. Setting default schema for the stream."
+            )
+            return schemaless_schema
 
-            if not inferred_schema:
-                raise InvalidSchemaError(
-                    FileBasedSourceError.INVALID_SCHEMA_ERROR,
-                    details=f"Empty schema. Please check that the files are valid for format {self.config.format}",
-                    stream=self.name,
-                )
+        max_n_files_for_schema_inference = (
+            self._discovery_policy.get_max_n_files_for_schema_inference(self.get_parser())
+        )
 
-            schema = {"type": "object", "properties": inferred_schema}
+        if first_n_files > max_n_files_for_schema_inference:
+            # Use the most recent files for schema inference, so we pick up schema changes during discovery.
+            self.logger.warning(
+                msg=f"Refusing to infer schema for {first_n_files} files; using {max_n_files_for_schema_inference} files."
+            )
+            first_n_files = max_n_files_for_schema_inference
+
+        files = sorted(files, key=lambda x: x.last_modified, reverse=True)[:first_n_files]
+
+        inferred_schema = self.infer_schema(files)
+
+        if not inferred_schema:
+            raise InvalidSchemaError(
+                FileBasedSourceError.INVALID_SCHEMA_ERROR,
+                details=f"Empty schema. Please check that the files are valid for format {self.config.format}",
+                stream=self.name,
+            )
+
+        schema = {"type": "object", "properties": inferred_schema}
 
         return schema
 
@@ -222,7 +337,9 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
         """
         Return all files that belong to the stream as defined by the stream's globs.
         """
-        return self.stream_reader.get_matching_files(self.config.globs or [], self.config.legacy_prefix, self.logger)
+        return self.stream_reader.get_matching_files(
+            self.config.globs or [], self.config.legacy_prefix, self.logger
+        )
 
     def infer_schema(self, files: List[RemoteFile]) -> Mapping[str, Any]:
         loop = asyncio.get_event_loop()
@@ -260,23 +377,36 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
         n_started, n_files = 0, len(files)
         files_iterator = iter(files)
         while pending_tasks or n_started < n_files:
-            while len(pending_tasks) <= self._discovery_policy.n_concurrent_requests and (file := next(files_iterator, None)):
+            while len(pending_tasks) <= self._discovery_policy.n_concurrent_requests and (
+                file := next(files_iterator, None)
+            ):
                 pending_tasks.add(asyncio.create_task(self._infer_file_schema(file)))
                 n_started += 1
             # Return when the first task is completed so that we can enqueue a new task as soon as the
             # number of concurrent tasks drops below the number allowed.
-            done, pending_tasks = await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED)
+            done, pending_tasks = await asyncio.wait(
+                pending_tasks, return_when=asyncio.FIRST_COMPLETED
+            )
             for task in done:
                 try:
                     base_schema = merge_schemas(base_schema, task.result())
+                except AirbyteTracedException as ate:
+                    raise ate
                 except Exception as exc:
-                    self.logger.error(f"An error occurred inferring the schema. \n {traceback.format_exc()}", exc_info=exc)
+                    self.logger.error(
+                        f"An error occurred inferring the schema. \n {traceback.format_exc()}",
+                        exc_info=exc,
+                    )
 
         return base_schema
 
     async def _infer_file_schema(self, file: RemoteFile) -> SchemaType:
         try:
-            return await self.get_parser().infer_schema(self.config, file, self.stream_reader, self.logger)
+            return await self.get_parser().infer_schema(
+                self.config, file, self.stream_reader, self.logger
+            )
+        except AirbyteTracedException as ate:
+            raise ate
         except Exception as exc:
             raise SchemaInferenceError(
                 FileBasedSourceError.SCHEMA_INFERENCE_ERROR,

airbyte_cdk/sources/http_logger.py

@@ -5,11 +5,16 @@
 from typing import Optional, Union
 
 import requests
+
 from airbyte_cdk.sources.message import LogMessage
 
 
 def format_http_message(
-    response: requests.Response, title: str, description: str, stream_name: Optional[str], is_auxiliary: bool = None
+    response: requests.Response,
+    title: str,
+    description: str,
+    stream_name: Optional[str],
+    is_auxiliary: bool | None = None,
 ) -> LogMessage:
     request = response.request
     log_message = {
@@ -37,10 +42,10 @@ def format_http_message(
         "url": {"full": request.url},
     }
     if is_auxiliary is not None:
-        log_message["http"]["is_auxiliary"] = is_auxiliary
+        log_message["http"]["is_auxiliary"] = is_auxiliary  # type: ignore [index]
     if stream_name:
         log_message["airbyte_cdk"] = {"stream": {"name": stream_name}}
-    return log_message
+    return log_message  # type: ignore [return-value]  # got "dict[str, object]", expected "dict[str, JsonType]"
 
 
 def _normalize_body_string(body_str: Optional[Union[str, bytes]]) -> Optional[str]:

airbyte_cdk/sources/message/__init__.py

@@ -10,4 +10,10 @@ from .repository import (
     NoopMessageRepository,
 )
 
-__all__ = ["InMemoryMessageRepository", "LogAppenderMessageRepositoryDecorator", "LogMessage", "MessageRepository", "NoopMessageRepository"]
+__all__ = [
+    "InMemoryMessageRepository",
+    "LogAppenderMessageRepositoryDecorator",
+    "LogMessage",
+    "MessageRepository",
+    "NoopMessageRepository",
+]

airbyte_cdk/sources/message/repository.py

@@ -28,7 +28,9 @@ _SEVERITY_BY_LOG_LEVEL = {
 
 def _is_severe_enough(threshold: Level, level: Level) -> bool:
     if threshold not in _SEVERITY_BY_LOG_LEVEL:
-        _LOGGER.warning(f"Log level {threshold} for threshold is not supported. This is probably a CDK bug. Please contact Airbyte.")
+        _LOGGER.warning(
+            f"Log level {threshold} for threshold is not supported. This is probably a CDK bug. Please contact Airbyte."
+        )
         return True
 
     if level not in _SEVERITY_BY_LOG_LEVEL:
@@ -80,7 +82,12 @@ class InMemoryMessageRepository(MessageRepository):
     def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
         if _is_severe_enough(self._log_level, level):
             self.emit_message(
-                AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=level, message=filter_secrets(json.dumps(message_provider()))))
+                AirbyteMessage(
+                    type=Type.LOG,
+                    log=AirbyteLogMessage(
+                        level=level, message=filter_secrets(json.dumps(message_provider()))
+                    ),
+                )
             )
 
     def consume_queue(self) -> Iterable[AirbyteMessage]:
@@ -89,7 +96,12 @@ class InMemoryMessageRepository(MessageRepository):
 
 
 class LogAppenderMessageRepositoryDecorator(MessageRepository):
-    def __init__(self, dict_to_append: LogMessage, decorated: MessageRepository, log_level: Level = Level.INFO):
+    def __init__(
+        self,
+        dict_to_append: LogMessage,
+        decorated: MessageRepository,
+        log_level: Level = Level.INFO,
+    ):
         self._dict_to_append = dict_to_append
         self._decorated = decorated
         self._log_level = log_level
@@ -106,7 +118,9 @@ class LogAppenderMessageRepositoryDecorator(MessageRepository):
     def consume_queue(self) -> Iterable[AirbyteMessage]:
         return self._decorated.consume_queue()
 
-    def _append_second_to_first(self, first: LogMessage, second: LogMessage, path: Optional[List[str]] = None) -> LogMessage:
+    def _append_second_to_first(
+        self, first: LogMessage, second: LogMessage, path: Optional[List[str]] = None
+    ) -> LogMessage:
         if path is None:
             path = []
 

airbyte_cdk/sources/source.py

@@ -5,27 +5,41 @@
 
 import logging
 from abc import ABC, abstractmethod
-from collections import defaultdict
-from typing import Any, Dict, Generic, Iterable, List, Mapping, MutableMapping, Optional, TypeVar, Union
+from typing import Any, Generic, Iterable, List, Mapping, Optional, TypeVar
 
 from airbyte_cdk.connector import BaseConnector, DefaultConnectorMixin, TConfig
-from airbyte_cdk.models import AirbyteCatalog, AirbyteMessage, AirbyteStateMessage, AirbyteStateType, ConfiguredAirbyteCatalog
+from airbyte_cdk.models import (
+    AirbyteCatalog,
+    AirbyteMessage,
+    AirbyteStateMessage,
+    AirbyteStateMessageSerializer,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
+)
 
 TState = TypeVar("TState")
 TCatalog = TypeVar("TCatalog")
 
 
+class ExperimentalClassWarning(DeprecationWarning):
+    pass
+
+
 class BaseSource(BaseConnector[TConfig], ABC, Generic[TConfig, TState, TCatalog]):
     @abstractmethod
-    def read_state(self, state_path: str) -> TState:
-        ...
+    def read_state(self, state_path: str) -> TState: ...
 
     @abstractmethod
-    def read_catalog(self, catalog_path: str) -> TCatalog:
-        ...
+    def read_catalog(self, catalog_path: str) -> TCatalog: ...
 
     @abstractmethod
-    def read(self, logger: logging.Logger, config: TConfig, catalog: TCatalog, state: Optional[TState] = None) -> Iterable[AirbyteMessage]:
+    def read(
+        self,
+        logger: logging.Logger,
+        config: TConfig,
+        catalog: TCatalog,
+        state: Optional[TState] = None,
+    ) -> Iterable[AirbyteMessage]:
         """
         Returns a generator of the AirbyteMessages generated by reading the source with the given configuration, catalog, and state.
         """
@@ -40,12 +54,12 @@ class BaseSource(BaseConnector[TConfig], ABC, Generic[TConfig, TState, TCatalog]
 
 class Source(
     DefaultConnectorMixin,
-    BaseSource[Mapping[str, Any], Union[List[AirbyteStateMessage], MutableMapping[str, Any]], ConfiguredAirbyteCatalog],
+    BaseSource[Mapping[str, Any], List[AirbyteStateMessage], ConfiguredAirbyteCatalog],
     ABC,
 ):
     # can be overridden to change an input state.
     @classmethod
-    def read_state(cls, state_path: str) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
+    def read_state(cls, state_path: str) -> List[AirbyteStateMessage]:
         """
         Retrieves the input state of a sync by reading from the specified JSON file. Incoming state can be deserialized into either
         a JSON object for legacy state input or as a list of AirbyteStateMessages for the per-stream state format. Regardless of the
@@ -53,39 +67,29 @@ class Source(
         :param state_path: The filepath to where the stream states are located
         :return: The complete stream state based on the connector's previous sync
         """
+        parsed_state_messages = []
         if state_path:
             state_obj = BaseConnector._read_json_file(state_path)
-            if not state_obj:
-                return cls._emit_legacy_state_format({})
-            if isinstance(state_obj, List):
-                parsed_state_messages = []
+            if state_obj:
                 for state in state_obj:  # type: ignore  # `isinstance(state_obj, List)` ensures that this is a list
-                    parsed_message = AirbyteStateMessage.parse_obj(state)
-                    if not parsed_message.stream and not parsed_message.data and not parsed_message.global_:
-                        raise ValueError("AirbyteStateMessage should contain either a stream, global, or state field")
+                    parsed_message = AirbyteStateMessageSerializer.load(state)
+                    if (
+                        not parsed_message.stream
+                        and not parsed_message.data
+                        and not parsed_message.global_
+                    ):
+                        raise ValueError(
+                            "AirbyteStateMessage should contain either a stream, global, or state field"
+                        )
                     parsed_state_messages.append(parsed_message)
-                return parsed_state_messages
-            else:
-                return cls._emit_legacy_state_format(state_obj)  # type: ignore  # assuming it is a dict
-        return cls._emit_legacy_state_format({})
-
-    @classmethod
-    def _emit_legacy_state_format(cls, state_obj: Dict[str, Any]) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
-        """
-        Existing connectors that override read() might not be able to interpret the new state format. We temporarily
-        send state in the old format for these connectors, but once all have been upgraded, this method can be removed,
-        and we can then emit state in the list format.
-        """
-        # vars(self.__class__) checks if the current class directly overrides the read() function
-        if "read" in vars(cls):
-            return defaultdict(dict, state_obj)
-        else:
-            if state_obj:
-                return [AirbyteStateMessage(type=AirbyteStateType.LEGACY, data=state_obj)]
-            else:
-                return []
+        return parsed_state_messages
 
     # can be overridden to change an input catalog
     @classmethod
     def read_catalog(cls, catalog_path: str) -> ConfiguredAirbyteCatalog:
-        return ConfiguredAirbyteCatalog.parse_obj(cls._read_json_file(catalog_path))
+        return ConfiguredAirbyteCatalogSerializer.load(cls._read_json_file(catalog_path))
+
+    @property
+    def name(self) -> str:
+        """Source name"""
+        return self.__class__.__name__

airbyte_cdk/sources/streams/__init__.py

@@ -3,6 +3,6 @@
 #
 
 # Initialize Streams Package
-from .core import FULL_REFRESH_SENTINEL_STATE_KEY, IncrementalMixin, Stream
+from .core import NO_CURSOR_STATE_KEY, CheckpointMixin, IncrementalMixin, Stream
 
-__all__ = ["FULL_REFRESH_SENTINEL_STATE_KEY", "IncrementalMixin", "Stream"]
+__all__ = ["NO_CURSOR_STATE_KEY", "IncrementalMixin", "CheckpointMixin", "Stream"]

airbyte_cdk/sources/streams/availability_strategy.py

@@ -5,9 +5,10 @@
 import logging
 import typing
 from abc import ABC, abstractmethod
-from typing import Optional, Tuple
+from typing import Any, Mapping, Optional, Tuple
 
-from airbyte_cdk.sources.streams import Stream
+from airbyte_cdk.models import SyncMode
+from airbyte_cdk.sources.streams.core import Stream, StreamData
 
 if typing.TYPE_CHECKING:
     from airbyte_cdk.sources import Source
@@ -19,7 +20,9 @@ class AvailabilityStrategy(ABC):
     """
 
     @abstractmethod
-    def check_availability(self, stream: Stream, logger: logging.Logger, source: Optional["Source"]) -> Tuple[bool, Optional[str]]:
+    def check_availability(
+        self, stream: Stream, logger: logging.Logger, source: Optional["Source"] = None
+    ) -> Tuple[bool, Optional[str]]:
         """
         Checks stream availability.
 
@@ -31,3 +34,51 @@ class AvailabilityStrategy(ABC):
           for some reason and the str should describe what went wrong and how to
           resolve the unavailability, if possible.
         """
+
+    @staticmethod
+    def get_first_stream_slice(stream: Stream) -> Optional[Mapping[str, Any]]:
+        """
+        Gets the first stream_slice from a given stream's stream_slices.
+        :param stream: stream
+        :raises StopIteration: if there is no first slice to return (the stream_slices generator is empty)
+        :return: first stream slice from 'stream_slices' generator (`None` is a valid stream slice)
+        """
+        # We wrap the return output of stream_slices() because some implementations return types that are iterable,
+        # but not iterators such as lists or tuples
+        slices = iter(
+            stream.stream_slices(
+                cursor_field=stream.cursor_field,  # type: ignore[arg-type]
+                sync_mode=SyncMode.full_refresh,
+            )
+        )
+        return next(slices)
+
+    @staticmethod
+    def get_first_record_for_slice(
+        stream: Stream, stream_slice: Optional[Mapping[str, Any]]
+    ) -> StreamData:
+        """
+        Gets the first record for a stream_slice of a stream.
+
+        :param stream: stream instance from which to read records
+        :param stream_slice: stream_slice parameters for slicing the stream
+        :raises StopIteration: if there is no first record to return (the read_records generator is empty)
+        :return: StreamData containing the first record in the slice
+        """
+        # Store the original value of exit_on_rate_limit
+        original_exit_on_rate_limit = stream.exit_on_rate_limit
+
+        try:
+            # Ensure exit_on_rate_limit is safely set to True if possible
+            stream.exit_on_rate_limit = True
+
+            # We wrap the return output of read_records() because some implementations return types that are iterable,
+            # but not iterators such as lists or tuples
+            records_for_slice = iter(
+                stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice)
+            )
+
+            return next(records_for_slice)
+        finally:
+            # Restore the original exit_on_rate_limit value
+            stream.exit_on_rate_limit = original_exit_on_rate_limit