file-mode-api: initial changes to handle include_files selection from configured catalog.

Author: aldogonzalez8

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -139,7 +139,7 @@ def read(
         catalog: ConfiguredAirbyteCatalog,
         state: Optional[List[AirbyteStateMessage]] = None,
     ) -> Iterator[AirbyteMessage]:
-        concurrent_streams, _ = self._group_streams(config=config)
+        concurrent_streams, _ = self._group_streams(config=config, catalog=catalog)
 
         # ConcurrentReadProcessor pops streams that are finished being read so before syncing, the names of
         # the concurrent streams must be saved so that they can be removed from the catalog before starting
@@ -180,7 +180,7 @@ def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> Airbyte
             ]
         )
 
-    def streams(self, config: Mapping[str, Any]) -> List[Stream]:
+    def streams(self, config: Mapping[str, Any], catalog: ConfiguredAirbyteCatalog | None = None) -> List[Stream]:
         """
         The `streams` method is used as part of the AbstractSource in the following cases:
         * ConcurrentDeclarativeSource.check -> ManifestDeclarativeSource.check -> AbstractSource.check -> DeclarativeSource.check_connection -> CheckStream.check_connection -> streams
@@ -189,10 +189,10 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
 
         In both case, we will assume that calling the DeclarativeStream is perfectly fine as the result for these is the same regardless of if it is a DeclarativeStream or a DefaultStream (concurrent). This should simply be removed once we have moved away from the mentioned code paths above.
         """
-        return super().streams(config)
+        return super().streams(config, catalog=catalog)
 
     def _group_streams(
-        self, config: Mapping[str, Any]
+        self, config: Mapping[str, Any], catalog: ConfiguredAirbyteCatalog | None = None
     ) -> Tuple[List[AbstractStream], List[Stream]]:
         concurrent_streams: List[AbstractStream] = []
         synchronous_streams: List[Stream] = []
@@ -205,7 +205,7 @@ def _group_streams(
 
         name_to_stream_mapping = {stream["name"]: stream for stream in streams}
 
-        for declarative_stream in self.streams(config=config):
+        for declarative_stream in self.streams(config=config, catalog=catalog):
             # Some low-code sources use a combination of DeclarativeStream and regular Python streams. We can't inspect
             # these legacy Python streams the way we do low-code streams to determine if they are concurrent compatible,
             # so we need to treat them as synchronous

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -20,6 +20,7 @@
     AirbyteMessage,
     AirbyteStateMessage,
     ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteStream,
     ConnectorSpecification,
     FailureType,
 )
@@ -141,7 +142,7 @@ def connection_checker(self) -> ConnectionChecker:
                 f"Expected to generate a ConnectionChecker component, but received {check_stream.__class__}"
             )
 
-    def streams(self, config: Mapping[str, Any]) -> List[Stream]:
+    def streams(self, config: Mapping[str, Any], catalog: Optional[ConfiguredAirbyteCatalog] = None) -> List[Stream]:
         self._emit_manifest_debug_message(
             extra_args={"source_name": self.name, "parsed_config": json.dumps(self._source_config)}
         )
@@ -154,6 +155,7 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         if api_budget_model:
             self._constructor.set_api_budget(api_budget_model, config)
 
+        catalog_with_streams_name = self._catalog_with_streams_name(catalog)
         source_streams = [
             self._constructor.create_component(
                 StateDelegatingStreamModel
@@ -162,12 +164,35 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
                 stream_config,
                 config,
                 emit_connector_builder_messages=self._emit_connector_builder_messages,
+                include_files=self._get_include_files(stream_config=stream_config, catalog_with_streams_name=catalog_with_streams_name),
             )
             for stream_config in self._initialize_cache_for_parent_streams(deepcopy(stream_configs))
         ]
 
         return source_streams
 
+    @staticmethod
+    def _get_include_files(
+        stream_config: Dict[str, Any], catalog_with_streams_name: Mapping[str, ConfiguredAirbyteStream] | None
+    ) -> Optional[bool]:
+        """
+        Returns the include_files for the stream if it exists in the catalog.
+        """
+        if catalog_with_streams_name:
+            stream_name = stream_config.get("name")
+            configured_catalog_stream = catalog_with_streams_name.get(stream_name)
+            return configured_catalog_stream and configured_catalog_stream.include_files
+        return False
+
+    @staticmethod
+    def _catalog_with_streams_name(catalog: ConfiguredAirbyteCatalog | None) -> Mapping[str, ConfiguredAirbyteStream] | None:
+        """
+        Returns a dict mapping stream names to their corresponding ConfiguredAirbyteStream objects.
+        """
+        if catalog:
+            return { configured_stream.stream.name: configured_stream for configured_stream in catalog.streams}
+        return None
+
     @staticmethod
     def _initialize_cache_for_parent_streams(
         stream_configs: List[Dict[str, Any]],

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -1852,8 +1852,9 @@ def create_declarative_stream(
                 )
         file_uploader = None
         if model.file_uploader:
+            include_files = kwargs.pop("include_files", False)
             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(
@@ -3597,7 +3598,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(
@@ -3618,7 +3619,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,

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,9 @@ 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 +140,21 @@ 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 +167,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 +190,52 @@ 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 +246,63 @@ 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_without_include_files(self) -> None:
+        """Test that article attachments can be read without including files, 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
+            assert (
+                file_reference.staging_file_url
+                == "/tmp/airbyte-file-transfer/article_attachments/12138758717583/some_image_name.png"
+            )
+
+            assert file_reference.source_file_relative_path
+            assert not re.match(
+                r"^article_attachments/[0-9a-fA-F-]{36}$", file_reference.source_file_relative_path
+            )
+            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 +320,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 +337,16 @@ 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