fix: Restore AirbyteRecordMessageFileReference import for file-based streams

- Add missing import in record_helper.py, file_transfer.py, and test files
- Restore proper type annotations for file reference parameters
- Uncomment file reference creation logic in default_file_uploader.py
- Fix FileStreamTest.test_get_article_attachments failure
- Maintains CI detection logic for manifest validation tests

Co-Authored-By: AJ Steers <[email protected]>

Author: devin-ai-integration[bot]

airbyte_cdk/models/__init__.py

@@ -19,6 +19,7 @@
     AirbyteMessage,
     AirbyteProtocol,
     AirbyteRecordMessage,
+    AirbyteRecordMessageFileReference,
     AirbyteStateBlob,
     AirbyteStateMessage,
     AirbyteStateStats,

airbyte_cdk/sources/declarative/retrievers/file_uploader/default_file_uploader.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 from typing import Any, Mapping, Optional, Union
 
+from airbyte_cdk.models import AirbyteRecordMessageFileReference
 from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import (
     InterpolatedString,
@@ -89,8 +90,8 @@ def upload(self, record: Record) -> None:
         logger.info(f"File size: {file_size_bytes / 1024} KB")
         logger.info(f"File relative path: {str(file_relative_path)}")
 
-        # record.file_reference = AirbyteRecordMessageFileReference(
-        #     staging_file_url=str(full_path),
-        #     source_file_relative_path=str(file_relative_path),
-        #     file_size_bytes=file_size_bytes,
-        # )
+        record.file_reference = AirbyteRecordMessageFileReference(
+            staging_file_url=str(full_path),
+            source_file_relative_path=str(file_relative_path),
+            file_size_bytes=file_size_bytes,
+        )

airbyte_cdk/sources/file_based/file_types/file_transfer.py

@@ -2,8 +2,9 @@
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
 #
 import logging
-from typing import Any, Iterable, Optional, Tuple
+from typing import Iterable, Tuple
 
+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
@@ -19,7 +20,7 @@ def upload(
         file: RemoteFile,
         stream_reader: AbstractFileBasedStreamReader,
         logger: logging.Logger,
-    ) -> Iterable[Tuple[FileRecordData, Optional[Any]]]:
+    ) -> Iterable[Tuple[FileRecordData, AirbyteRecordMessageFileReference]]:
         try:
             yield stream_reader.upload(
                 file=file, local_directory=self._local_directory, logger=logger

airbyte_cdk/sources/utils/record_helper.py

@@ -9,6 +9,7 @@
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteRecordMessage,
+    AirbyteRecordMessageFileReference,
     AirbyteTraceMessage,
 )
 from airbyte_cdk.models import Type as MessageType
@@ -21,7 +22,7 @@ def stream_data_to_airbyte_message(
     data_or_message: StreamData,
     transformer: TypeTransformer = TypeTransformer(TransformConfig.NoTransform),
     schema: Optional[Mapping[str, Any]] = None,
-    file_reference: Optional[Any] = None,
+    file_reference: Optional[AirbyteRecordMessageFileReference] = None,
 ) -> AirbyteMessage:
     if schema is None:
         schema = {}

unit_tests/sources/file_based/stream/test_default_file_based_stream.py

@@ -15,6 +15,7 @@
 from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
+    AirbyteRecordMessageFileReference,
     AirbyteStream,
     Level,
 )
@@ -288,7 +289,11 @@ class DefaultFileBasedStreamFileTransferTest(unittest.TestCase):
         bytes=10,
         source_uri="file:///absolute/path/file.csv",
     )
-    _A_FILE_REFERENCE_MESSAGE = None  # AirbyteRecordMessageFileReference removed
+    _A_FILE_REFERENCE_MESSAGE = AirbyteRecordMessageFileReference(
+        file_size_bytes=10,
+        source_file_relative_path="relative/path/file.csv",
+        staging_file_url="/absolute/path/file.csv",
+    )
 
     def setUp(self) -> None:
         self._stream_config = Mock()
@@ -470,7 +475,11 @@ def test_when_compute_slices_with_duplicates(self) -> None:
 
 class DefaultFileBasedStreamSchemaTest(unittest.TestCase):
     _NOW = datetime(2022, 10, 22, tzinfo=timezone.utc)
-    _A_FILE_REFERENCE_MESSAGE = None  # AirbyteRecordMessageFileReference removed
+    _A_FILE_REFERENCE_MESSAGE = AirbyteRecordMessageFileReference(
+        file_size_bytes=10,
+        source_file_relative_path="relative/path/file.csv",
+        staging_file_url="/absolute/path/file.csv",
+    )
 
     def setUp(self) -> None:
         self._stream_config = Mock(spec=FileBasedStreamConfig)