file-mode-api: limit schema on cdk side with a model

Author: aldogonzalez8

airbyte_cdk/sources/file_based/file_based_stream_reader.py

@@ -8,7 +8,7 @@
 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, Dict, Iterable, List, Optional, Set, Tuple
 
 from wcmatch.glob import GLOBSTAR, globmatch
 
@@ -19,6 +19,7 @@
     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
 
 
@@ -151,7 +152,7 @@ def include_identities_stream(self) -> bool:
     @abstractmethod
     def upload(
         self, file: RemoteFile, local_directory: str, logger: logging.Logger
-    ) -> AirbyteRecordMessageFileReference:
+    ) -> 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

airbyte_cdk/sources/file_based/file_record_data.py

@@ -0,0 +1,24 @@
+#
+# Copyright (c) 2023 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 = None
+
+    id: Optional[str] = None
+    created_at: Optional[int] = None
+    updated_at: Optional[int] = None
+    mime_type: Optional[str] = None
+    description: Optional[str] = None

airbyte_cdk/sources/file_based/file_types/file_transfer.py

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

airbyte_cdk/sources/file_based/schema_helpers.py

@@ -19,13 +19,17 @@
 
 schemaless_schema = {"type": "object", "properties": {"data": {"type": "object"}}}
 
-# we use the File attributes for metadata records.
 file_transfer_schema = {
     "type": "object",
     "properties": {
-        "uri": {"type": "string"},
-        "last_modified": {"type": "string"},
+        "folder": {"type": "string"},
+        "file_name": {"type": "string"},
+        "bytes": {"type": "integer"},
+        "id": {"type": ["null", "string"]},
+        "created_at": {"type": ["null", "integer"]},
+        "updated_at": {"type": ["null", "integer"]},
         "mime_type": {"type": ["null", "string"]},
+        "description": {"type": ["null", "string"]},
     },
 }
 

airbyte_cdk/sources/file_based/stream/default_file_based_stream.py

@@ -97,14 +97,7 @@ 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"},
-                },
-            }
+            return file_transfer_schema
         else:
             return super()._filter_schema_invalid_properties(configured_catalog_json_schema)
 
@@ -166,12 +159,13 @@ def read_records_from_slice(self, stream_slice: StreamSlice) -> Iterable[Airbyte
 
             try:
                 if self.use_file_transfer:
-                    for file_reference in file_transfer.upload(
+                    for file_record_data, file_reference in file_transfer.upload(
                         file=file, stream_reader=self.stream_reader, logger=self.logger
                     ):
-                        record = self.transform_record({}, file, file_datetime_string)
                         yield stream_data_to_airbyte_message(
-                            self.name, record, file_reference=file_reference
+                            self.name,
+                            file_record_data.dict(exclude_none=True),
+                            file_reference=file_reference,
                         )
                 else:
                     for record in parser.parse_records(

unit_tests/sources/file_based/stream/test_default_file_based_stream.py

@@ -30,6 +30,7 @@
     FileBasedSourceError,
 )
 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.file_types import FileTransfer
 from airbyte_cdk.sources.file_based.file_types.file_type_parser import FileTypeParser
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
@@ -281,6 +282,11 @@ def test_yield_and_raise_collected(self) -> None:
 
 class DefaultFileBasedStreamFileTransferTest(unittest.TestCase):
     _NOW = datetime(2022, 10, 22, tzinfo=timezone.utc)
+    _A_FILE_RECORD_DATA = FileRecordData(
+        folder="/absolute/path/",
+        filename="file.csv",
+        bytes=10,
+    )
     _A_FILE_REFERENCE_MESSAGE = AirbyteRecordMessageFileReference(
         file_size_bytes=10,
         source_file_relative_path="relative/path/file.csv",
@@ -316,7 +322,9 @@ def setUp(self) -> None:
     def test_when_read_records_from_slice_then_return_records(self) -> None:
         """Verify that we have the new file method and data is empty"""
         with mock.patch.object(
-            FileTransfer, "upload", return_value=[self._A_FILE_REFERENCE_MESSAGE]
+            FileTransfer,
+            "upload",
+            return_value=[(self._A_FILE_RECORD_DATA, self._A_FILE_REFERENCE_MESSAGE)],
         ):
             remote_file = RemoteFile(uri="uri", last_modified=self._NOW)
             messages = list(self._stream.read_records_from_slice({"files": [remote_file]}))
@@ -326,10 +334,9 @@ def test_when_read_records_from_slice_then_return_records(self) -> None:
             ]
             assert list(map(lambda message: message.record.data, messages)) == [
                 {
-                    "_ab_source_file_last_modified": remote_file.last_modified.strftime(
-                        "%Y-%m-%dT%H:%M:%S.%fZ"
-                    ),
-                    "_ab_source_file_url": remote_file.uri,
+                    "bytes": 10,
+                    "filename": "file.csv",
+                    "folder": "/absolute/path/",
                 }
             ]