airbyte-cdk 6.45.0__tar.gz → 6.45.0.dev4100__tar.gz

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 6.45.0
+Version: 6.45.0.dev4100
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://airbyte.com
 License: MIT
@@ -22,7 +22,7 @@ Provides-Extra: sql
 Provides-Extra: vector-db-based
 Requires-Dist: Jinja2 (>=3.1.2,<3.2.0)
 Requires-Dist: PyYAML (>=6.0.1,<7.0.0)
-Requires-Dist: airbyte-protocol-models-dataclasses (>=0.14,<0.15)
+Requires-Dist: airbyte-protocol-models-dataclasses (==0.14.4.dev1743719077)
 Requires-Dist: anyascii (>=0.3.2,<0.4.0)
 Requires-Dist: avro (>=1.11.2,<1.13.0) ; extra == "file-based"
 Requires-Dist: backoff

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -4,7 +4,7 @@
 
 
 from dataclasses import asdict, dataclass, field
-from typing import Any, Dict, List, Mapping
+from typing import Any, List, Mapping
 
 from airbyte_cdk.connector_builder.test_reader import TestReader
 from airbyte_cdk.models import (
@@ -27,34 +27,30 @@ from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE = 5
 DEFAULT_MAXIMUM_NUMBER_OF_SLICES = 5
 DEFAULT_MAXIMUM_RECORDS = 100
-DEFAULT_MAXIMUM_STREAMS = 100
 
 MAX_PAGES_PER_SLICE_KEY = "max_pages_per_slice"
 MAX_SLICES_KEY = "max_slices"
 MAX_RECORDS_KEY = "max_records"
-MAX_STREAMS_KEY = "max_streams"
 
 
 @dataclass
-class TestLimits:
+class TestReadLimits:
     max_records: int = field(default=DEFAULT_MAXIMUM_RECORDS)
     max_pages_per_slice: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
     max_slices: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
-    max_streams: int = field(default=DEFAULT_MAXIMUM_STREAMS)
 
 
-def get_limits(config: Mapping[str, Any]) -> TestLimits:
+def get_limits(config: Mapping[str, Any]) -> TestReadLimits:
     command_config = config.get("__test_read_config", {})
     max_pages_per_slice = (
         command_config.get(MAX_PAGES_PER_SLICE_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE
     )
     max_slices = command_config.get(MAX_SLICES_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_SLICES
     max_records = command_config.get(MAX_RECORDS_KEY) or DEFAULT_MAXIMUM_RECORDS
-    max_streams = command_config.get(MAX_STREAMS_KEY) or DEFAULT_MAXIMUM_STREAMS
-    return TestLimits(max_records, max_pages_per_slice, max_slices, max_streams)
+    return TestReadLimits(max_records, max_pages_per_slice, max_slices)
 
 
-def create_source(config: Mapping[str, Any], limits: TestLimits) -> ManifestDeclarativeSource:
+def create_source(config: Mapping[str, Any], limits: TestReadLimits) -> ManifestDeclarativeSource:
     manifest = config["__injected_declarative_manifest"]
     return ManifestDeclarativeSource(
         config=config,
@@ -75,7 +71,7 @@ def read_stream(
     config: Mapping[str, Any],
     configured_catalog: ConfiguredAirbyteCatalog,
     state: List[AirbyteStateMessage],
-    limits: TestLimits,
+    limits: TestReadLimits,
 ) -> AirbyteMessage:
     try:
         test_read_handler = TestReader(
@@ -121,40 +117,5 @@ def resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
         return error.as_airbyte_message()
 
 
-def full_resolve_manifest(source: ManifestDeclarativeSource, limits: TestLimits) -> AirbyteMessage:
-    try:
-        manifest = {**source.resolved_manifest}
-        streams = manifest.get("streams", [])
-        for stream in streams:
-            stream["dynamic_stream_name"] = None
-
-        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
-        for stream in source.dynamic_streams:
-            generated_streams = mapped_streams.setdefault(stream["dynamic_stream_name"], [])
-
-            if len(generated_streams) < limits.max_streams:
-                generated_streams += [stream]
-
-        for generated_streams_list in mapped_streams.values():
-            streams.extend(generated_streams_list)
-
-        manifest["streams"] = streams
-        return AirbyteMessage(
-            type=Type.RECORD,
-            record=AirbyteRecordMessage(
-                data={"manifest": manifest},
-                emitted_at=_emitted_at(),
-                stream="full_resolve_manifest",
-            ),
-        )
-    except AirbyteTracedException as exc:
-        return exc.as_airbyte_message()
-    except Exception as exc:
-        error = AirbyteTracedException.from_exception(
-            exc, message=f"Error full resolving manifest: {str(exc)}"
-        )
-        return error.as_airbyte_message()
-
-
 def _emitted_at() -> int:
     return ab_datetime_now().to_epoch_millis()

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/connector_builder/main.py

@@ -10,9 +10,8 @@ import orjson
 
 from airbyte_cdk.connector import BaseConnector
 from airbyte_cdk.connector_builder.connector_builder_handler import (
-    TestLimits,
+    TestReadLimits,
     create_source,
-    full_resolve_manifest,
     get_limits,
     read_stream,
     resolve_manifest,
@@ -73,7 +72,7 @@ def handle_connector_builder_request(
     config: Mapping[str, Any],
     catalog: Optional[ConfiguredAirbyteCatalog],
     state: List[AirbyteStateMessage],
-    limits: TestLimits,
+    limits: TestReadLimits,
 ) -> AirbyteMessage:
     if command == "resolve_manifest":
         return resolve_manifest(source)
@@ -82,8 +81,6 @@ def handle_connector_builder_request(
             catalog is not None
         ), "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
         return read_stream(source, config, catalog, state, limits)
-    elif command == "full_resolve_manifest":
-        return full_resolve_manifest(source, limits)
     else:
         raise ValueError(f"Unrecognized command {command}.")
 

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/models/__init__.py

@@ -19,6 +19,7 @@ from .airbyte_protocol import (
     AirbyteMessage,
     AirbyteProtocol,
     AirbyteRecordMessage,
+    AirbyteRecordMessageFileReference,
     AirbyteStateBlob,
     AirbyteStateMessage,
     AirbyteStateStats,

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/models/airbyte_protocol.py

@@ -8,8 +8,6 @@ from typing import Annotated, Any, Dict, List, Mapping, Optional, Union
 from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
 from serpyco_rs.metadata import Alias
 
-from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
-
 # ruff: noqa: F405  # ignore fuzzy import issues with 'import *'
 
 
@@ -84,7 +82,7 @@ class AirbyteMessage:
     spec: Optional[ConnectorSpecification] = None  # type: ignore [name-defined]
     connectionStatus: Optional[AirbyteConnectionStatus] = None  # type: ignore [name-defined]
     catalog: Optional[AirbyteCatalog] = None  # type: ignore [name-defined]
-    record: Optional[Union[AirbyteFileTransferRecordMessage, AirbyteRecordMessage]] = None  # type: ignore [name-defined]
+    record: Optional[AirbyteRecordMessage] = None  # type: ignore [name-defined]
     state: Optional[AirbyteStateMessage] = None
     trace: Optional[AirbyteTraceMessage] = None  # type: ignore [name-defined]
     control: Optional[AirbyteControlMessage] = None  # type: ignore [name-defined]

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -149,7 +149,7 @@ class ConcurrentReadProcessor:
         message = stream_data_to_airbyte_message(
             stream_name=record.stream_name,
             data_or_message=record.data,
-            is_file_transfer_message=record.is_file_transfer_message,
+            file_reference=record.file_reference,
         )
         stream = self._stream_name_to_instance[record.stream_name]
 

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/sources/declarative/async_job/job.py

@@ -34,12 +34,6 @@ class AsyncJob:
 
     def status(self) -> AsyncJobStatus:
         if self._timer.has_timed_out():
-            # TODO: we should account the fact that,
-            # certain APIs could send the `Timeout` status,
-            # thus we should not return `Timeout` in that case,
-            # but act based on the scenario.
-
-            # the default behavior is to return `Timeout` status and retry.
             return AsyncJobStatus.TIMED_OUT
         return self._status
 

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/sources/declarative/async_job/job_orchestrator.py

@@ -44,21 +44,16 @@ class AsyncPartition:
     This bucket of api_jobs is a bit useless for this iteration but should become interesting when we will be able to split jobs
     """
 
-    _DEFAULT_MAX_JOB_RETRY = 3
+    _MAX_NUMBER_OF_ATTEMPTS = 3
 
-    def __init__(
-        self, jobs: List[AsyncJob], stream_slice: StreamSlice, job_max_retry: Optional[int] = None
-    ) -> None:
+    def __init__(self, jobs: List[AsyncJob], stream_slice: StreamSlice) -> None:
         self._attempts_per_job = {job: 1 for job in jobs}
         self._stream_slice = stream_slice
-        self._job_max_retry = (
-            job_max_retry if job_max_retry is not None else self._DEFAULT_MAX_JOB_RETRY
-        )
 
     def has_reached_max_attempt(self) -> bool:
         return any(
             map(
-                lambda attempt_count: attempt_count >= self._job_max_retry,
+                lambda attempt_count: attempt_count >= self._MAX_NUMBER_OF_ATTEMPTS,
                 self._attempts_per_job.values(),
             )
         )
@@ -67,7 +62,7 @@ class AsyncPartition:
         current_attempt_count = self._attempts_per_job.pop(job_to_replace, None)
         if current_attempt_count is None:
             raise ValueError("Could not find job to replace")
-        elif current_attempt_count >= self._job_max_retry:
+        elif current_attempt_count >= self._MAX_NUMBER_OF_ATTEMPTS:
             raise ValueError(f"Max attempt reached for job in partition {self._stream_slice}")
 
         new_attempt_count = current_attempt_count + 1
@@ -160,7 +155,6 @@ class AsyncJobOrchestrator:
         message_repository: MessageRepository,
         exceptions_to_break_on: Iterable[Type[Exception]] = tuple(),
         has_bulk_parent: bool = False,
-        job_max_retry: Optional[int] = None,
     ) -> None:
         """
         If the stream slices provided as a parameters relies on a async job streams that relies on the same JobTracker, `has_bulk_parent`
@@ -181,12 +175,11 @@ class AsyncJobOrchestrator:
         self._message_repository = message_repository
         self._exceptions_to_break_on: Tuple[Type[Exception], ...] = tuple(exceptions_to_break_on)
         self._has_bulk_parent = has_bulk_parent
-        self._job_max_retry = job_max_retry
 
         self._non_breaking_exceptions: List[Exception] = []
 
     def _replace_failed_jobs(self, partition: AsyncPartition) -> None:
-        failed_status_jobs = (AsyncJobStatus.FAILED, AsyncJobStatus.TIMED_OUT)
+        failed_status_jobs = (AsyncJobStatus.FAILED,)
         jobs_to_replace = [job for job in partition.jobs if job.status() in failed_status_jobs]
         for job in jobs_to_replace:
             new_job = self._start_job(job.job_parameters(), job.api_job_id())
@@ -221,7 +214,7 @@ class AsyncJobOrchestrator:
             for _slice in self._slice_iterator:
                 at_least_one_slice_consumed_from_slice_iterator_during_current_iteration = True
                 job = self._start_job(_slice)
-                self._running_partitions.append(AsyncPartition([job], _slice, self._job_max_retry))
+                self._running_partitions.append(AsyncPartition([job], _slice))
                 if self._has_bulk_parent and self._slice_iterator.has_next():
                     break
         except ConcurrentJobLimitReached:
@@ -370,7 +363,7 @@ class AsyncJobOrchestrator:
                     self._reallocate_partition(current_running_partitions, partition)
 
             # We only remove completed / timeout jobs jobs as we want failed jobs to be re-allocated in priority
-            self._remove_completed_jobs(partition)
+            self._remove_completed_or_timed_out_jobs(partition)
 
         # update the referenced list with running partitions
         self._running_partitions = current_running_partitions
@@ -385,7 +378,11 @@ class AsyncJobOrchestrator:
     def _stop_timed_out_jobs(self, partition: AsyncPartition) -> None:
         for job in partition.jobs:
             if job.status() == AsyncJobStatus.TIMED_OUT:
-                self._abort_job(job, free_job_allocation=False)
+                self._abort_job(job, free_job_allocation=True)
+                raise AirbyteTracedException(
+                    internal_message=f"Job {job.api_job_id()} has timed out. Try increasing the `polling job timeout`.",
+                    failure_type=FailureType.config_error,
+                )
 
     def _abort_job(self, job: AsyncJob, free_job_allocation: bool = True) -> None:
         try:
@@ -395,7 +392,7 @@ class AsyncJobOrchestrator:
         except Exception as exception:
             LOGGER.warning(f"Could not free budget for job {job.api_job_id()}: {exception}")
 
-    def _remove_completed_jobs(self, partition: AsyncPartition) -> None:
+    def _remove_completed_or_timed_out_jobs(self, partition: AsyncPartition) -> None:
         """
         Remove completed or timed out jobs from the partition.
 
@@ -403,7 +400,7 @@ class AsyncJobOrchestrator:
             partition (AsyncPartition): The partition to process.
         """
         for job in partition.jobs:
-            if job.status() == AsyncJobStatus.COMPLETED:
+            if job.status() in [AsyncJobStatus.COMPLETED, AsyncJobStatus.TIMED_OUT]:
                 self._job_tracker.remove_job(job.api_job_id())
 
     def _reallocate_partition(
@@ -418,7 +415,10 @@ class AsyncJobOrchestrator:
             current_running_partitions (list): The list of currently running partitions.
             partition (AsyncPartition): The partition to reallocate.
         """
-        current_running_partitions.insert(0, partition)
+        for job in partition.jobs:
+            if job.status() != AsyncJobStatus.TIMED_OUT:
+                # allow the FAILED jobs to be re-allocated for partition
+                current_running_partitions.insert(0, partition)
 
     def _process_partitions_with_errors(self, partition: AsyncPartition) -> None:
         """

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/sources/declarative/async_job/job_tracker.py

@@ -3,11 +3,9 @@
 import logging
 import threading
 import uuid
-from dataclasses import dataclass, field
-from typing import Any, Mapping, Set, Union
+from typing import Set
 
 from airbyte_cdk.logger import lazy_log
-from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 
 LOGGER = logging.getLogger("airbyte")
 
@@ -16,29 +14,15 @@ class ConcurrentJobLimitReached(Exception):
     pass
 
 
-@dataclass
 class JobTracker:
-    limit: Union[int, str]
-    config: Mapping[str, Any] = field(default_factory=dict)
-
-    def __post_init__(self) -> None:
+    def __init__(self, limit: int):
         self._jobs: Set[str] = set()
-        self._lock = threading.Lock()
-        if isinstance(self.limit, str):
-            try:
-                self.limit = int(
-                    InterpolatedString(self.limit, parameters={}).eval(config=self.config)
-                )
-            except Exception as e:
-                LOGGER.warning(
-                    f"Error interpolating max job count: {self.limit}. Setting to 1. {e}"
-                )
-                self.limit = 1
-        if self.limit < 1:
+        if limit < 1:
             LOGGER.warning(
-                f"The `max_concurrent_async_job_count` property is less than 1: {self.limit}. Setting to 1. Please update the source manifest to set a valid value."
+                f"The `max_concurrent_async_job_count` property is less than 1: {limit}. Setting to 1. Please update the source manifest to set a valid value."
             )
-        self._limit = self.limit if self.limit >= 1 else 1
+        self._limit = 1 if limit < 1 else limit
+        self._lock = threading.Lock()
 
     def try_to_get_intent(self) -> str:
         lazy_log(

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/sources/declarative/checks/__init__.py

@@ -7,10 +7,7 @@ from typing import Mapping
 from pydantic.v1 import BaseModel
 
 from airbyte_cdk.sources.declarative.checks.check_dynamic_stream import CheckDynamicStream
-from airbyte_cdk.sources.declarative.checks.check_stream import (
-    CheckStream,
-    DynamicStreamCheckConfig,
-)
+from airbyte_cdk.sources.declarative.checks.check_stream import CheckStream
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 from airbyte_cdk.sources.declarative.models import (
     CheckDynamicStream as CheckDynamicStreamModel,
@@ -24,4 +21,4 @@ COMPONENTS_CHECKER_TYPE_MAPPING: Mapping[str, type[BaseModel]] = {
     "CheckDynamicStream": CheckDynamicStreamModel,
 }
 
-__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker", "DynamicStreamCheckConfig"]
+__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker"]

airbyte_cdk-6.45.0.dev4100/airbyte_cdk/sources/declarative/checks/check_stream.py

@@ -0,0 +1,56 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+import traceback
+from dataclasses import InitVar, dataclass
+from typing import Any, List, Mapping, Tuple
+
+from airbyte_cdk import AbstractSource
+from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
+from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy
+
+
+@dataclass
+class CheckStream(ConnectionChecker):
+    """
+    Checks the connections by checking availability of one or many streams selected by the developer
+
+    Attributes:
+        stream_name (List[str]): names of streams to check
+    """
+
+    stream_names: List[str]
+    parameters: InitVar[Mapping[str, Any]]
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        self._parameters = parameters
+
+    def check_connection(
+        self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> Tuple[bool, Any]:
+        streams = source.streams(config=config)
+        stream_name_to_stream = {s.name: s for s in streams}
+        if len(streams) == 0:
+            return False, f"No streams to connect to from source {source}"
+        for stream_name in self.stream_names:
+            if stream_name not in stream_name_to_stream.keys():
+                raise ValueError(
+                    f"{stream_name} is not part of the catalog. Expected one of {stream_name_to_stream.keys()}."
+                )
+
+            stream = stream_name_to_stream[stream_name]
+            availability_strategy = HttpAvailabilityStrategy()
+            try:
+                stream_is_available, reason = availability_strategy.check_availability(
+                    stream, logger
+                )
+                if not stream_is_available:
+                    return False, reason
+            except Exception as error:
+                logger.error(
+                    f"Encountered an error trying to connect to stream {stream_name}. Error: \n {traceback.format_exc()}"
+                )
+                return False, f"Unable to connect to stream {stream_name} - {error}"
+        return True, None

{airbyte_cdk-6.45.0 → airbyte_cdk-6.45.0.dev4100}/airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -25,6 +25,7 @@ from airbyte_cdk.sources.declarative.incremental.per_partition_with_global impor
     PerPartitionWithGlobalCursor,
 )
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
+from airbyte_cdk.sources.declarative.models import FileUploader
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConcurrencyLevel as ConcurrencyLevelModel,
 )
@@ -206,6 +207,10 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
             # these legacy Python streams the way we do low-code streams to determine if they are concurrent compatible,
             # so we need to treat them as synchronous
 
+            supports_file_transfer = (
+                "file_uploader" in name_to_stream_mapping[declarative_stream.name]
+            )
+
             if (
                 isinstance(declarative_stream, DeclarativeStream)
                 and name_to_stream_mapping[declarative_stream.name]["type"]
@@ -322,6 +327,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                             else None,
                             logger=self.logger,
                             cursor=cursor,
+                            supports_file_transfer=supports_file_transfer,
                         )
                     )
                 elif (
@@ -353,6 +359,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                             cursor_field=None,
                             logger=self.logger,
                             cursor=final_state_cursor,
+                            supports_file_transfer=supports_file_transfer,
                         )
                     )
                 elif (
@@ -406,6 +413,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                             cursor_field=perpartition_cursor.cursor_field.cursor_field_key,
                             logger=self.logger,
                             cursor=perpartition_cursor,
+                            supports_file_transfer=supports_file_transfer,
                         )
                     )
                 else: