airbyte-cdk 6.45.0.post24.dev14387467928__py3-none-any.whl → 6.45.1.post42.dev14452300548__py3-none-any.whl

airbyte_cdk/sources/file_based/file_types/file_transfer.py

@@ -2,34 +2,27 @@
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
 #
 import logging
-import os
-from typing import Any, Dict, Iterable
+from typing import Iterable, Tuple
 
-from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
+from airbyte_cdk.models import AirbyteRecordMessageFileReference
 from airbyte_cdk.sources.file_based.file_based_stream_reader import AbstractFileBasedStreamReader
+from airbyte_cdk.sources.file_based.file_record_data import FileRecordData
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
-
-AIRBYTE_STAGING_DIRECTORY = os.getenv("AIRBYTE_STAGING_DIRECTORY", "/staging/files")
-DEFAULT_LOCAL_DIRECTORY = "/tmp/airbyte-file-transfer"
+from airbyte_cdk.sources.utils.files_directory import get_files_directory
 
 
 class FileTransfer:
     def __init__(self) -> None:
-        self._local_directory = (
-            AIRBYTE_STAGING_DIRECTORY
-            if os.path.exists(AIRBYTE_STAGING_DIRECTORY)
-            else DEFAULT_LOCAL_DIRECTORY
-        )
+        self._local_directory = get_files_directory()
 
-    def get_file(
+    def upload(
         self,
-        config: FileBasedStreamConfig,
         file: RemoteFile,
         stream_reader: AbstractFileBasedStreamReader,
         logger: logging.Logger,
-    ) -> Iterable[Dict[str, Any]]:
+    ) -> Iterable[Tuple[FileRecordData, AirbyteRecordMessageFileReference]]:
         try:
-            yield stream_reader.get_file(
+            yield stream_reader.upload(
                 file=file, local_directory=self._local_directory, logger=logger
             )
         except Exception as ex:

airbyte_cdk/sources/file_based/schema_helpers.py

@@ -18,9 +18,17 @@ JsonSchemaSupportedType = Union[List[str], Literal["string"], str]
 SchemaType = Mapping[str, Mapping[str, JsonSchemaSupportedType]]
 
 schemaless_schema = {"type": "object", "properties": {"data": {"type": "object"}}}
+
 file_transfer_schema = {
     "type": "object",
-    "properties": {"data": {"type": "object"}, "file": {"type": "object"}},
+    "properties": {
+        "folder": {"type": "string"},
+        "file_name": {"type": "string"},
+        "bytes": {"type": "integer"},
+        "id": {"type": ["null", "string"]},
+        "updated_at": {"type": ["null", "string"]},
+        "mime_type": {"type": ["null", "string"]},
+    },
 }
 
 

airbyte_cdk/sources/file_based/stream/concurrent/adapters.py

@@ -4,7 +4,7 @@
 
 import copy
 import logging
-from functools import cache, lru_cache
+from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, MutableMapping, Optional, Union
 
 from typing_extensions import deprecated
@@ -258,19 +258,14 @@ class FileBasedStreamPartition(Partition):
                     and record_data.record is not None
                 ):
                     # `AirbyteMessage`s of type `Record` should also be yielded so they are enqueued
-                    # If stream is flagged for file_transfer the record should data in file key
-                    record_message_data = (
-                        record_data.record.file
-                        if self._use_file_transfer()
-                        else record_data.record.data
-                    )
+                    record_message_data = record_data.record.data
                     if not record_message_data:
                         raise ExceptionWithDisplayMessage("A record without data was found")
                     else:
                         yield Record(
                             data=record_message_data,
                             stream_name=self.stream_name(),
-                            is_file_transfer_message=self._use_file_transfer(),
+                            file_reference=record_data.record.file_reference,
                         )
                 else:
                     self._message_repository.emit_message(record_data)
@@ -284,9 +279,9 @@ class FileBasedStreamPartition(Partition):
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         if self._slice is None:
             return None
-        assert (
-            len(self._slice["files"]) == 1
-        ), f"Expected 1 file per partition but got {len(self._slice['files'])} for stream {self.stream_name()}"
+        assert len(self._slice["files"]) == 1, (
+            f"Expected 1 file per partition but got {len(self._slice['files'])} for stream {self.stream_name()}"
+        )
         file = self._slice["files"][0]
         return {"files": [file]}
 
@@ -306,10 +301,6 @@ class FileBasedStreamPartition(Partition):
     def stream_name(self) -> str:
         return self._stream.name
 
-    @cache
-    def _use_file_transfer(self) -> bool:
-        return hasattr(self._stream, "use_file_transfer") and self._stream.use_file_transfer
-
     def __repr__(self) -> str:
         return f"FileBasedStreamPartition({self._stream.name}, {self._slice})"
 

airbyte_cdk/sources/file_based/stream/default_file_based_stream.py

@@ -11,7 +11,7 @@ from functools import cache
 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 AirbyteLogMessage, AirbyteMessage, AirbyteStream, 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 (
@@ -56,6 +56,7 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
     airbyte_columns = [ab_last_mod_col, ab_file_name_col]
     use_file_transfer = False
     preserve_directory_structure = True
+    _file_transfer = FileTransfer()
 
     def __init__(self, **kwargs: Any):
         if self.FILE_TRANSFER_KW in kwargs:
@@ -93,21 +94,6 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             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]]]:
@@ -145,14 +131,6 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
         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`.
@@ -173,19 +151,13 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
 
             try:
                 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
+                    for file_record_data, file_reference in self._file_transfer.upload(
+                        file=file, stream_reader=self.stream_reader, logger=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
+                            self.name,
+                            file_record_data.dict(exclude_none=True),
+                            file_reference=file_reference,
                         )
                 else:
                     for record in parser.parse_records(
@@ -259,6 +231,8 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
 
     @cache
     def get_json_schema(self) -> JsonSchema:
+        if self.use_file_transfer:
+            return file_transfer_schema
         extra_fields = {
             self.ab_last_mod_col: {"type": "string"},
             self.ab_file_name_col: {"type": "string"},
@@ -282,9 +256,7 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             return {"type": "object", "properties": {**extra_fields, **schema["properties"]}}
 
     def _get_raw_json_schema(self) -> JsonSchema:
-        if self.use_file_transfer:
-            return file_transfer_schema
-        elif self.config.input_schema:
+        if self.config.input_schema:
             return self.config.get_input_schema()  # type: ignore
         elif self.config.schemaless:
             return schemaless_schema
@@ -341,6 +313,11 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             self.config.globs or [], self.config.legacy_prefix, self.logger
         )
 
+    def as_airbyte_stream(self) -> AirbyteStream:
+        file_stream = super().as_airbyte_stream()
+        file_stream.is_file_based = self.use_file_transfer
+        return file_stream
+
     def infer_schema(self, files: List[RemoteFile]) -> Mapping[str, Any]:
         loop = asyncio.get_event_loop()
         schema = loop.run_until_complete(self._infer_schema(files))

airbyte_cdk/sources/file_based/stream/permissions_file_based_stream.py

@@ -61,9 +61,7 @@ class PermissionsFileBasedStream(DefaultFileBasedStream):
                 permissions_record = self.transform_record(
                     permissions_record, file, file_datetime_string
                 )
-                yield stream_data_to_airbyte_message(
-                    self.name, permissions_record, is_file_transfer_message=False
-                )
+                yield stream_data_to_airbyte_message(self.name, permissions_record)
             except Exception as e:
                 self.logger.error(f"Failed to retrieve permissions for file {file.uri}: {str(e)}")
                 yield AirbyteMessage(

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -29,6 +29,7 @@ class DefaultStream(AbstractStream):
         logger: Logger,
         cursor: Cursor,
         namespace: Optional[str] = None,
+        supports_file_transfer: bool = False,
     ) -> None:
         self._stream_partition_generator = partition_generator
         self._name = name
@@ -39,6 +40,7 @@ class DefaultStream(AbstractStream):
         self._logger = logger
         self._cursor = cursor
         self._namespace = namespace
+        self._supports_file_transfer = supports_file_transfer
 
     def generate_partitions(self) -> Iterable[Partition]:
         yield from self._stream_partition_generator.generate()
@@ -68,6 +70,7 @@ class DefaultStream(AbstractStream):
             json_schema=dict(self._json_schema),
             supported_sync_modes=[SyncMode.full_refresh],
             is_resumable=False,
+            is_file_based=self._supports_file_transfer,
         )
 
         if self._namespace:

airbyte_cdk/sources/types.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 from typing import Any, ItemsView, Iterator, KeysView, List, Mapping, Optional, ValuesView
 
+from airbyte_cdk.models import AirbyteRecordMessageFileReference
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
 # A FieldPointer designates a path to a field inside a mapping. For example, retrieving ["k1", "k1.2"] in the object {"k1" :{"k1.2":
@@ -23,12 +24,12 @@ class Record(Mapping[str, Any]):
         data: Mapping[str, Any],
         stream_name: str,
         associated_slice: Optional[StreamSlice] = None,
-        is_file_transfer_message: bool = False,
+        file_reference: Optional[AirbyteRecordMessageFileReference] = None,
     ):
         self._data = data
         self._associated_slice = associated_slice
         self.stream_name = stream_name
-        self.is_file_transfer_message = is_file_transfer_message
+        self._file_reference = file_reference
 
     @property
     def data(self) -> Mapping[str, Any]:
@@ -38,6 +39,14 @@ class Record(Mapping[str, Any]):
     def associated_slice(self) -> Optional[StreamSlice]:
         return self._associated_slice
 
+    @property
+    def file_reference(self) -> AirbyteRecordMessageFileReference:
+        return self._file_reference
+
+    @file_reference.setter
+    def file_reference(self, value: AirbyteRecordMessageFileReference) -> None:
+        self._file_reference = value
+
     def __repr__(self) -> str:
         return repr(self._data)
 

airbyte_cdk/sources/utils/files_directory.py

@@ -0,0 +1,15 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+import os
+
+AIRBYTE_STAGING_DIRECTORY = os.getenv("AIRBYTE_STAGING_DIRECTORY", "/staging/files")
+DEFAULT_LOCAL_DIRECTORY = "/tmp/airbyte-file-transfer"
+
+
+def get_files_directory() -> str:
+    return (
+        AIRBYTE_STAGING_DIRECTORY
+        if os.path.exists(AIRBYTE_STAGING_DIRECTORY)
+        else DEFAULT_LOCAL_DIRECTORY
+    )

airbyte_cdk/sources/utils/record_helper.py

@@ -9,10 +9,10 @@ from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteRecordMessage,
+    AirbyteRecordMessageFileReference,
     AirbyteTraceMessage,
 )
 from airbyte_cdk.models import Type as MessageType
-from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 
@@ -22,7 +22,7 @@ def stream_data_to_airbyte_message(
     data_or_message: StreamData,
     transformer: TypeTransformer = TypeTransformer(TransformConfig.NoTransform),
     schema: Optional[Mapping[str, Any]] = None,
-    is_file_transfer_message: bool = False,
+    file_reference: Optional[AirbyteRecordMessageFileReference] = None,
 ) -> AirbyteMessage:
     if schema is None:
         schema = {}
@@ -36,12 +36,12 @@ def stream_data_to_airbyte_message(
             # taken unless configured. See
             # docs/connector-development/cdk-python/schemas.md for details.
             transformer.transform(data, schema)
-            if is_file_transfer_message:
-                message = AirbyteFileTransferRecordMessage(
-                    stream=stream_name, file=data, emitted_at=now_millis, data={}
-                )
-            else:
-                message = AirbyteRecordMessage(stream=stream_name, data=data, emitted_at=now_millis)
+            message = AirbyteRecordMessage(
+                stream=stream_name,
+                data=data,
+                emitted_at=now_millis,
+                file_reference=file_reference,
+            )
             return AirbyteMessage(type=MessageType.RECORD, record=message)
         case AirbyteTraceMessage():
             return AirbyteMessage(type=MessageType.TRACE, trace=data_or_message)

airbyte_cdk/sql/shared/sql_processor.py

@@ -326,9 +326,9 @@ class SqlProcessorBase(abc.ABC):
 
         if DEBUG_MODE:
             found_schemas = schemas_list
-            assert (
-                schema_name in found_schemas
-            ), f"Schema {schema_name} was not created. Found: {found_schemas}"
+            assert schema_name in found_schemas, (
+                f"Schema {schema_name} was not created. Found: {found_schemas}"
+            )
 
     def _quote_identifier(self, identifier: str) -> str:
         """Return the given identifier, quoted."""
@@ -617,10 +617,10 @@ class SqlProcessorBase(abc.ABC):
         self._execute_sql(
             f"""
             INSERT INTO {self._fully_qualified(final_table_name)} (
-            {f',{nl}  '.join(columns)}
+            {f",{nl}  ".join(columns)}
             )
             SELECT
-            {f',{nl}  '.join(columns)}
+            {f",{nl}  ".join(columns)}
             FROM {self._fully_qualified(temp_table_name)}
             """,
         )
@@ -645,8 +645,7 @@ class SqlProcessorBase(abc.ABC):
         deletion_name = f"{final_table_name}_deleteme"
         commands = "\n".join(
             [
-                f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME "
-                f"TO {deletion_name};",
+                f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME TO {deletion_name};",
                 f"ALTER TABLE {self._fully_qualified(temp_table_name)} RENAME "
                 f"TO {final_table_name};",
                 f"DROP TABLE {self._fully_qualified(deletion_name)};",
@@ -686,10 +685,10 @@ class SqlProcessorBase(abc.ABC):
                 {set_clause}
             WHEN NOT MATCHED THEN INSERT
             (
-                {f',{nl}    '.join(columns)}
+                {f",{nl}    ".join(columns)}
             )
             VALUES (
-                tmp.{f',{nl}    tmp.'.join(columns)}
+                tmp.{f",{nl}    tmp.".join(columns)}
             );
             """,
         )

airbyte_cdk/test/entrypoint_wrapper.py

@@ -82,10 +82,6 @@ class EntrypointOutput:
     def state_messages(self) -> List[AirbyteMessage]:
         return self._get_message_by_types([Type.STATE])
 
-    @property
-    def connection_status_messages(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.CONNECTION_STATUS])
-
     @property
     def most_recent_state(self) -> Any:
         state_messages = self._get_message_by_types([Type.STATE])

airbyte_cdk/test/mock_http/response_builder.py

@@ -198,6 +198,14 @@ def find_template(resource: str, execution_folder: str) -> Dict[str, Any]:
         return json.load(template_file)  # type: ignore  # we assume the dev correctly set up the resource file
 
 
+def find_binary_response(resource: str, execution_folder: str) -> bytes:
+    response_filepath = str(
+        get_unit_test_folder(execution_folder) / "resource" / "http" / "response" / f"{resource}"
+    )
+    with open(response_filepath, "rb") as response_file:
+        return response_file.read()  # type: ignore  # we assume the dev correctly set up the resource file
+
+
 def create_record_builder(
     response_template: Dict[str, Any],
     records_path: Union[FieldPath, NestedPath],

{airbyte_cdk-6.45.0.post24.dev14387467928.dist-info → airbyte_cdk-6.45.1.post42.dev14452300548.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 6.45.0.post24.dev14387467928
+Version: 6.45.1.post42.dev14452300548
 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.15,<0.16)
 Requires-Dist: anyascii (>=0.3.2,<0.4.0)
 Requires-Dist: avro (>=1.11.2,<1.13.0) ; extra == "file-based"
 Requires-Dist: backoff