feat(cdk): handle include_files from configured catalog

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -88,6 +88,7 @@ def __init__(
             emit_connector_builder_messages=emit_connector_builder_messages,
             disable_resumable_full_refresh=True,
             connector_state_manager=self._connector_state_manager,
+            catalog=catalog,
         )
 
         super().__init__(

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -99,6 +99,7 @@ def __init__(
         component_factory: Optional[ModelToComponentFactory] = None,
         migrate_manifest: Optional[bool] = False,
         normalize_manifest: Optional[bool] = False,
+        catalog: Optional[ConfiguredAirbyteCatalog] = None,
     ) -> None:
         """
         Args:
@@ -124,6 +125,7 @@ def __init__(
             else ModelToComponentFactory(
                 emit_connector_builder_messages,
                 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"),
+                catalog=catalog,
             )
         )
         self._message_repository = self._constructor.get_message_repository()

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -30,7 +30,7 @@
 from airbyte_cdk.connector_builder.models import (
     LogMessage as ConnectorBuilderLogMessage,
 )
-from airbyte_cdk.models import FailureType, Level
+from airbyte_cdk.models import ConfiguredAirbyteCatalog, ConfiguredAirbyteStream, FailureType, Level
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.declarative.async_job.job_orchestrator import AsyncJobOrchestrator
 from airbyte_cdk.sources.declarative.async_job.job_tracker import JobTracker
@@ -583,6 +583,7 @@ def __init__(
         disable_retries: bool = False,
         disable_cache: bool = False,
         disable_resumable_full_refresh: bool = False,
+        catalog: ConfiguredAirbyteCatalog = None,
         message_repository: Optional[MessageRepository] = None,
         connector_state_manager: Optional[ConnectorStateManager] = None,
         max_concurrent_async_job_count: Optional[int] = None,
@@ -602,6 +603,9 @@ def __init__(
         self._job_tracker: JobTracker = JobTracker(max_concurrent_async_job_count or 1)
         # placeholder for deprecation warnings
         self._collected_deprecation_logs: List[ConnectorBuilderLogMessage] = []
+        self._catalog_with_streams_name = (
+            self._get_catalog_with_streams_name(catalog) if catalog else None
+        )
 
     def _init_mappings(self) -> None:
         self.PYDANTIC_MODEL_TO_CONSTRUCTOR: Mapping[Type[BaseModel], Callable[..., Any]] = {
@@ -1887,8 +1891,9 @@ def create_declarative_stream(
                 )
         file_uploader = None
         if model.file_uploader:
+            include_files = self._get_include_files(model)
             file_uploader = self._create_component_from_model(
-                model=model.file_uploader, config=config
+                model=model.file_uploader, config=config, include_files=include_files
             )
 
         retriever = self._create_component_from_model(
@@ -3679,7 +3684,7 @@ def create_fixed_window_call_rate_policy(
         )
 
     def create_file_uploader(
-        self, model: FileUploaderModel, config: Config, **kwargs: Any
+        self, model: FileUploaderModel, config: Config, include_files: bool, **kwargs: Any
     ) -> FileUploader:
         name = "File Uploader"
         requester = self._create_component_from_model(
@@ -3700,7 +3705,7 @@ def create_file_uploader(
             download_target_extractor=download_target_extractor,
             config=config,
             file_writer=NoopFileWriter()
-            if emit_connector_builder_messages
+            if emit_connector_builder_messages or not include_files
             else LocalFileSystemFileWriter(),
             parameters=model.parameters or {},
             filename_extractor=model.filename_extractor if model.filename_extractor else None,
@@ -3792,3 +3797,27 @@ def create_grouping_partition_router(
             deduplicate=model.deduplicate if model.deduplicate is not None else True,
             config=config,
         )
+
+    @staticmethod
+    def _get_catalog_with_streams_name(
+        catalog: ConfiguredAirbyteCatalog,
+    ) -> Mapping[str, ConfiguredAirbyteStream]:
+        """
+        Returns a dict mapping stream names to their corresponding ConfiguredAirbyteStream objects.
+        """
+        return {
+            configured_stream.stream.name: configured_stream
+            for configured_stream in catalog.streams
+        }
+
+    def _get_include_files(
+        self,
+        stream_model: DeclarativeStreamModel,
+    ) -> bool:
+        """
+        Returns the include_files for the stream if it exists in the catalog.
+        """
+        if stream_model.name and self._catalog_with_streams_name:
+            configured_catalog_stream = self._catalog_with_streams_name.get(stream_model.name)
+            return bool(configured_catalog_stream and configured_catalog_stream.include_files)
+        return False

airbyte_cdk/test/catalog_builder.py

@@ -41,6 +41,14 @@ def with_json_schema(self, json_schema: Dict[str, Any]) -> "ConfiguredAirbyteStr
         self._stream["stream"]["json_schema"] = json_schema
         return self
 
+    def with_include_files(self, include_files: bool) -> "ConfiguredAirbyteStreamBuilder":
+        """
+        Set whether the stream should include files in the sync.
+        :param include_files: True if files should be included, False otherwise.
+        """
+        self._stream["include_files"] = include_files
+        return self
+
     def build(self) -> ConfiguredAirbyteStream:
         return ConfiguredAirbyteStreamSerializer.load(self._stream)
 

unit_tests/sources/declarative/file/test_file_stream.py

@@ -120,7 +120,15 @@ def test_get_articles(self) -> None:
             assert output.records
 
     def test_get_article_attachments(self) -> None:
-        with HttpMocker() as http_mocker:
+        with (
+            HttpMocker() as http_mocker,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.noop_file_writer.NoopFileWriter.write"
+            ) as mock_noop_write,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.local_file_system_file_writer.LocalFileSystemFileWriter.write"
+            ) as mock_file_system_write,
+        ):
             http_mocker.get(
                 HttpRequest(url=STREAM_URL),
                 HttpResponse(json.dumps(find_template("file_api/articles", __file__)), 200),
@@ -138,14 +146,25 @@ def test_get_article_attachments(self) -> None:
                 ),
             )
 
+            file_size = 12345
+            mock_file_system_write.return_value = file_size  # Simulate a file size
+
             output = read(
                 self._config(),
                 CatalogBuilder()
-                .with_stream(ConfiguredAirbyteStreamBuilder().with_name("article_attachments"))
+                .with_stream(
+                    ConfiguredAirbyteStreamBuilder()
+                    .with_name("article_attachments")
+                    .with_include_files(True)
+                )
                 .build(),
             )
 
             assert output.records
+            # Ensure that FileSystemFileWriter is called.
+            mock_file_system_write.assert_called()
+            # Ensure that NoopFileWriter is not called.
+            mock_noop_write.assert_not_called()
             file_reference = output.records[0].record.file_reference
             assert file_reference
             assert file_reference.staging_file_url
@@ -158,7 +177,8 @@ def test_get_article_attachments(self) -> None:
             )
             assert file_reference.file_size_bytes
 
-    def test_get_article_attachments_with_filename_extractor(self) -> None:
+    def test_get_article_attachments_and_file_is_uploaded(self) -> None:
+        """Test that article attachments can be read and the file is uploaded to the staging directory"""
         with HttpMocker() as http_mocker:
             http_mocker.get(
                 HttpRequest(url=STREAM_URL),
@@ -180,12 +200,68 @@ def test_get_article_attachments_with_filename_extractor(self) -> None:
             output = read(
                 self._config(),
                 CatalogBuilder()
-                .with_stream(ConfiguredAirbyteStreamBuilder().with_name("article_attachments"))
+                .with_stream(
+                    ConfiguredAirbyteStreamBuilder()
+                    .with_name("article_attachments")
+                    .with_include_files(True)
+                )
+                .build(),
+                yaml_file="test_file_stream_with_filename_extractor.yaml",
+            )
+            file_reference = output.records[0].record.file_reference
+            assert file_reference.file_size_bytes
+            assert Path(file_reference.staging_file_url).exists(), (
+                "File should be uploaded to the staging directory"
+            )
+
+    def test_get_article_attachments_with_filename_extractor(self) -> None:
+        """Test that article attachments can be read with filename extractor and file system writer is called"""
+        with (
+            HttpMocker() as http_mocker,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.noop_file_writer.NoopFileWriter.write"
+            ) as mock_noop_write,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.local_file_system_file_writer.LocalFileSystemFileWriter.write"
+            ) as mock_file_system_write,
+        ):
+            http_mocker.get(
+                HttpRequest(url=STREAM_URL),
+                HttpResponse(json.dumps(find_template("file_api/articles", __file__)), 200),
+            )
+            http_mocker.get(
+                HttpRequest(url=STREAM_ATTACHMENTS_URL),
+                HttpResponse(
+                    json.dumps(find_template("file_api/article_attachments", __file__)), 200
+                ),
+            )
+            http_mocker.get(
+                HttpRequest(url=STREAM_ATTACHMENT_CONTENT_URL),
+                HttpResponse(
+                    find_binary_response("file_api/article_attachment_content.png", __file__), 200
+                ),
+            )
+
+            file_size = 12345
+            mock_file_system_write.return_value = file_size  # Simulate a file size
+
+            output = read(
+                self._config(),
+                CatalogBuilder()
+                .with_stream(
+                    ConfiguredAirbyteStreamBuilder()
+                    .with_name("article_attachments")
+                    .with_include_files(True)
+                )
                 .build(),
                 yaml_file="test_file_stream_with_filename_extractor.yaml",
             )
 
             assert len(output.records) == 1
+            # Ensure that FileSystemFileWriter is called.
+            mock_file_system_write.assert_called()
+            # Ensure that NoopFileWriter is not called.
+            mock_noop_write.assert_not_called()
             file_reference = output.records[0].record.file_reference
             assert file_reference
             assert (
@@ -196,10 +272,115 @@ def test_get_article_attachments_with_filename_extractor(self) -> None:
             assert not re.match(
                 r"^article_attachments/[0-9a-fA-F-]{36}$", file_reference.source_file_relative_path
             )
-            assert file_reference.file_size_bytes
+            assert file_reference.file_size_bytes == file_size
+
+    def test_get_article_attachments_with_include_files_false(self) -> None:
+        """Test that article attachments can be read with including files False, it can be opt-out by configured catalog"""
+        include_files = False
+        with (
+            HttpMocker() as http_mocker,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.noop_file_writer.NoopFileWriter.write"
+            ) as mock_noop_write,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.local_file_system_file_writer.LocalFileSystemFileWriter.write"
+            ) as mock_file_system_write,
+        ):
+            http_mocker.get(
+                HttpRequest(url=STREAM_URL),
+                HttpResponse(json.dumps(find_template("file_api/articles", __file__)), 200),
+            )
+            http_mocker.get(
+                HttpRequest(url=STREAM_ATTACHMENTS_URL),
+                HttpResponse(
+                    json.dumps(find_template("file_api/article_attachments", __file__)), 200
+                ),
+            )
+            http_mocker.get(
+                HttpRequest(url=STREAM_ATTACHMENT_CONTENT_URL),
+                HttpResponse(
+                    find_binary_response("file_api/article_attachment_content.png", __file__), 200
+                ),
+            )
+
+            mock_noop_write.return_value = NoopFileWriter.NOOP_FILE_SIZE
+
+            output = read(
+                self._config(),
+                CatalogBuilder()
+                .with_stream(
+                    ConfiguredAirbyteStreamBuilder()
+                    .with_name("article_attachments")
+                    .with_include_files(include_files)
+                )
+                .build(),
+                yaml_file="test_file_stream_with_filename_extractor.yaml",
+            )
+
+            assert len(output.records) == 1
+            # Ensure that LocalFileSystemFileWriter is not called when include_files is False
+            mock_file_system_write.assert_not_called()
+            # Ensure that NoopFileWriter is called to simulate file writing
+            mock_noop_write.assert_called()
+            file_reference = output.records[0].record.file_reference
+            assert file_reference.file_size_bytes == NoopFileWriter.NOOP_FILE_SIZE
+
+    def test_get_article_attachments_without_include_files(self) -> None:
+        """Test that article attachments can be read without including files, it can be opt-out by configured catalog"""
+        with (
+            HttpMocker() as http_mocker,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.noop_file_writer.NoopFileWriter.write"
+            ) as mock_noop_write,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.local_file_system_file_writer.LocalFileSystemFileWriter.write"
+            ) as mock_file_system_write,
+        ):
+            http_mocker.get(
+                HttpRequest(url=STREAM_URL),
+                HttpResponse(json.dumps(find_template("file_api/articles", __file__)), 200),
+            )
+            http_mocker.get(
+                HttpRequest(url=STREAM_ATTACHMENTS_URL),
+                HttpResponse(
+                    json.dumps(find_template("file_api/article_attachments", __file__)), 200
+                ),
+            )
+            http_mocker.get(
+                HttpRequest(url=STREAM_ATTACHMENT_CONTENT_URL),
+                HttpResponse(
+                    find_binary_response("file_api/article_attachment_content.png", __file__), 200
+                ),
+            )
+
+            mock_noop_write.return_value = NoopFileWriter.NOOP_FILE_SIZE
+
+            output = read(
+                self._config(),
+                CatalogBuilder()
+                .with_stream(ConfiguredAirbyteStreamBuilder().with_name("article_attachments"))
+                .build(),
+                yaml_file="test_file_stream_with_filename_extractor.yaml",
+            )
+
+            assert len(output.records) == 1
+            # Ensure that LocalFileSystemFileWriter is not called when include_files is False
+            mock_file_system_write.assert_not_called()
+            # Ensure that NoopFileWriter is called to simulate file writing
+            mock_noop_write.assert_called()
+            file_reference = output.records[0].record.file_reference
+            assert file_reference.file_size_bytes == NoopFileWriter.NOOP_FILE_SIZE
 
     def test_get_article_attachments_messages_for_connector_builder(self) -> None:
-        with HttpMocker() as http_mocker:
+        with (
+            HttpMocker() as http_mocker,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.noop_file_writer.NoopFileWriter.write"
+            ) as mock_noop_write,
+            patch(
+                "airbyte_cdk.sources.declarative.retrievers.file_uploader.local_file_system_file_writer.LocalFileSystemFileWriter.write"
+            ) as mock_file_system_write,
+        ):
             http_mocker.get(
                 HttpRequest(url=STREAM_URL),
                 HttpResponse(json.dumps(find_template("file_api/articles", __file__)), 200),
@@ -217,6 +398,9 @@ def test_get_article_attachments_messages_for_connector_builder(self) -> None:
                 ),
             )
 
+            file_size = NoopFileWriter.NOOP_FILE_SIZE
+            mock_noop_write.return_value = file_size  # Simulate a file size
+
             # Define a mock factory that forces emit_connector_builder_messages=True
             class MockModelToComponentFactory(OriginalModelToComponentFactory):
                 def __init__(self, *args, **kwargs):
@@ -231,12 +415,20 @@ def __init__(self, *args, **kwargs):
                 output = read(
                     self._config(),
                     CatalogBuilder()
-                    .with_stream(ConfiguredAirbyteStreamBuilder().with_name("article_attachments"))
+                    .with_stream(
+                        ConfiguredAirbyteStreamBuilder()
+                        .with_name("article_attachments")
+                        .with_include_files(True)
+                    )
                     .build(),
                     yaml_file="test_file_stream_with_filename_extractor.yaml",
                 )
 
                 assert len(output.records) == 1
+                # Ensure that NoopFileWriter is called.
+                mock_noop_write.assert_called()
+                # Ensure that LocalFileSystemFileWriter is not called.
+                mock_file_system_write.assert_not_called()
                 file_reference = output.records[0].record.file_reference
                 assert file_reference
                 assert file_reference.staging_file_url