feat(PoC): adjust file-based and file uploader component to latest protocol changes. (#457)

Co-authored-by: Maxime Carbonneau-Leclerc <[email protected]>
Co-authored-by: octavia-squidington-iii <[email protected]>
Co-authored-by: Aaron ("AJ") Steers <[email protected]>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: 5 people

airbyte_cdk/models/airbyte_protocol.py

@@ -8,8 +8,6 @@
 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/models/file_transfer_record_message.py

@@ -149,7 +149,6 @@ def on_record(self, record: Record) -> Iterable[AirbyteMessage]:
         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/sources/concurrent_source/concurrent_read_processor.py

@@ -83,7 +83,7 @@ def upload(self, record: Record) -> None:
             logger.info(f"File relative path: {str(file_relative_path)}")
 
             record.file_reference = AirbyteRecordMessageFileReference(
-                file_url=str(full_path),
-                file_relative_path=str(file_relative_path),
+                staging_file_url=str(full_path),
+                source_file_relative_path=str(file_relative_path),
                 file_size_bytes=file_size_bytes,
             )

airbyte_cdk/sources/declarative/retrievers/file_uploader.py

@@ -8,16 +8,18 @@
 from enum import Enum
 from io import IOBase
 from os import makedirs, path
-from typing import Any, Dict, Iterable, List, Optional, Set
+from typing import Any, Callable, Iterable, List, MutableMapping, Optional, Set, Tuple
 
 from wcmatch.glob import GLOBSTAR, globmatch
 
+from airbyte_cdk.models import AirbyteRecordMessageFileReference
 from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import AbstractFileBasedSpec
 from airbyte_cdk.sources.file_based.config.validate_config_transfer_modes import (
     include_identities_stream,
     preserve_directory_structure,
     use_file_transfer,
 )
+from airbyte_cdk.sources.file_based.file_record_data import FileRecordData
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
 
 
@@ -28,6 +30,10 @@ class FileReadMode(Enum):
 
 class AbstractFileBasedStreamReader(ABC):
     DATE_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
+    FILE_RELATIVE_PATH = "file_relative_path"
+    FILE_NAME = "file_name"
+    LOCAL_FILE_PATH = "local_file_path"
+    FILE_FOLDER = "file_folder"
 
     def __init__(self) -> None:
         self._config = None
@@ -148,9 +154,9 @@ def include_identities_stream(self) -> bool:
         return False
 
     @abstractmethod
-    def get_file(
+    def upload(
         self, file: RemoteFile, local_directory: str, logger: logging.Logger
-    ) -> Dict[str, Any]:
+    ) -> Tuple[FileRecordData, AirbyteRecordMessageFileReference]:
         """
         This is required for connectors that will support writing to
         files. It will handle the logic to download,get,read,acquire or
@@ -162,25 +168,41 @@ def get_file(
                logger (logging.Logger): Logger for logging information and errors.
 
            Returns:
-               dict: A dictionary containing the following:
-                   - "file_url" (str): The absolute path of the downloaded file.
-                   - "bytes" (int): The file size in bytes.
-                   - "file_relative_path" (str): The relative path of the file for local storage. Is relative to local_directory as
-                   this a mounted volume in the pod container.
-
+               AirbyteRecordMessageFileReference: A file reference object containing:
+                   - staging_file_url (str): The absolute path to the referenced file in the staging area.
+                   - file_size_bytes (int): The size of the referenced file in bytes.
+                   - source_file_relative_path (str): The relative path to the referenced file in source.
         """
         ...
 
-    def _get_file_transfer_paths(self, file: RemoteFile, local_directory: str) -> List[str]:
+    def _get_file_transfer_paths(
+        self, source_file_relative_path: str, staging_directory: str
+    ) -> MutableMapping[str, Any]:
+        """
+        This method is used to get the file transfer paths for a given source file relative path and local directory.
+        It returns a dictionary with the following keys:
+            - FILE_RELATIVE_PATH: The relative path to file in reference to the staging directory.
+            - LOCAL_FILE_PATH: The absolute path to the file.
+            - FILE_NAME: The name of the referenced file.
+            - FILE_FOLDER: The folder of the referenced file.
+        """
         preserve_directory_structure = self.preserve_directory_structure()
+
+        file_name = path.basename(source_file_relative_path)
+        file_folder = path.dirname(source_file_relative_path)
         if preserve_directory_structure:
             # Remove left slashes from source path format to make relative path for writing locally
-            file_relative_path = file.uri.lstrip("/")
+            file_relative_path = source_file_relative_path.lstrip("/")
         else:
-            file_relative_path = path.basename(file.uri)
-        local_file_path = path.join(local_directory, file_relative_path)
-
+            file_relative_path = file_name
+        local_file_path = path.join(staging_directory, file_relative_path)
         # Ensure the local directory exists
         makedirs(path.dirname(local_file_path), exist_ok=True)
-        absolute_file_path = path.abspath(local_file_path)
-        return [file_relative_path, local_file_path, absolute_file_path]
+
+        file_paths = {
+            self.FILE_RELATIVE_PATH: file_relative_path,
+            self.LOCAL_FILE_PATH: local_file_path,
+            self.FILE_NAME: file_name,
+            self.FILE_FOLDER: file_folder,
+        }
+        return file_paths

airbyte_cdk/sources/file_based/file_based_stream_reader.py

@@ -0,0 +1,23 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic.v1 import BaseModel
+
+
+class FileRecordData(BaseModel):
+    """
+    A record in a file-based stream.
+    """
+
+    folder: str
+    filename: str
+    bytes: int
+    source_uri: str
+    id: Optional[str] = None
+    created_at: Optional[str] = None
+    updated_at: Optional[str] = None
+    mime_type: Optional[str] = None

airbyte_cdk/sources/file_based/file_record_data.py

@@ -2,11 +2,11 @@
 # 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
 from airbyte_cdk.sources.utils.files_directory import get_files_directory
 
@@ -15,15 +15,14 @@ class FileTransfer:
     def __init__(self) -> None:
         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/file_types/file_transfer.py

@@ -18,9 +18,19 @@
 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"},
+        "source_uri": {"type": "string"},
+        "bytes": {"type": "integer"},
+        "id": {"type": ["null", "string"]},
+        "created_at": {"type": ["null", "string"]},
+        "updated_at": {"type": ["null", "string"]},
+        "mime_type": {"type": ["null", "string"]},
+    },
 }
 
 

airbyte_cdk/sources/file_based/schema_helpers.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 @@ def read(self) -> Iterable[Record]:
                     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)
@@ -306,10 +301,6 @@ def __hash__(self) -> int:
     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/concurrent/adapters.py

@@ -11,7 +11,7 @@
 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 @@ def primary_key(self) -> PrimaryKeyType:
             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 @@ def transform_record(
         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 @@ def read_records_from_slice(self, stream_slice: StreamSlice) -> Iterable[Airbyte
 
             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 @@ def cursor_field(self) -> Union[str, List[str]]:
 
     @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 @@ def get_json_schema(self) -> JsonSchema:
             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 @@ def get_files(self) -> Iterable[RemoteFile]:
             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))