clean up the code and consolidate concurrent and legacy flow under a single path

Author: brianjlai

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -202,7 +202,7 @@ def _group_streams(
 
         # Combine streams and dynamic_streams. Note: both cannot be empty at the same time,
         # and this is validated during the initialization of the source.
-        streams = self._stream_configs(self._source_config) + self._dynamic_stream_configs(
+        streams = self._stream_configs(self._source_config, config) + self._dynamic_stream_configs(
             self._source_config, config
         )
 

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -35,6 +35,7 @@
 from airbyte_cdk.sources.declarative.checks import COMPONENTS_CHECKER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
+from airbyte_cdk.sources.declarative.interpolation import InterpolatedBoolean
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConditionalStreams as ConditionalStreamsModel,
 )
@@ -303,43 +304,27 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
             }
         )
 
-        stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams
+        stream_configs = (
+            self._stream_configs(self._source_config, config=config) + self.dynamic_streams
+        )
 
         api_budget_model = self._source_config.get("api_budget")
         if api_budget_model:
             self._constructor.set_api_budget(api_budget_model, config)
 
-        source_streams = []
-        for stream_config in self._initialize_cache_for_parent_streams(deepcopy(stream_configs)):
-            match stream_config.get("type"):
-                case StateDelegatingStreamModel.__name__:
-                    source_streams.append(
-                        self._constructor.create_component(
-                            StateDelegatingStreamModel,
-                            stream_config,
-                            config,
-                            emit_connector_builder_messages=self._emit_connector_builder_messages,
-                        )
-                    )
-                case ConditionalStreamsModel.__name__:
-                    # ConditionalStreamsModel returns a list of DeclarativeStreams so it must be extended
-                    source_streams.extend(
-                        self._constructor.create_component(
-                            ConditionalStreamsModel,
-                            stream_config,
-                            config,
-                            emit_connector_builder_messages=self._emit_connector_builder_messages,
-                        )
-                    )
-                case DeclarativeStreamModel.__name__:
-                    source_streams.append(
-                        self._constructor.create_component(
-                            DeclarativeStreamModel,
-                            stream_config,
-                            config,
-                            emit_connector_builder_messages=self._emit_connector_builder_messages,
-                        )
-                    )
+        source_streams = [
+            self._constructor.create_component(
+                (
+                    StateDelegatingStreamModel
+                    if stream_config.get("type") == StateDelegatingStreamModel.__name__
+                    else DeclarativeStreamModel
+                ),
+                stream_config,
+                config,
+                emit_connector_builder_messages=self._emit_connector_builder_messages,
+            )
+            for stream_config in self._initialize_cache_for_parent_streams(deepcopy(stream_configs))
+        ]
         return source_streams
 
     @staticmethod
@@ -383,26 +368,16 @@ def update_with_cache_parent_configs(
                             update_with_cache_parent_configs(router["parent_stream_configs"])
 
         for stream_config in stream_configs:
-            current_stream_configs = []
-            if stream_config.get("type") == ConditionalStreamsModel.__name__:
-                current_stream_configs = [
-                    conditional_stream_config
-                    for conditional_stream_config in stream_config.get("streams") or []
-                    if conditional_stream_config["name"] in parent_streams
-                ]
-            elif stream_config["name"] in parent_streams:
-                current_stream_configs = [stream_config]
-
-            for current_stream_config in current_stream_configs:
-                if current_stream_config["type"] == "StateDelegatingStream":
-                    current_stream_config["full_refresh_stream"]["retriever"]["requester"][
-                        "use_cache"
-                    ] = True
-                    current_stream_config["incremental_stream"]["retriever"]["requester"][
-                        "use_cache"
-                    ] = True
+            if stream_config["name"] in parent_streams:
+                if stream_config["type"] == "StateDelegatingStream":
+                    stream_config["full_refresh_stream"]["retriever"]["requester"]["use_cache"] = (
+                        True
+                    )
+                    stream_config["incremental_stream"]["retriever"]["requester"]["use_cache"] = (
+                        True
+                    )
                 else:
-                    current_stream_config["retriever"]["requester"]["use_cache"] = True
+                    stream_config["retriever"]["requester"]["use_cache"] = True
         return stream_configs
 
     def spec(self, logger: logging.Logger) -> ConnectorSpecification:
@@ -506,12 +481,27 @@ def _parse_version(
             # No exception
             return parsed_version
 
-    def _stream_configs(self, manifest: Mapping[str, Any]) -> List[Dict[str, Any]]:
+    def _stream_configs(
+        self, manifest: Mapping[str, Any], config: Mapping[str, Any]
+    ) -> List[Dict[str, Any]]:
         # This has a warning flag for static, but after we finish part 4 we'll replace manifest with self._source_config
-        stream_configs: List[Dict[str, Any]] = manifest.get("streams", [])
-        for s in stream_configs:
-            if "type" not in s:
-                s["type"] = "DeclarativeStream"
+        stream_configs = []
+        for current_stream_config in manifest.get("streams", []):
+            if (
+                "type" in current_stream_config
+                and current_stream_config["type"] == "ConditionalStreams"
+            ):
+                interpolated_boolean = InterpolatedBoolean(
+                    condition=current_stream_config.get("condition"),
+                    parameters={},
+                )
+
+                if interpolated_boolean.eval(config=config):
+                    stream_configs.extend(current_stream_config.get("streams", []))
+            else:
+                if "type" not in current_stream_config:
+                    current_stream_config["type"] = "DeclarativeStream"
+                stream_configs.append(current_stream_config)
         return stream_configs
 
     def _dynamic_stream_configs(

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -659,7 +659,6 @@ def _init_mappings(self) -> None:
             CheckDynamicStreamModel: self.create_check_dynamic_stream,
             CompositeErrorHandlerModel: self.create_composite_error_handler,
             ConcurrencyLevelModel: self.create_concurrency_level,
-            ConditionalStreamsModel: self.create_conditional_streams,
             ConfigMigrationModel: self.create_config_migration,
             ConfigAddFieldsModel: self.create_config_add_fields,
             ConfigRemapFieldModel: self.create_config_remap_field,
@@ -1623,22 +1622,6 @@ def create_concurrent_cursor_from_perpartition_cursor(
             use_global_cursor=use_global_cursor,
         )
 
-    def create_conditional_streams(
-        self, model: ConditionalStreamsModel, config: Config, **kwargs: Any
-    ) -> List[DeclarativeStream]:
-        condition = InterpolatedBoolean(
-            condition=model.condition, parameters=model.parameters or {}
-        )
-        should_include_streams = condition.eval(config=config)
-        return (
-            [
-                self._create_component_from_model(stream, config=config, **kwargs)
-                for stream in model.streams
-            ]
-            if should_include_streams
-            else []
-        )
-
     @staticmethod
     def create_constant_backoff_strategy(
         model: ConstantBackoffStrategyModel, config: Config, **kwargs: Any

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -4802,109 +4802,3 @@ def test_create_stream_with_multiple_schema_loaders():
     assert len(schema_loader.schema_loaders) == 2
     assert isinstance(schema_loader.schema_loaders[0], InlineSchemaLoader)
     assert isinstance(schema_loader.schema_loaders[1], InlineSchemaLoader)
-
-
-def test_conditional_streams():
-    content = """
-    lists_stream:
-      type: DeclarativeStream
-      name: lists
-      primary_key: id
-      retriever:
-        paginator:
-        type: "DefaultPaginator"
-        page_size_option:
-          type: RequestOption
-          inject_into: request_parameter
-          field_name: page_size
-        page_token_option:
-          type: RequestPath
-        pagination_strategy:
-          type: "CursorPagination"
-          cursor_value: "{{ response._metadata.next }}"
-          page_size: 10
-        requester:
-          url_base: "https://testing.com"
-          path: "/api/v1/lists"
-          authenticator:
-            type: "BearerAuthenticator"
-            api_token: "{{ config.apikey }}"
-          request_parameters:
-            page_size: 10
-          record_selector:
-          extractor:
-            field_path: ["result"]
-      schema_loader:
-        type: JsonFileSchemaLoader
-        file_path: "./source_test/schemas/lists.yaml"
-    conditions_stream:
-      type: DeclarativeStream
-      name: sometimes
-      primary_key: id
-      retriever:
-        paginator:
-          type: "DefaultPaginator"
-          page_size_option:
-            type: RequestOption
-            inject_into: request_parameter
-            field_name: page_size
-          page_token_option:
-            type: RequestPath
-          pagination_strategy:
-            type: "CursorPagination"
-            cursor_value: "{{ response._metadata.next }}"
-            page_size: 10
-        requester:
-          type: HttpRequester
-          url_base: "https://testing.com"
-          path: "/api/v1/sometimes"
-          authenticator:
-            type: "BearerAuthenticator"
-            api_token: "{{ config.apikey }}"
-          request_parameters:
-            page_size: 10
-        record_selector:
-          type: RecordSelector
-          extractor:
-            type: DpathExtractor
-            field_path: ["result"]
-      schema_loader:
-        type: JsonFileSchemaLoader
-        file_path: "./source_test/schemas/sometimes.yaml"
-    streams:
-      - "#/lists_stream"
-      - type: ConditionalStreams
-        condition: "{{ config['is_sandbox'] }}"
-        streams:
-          - "#/conditions_stream"
-    """
-    parsed_manifest = YamlDeclarativeSource._parse(content)
-    resolved_manifest = resolver.preprocess_manifest(parsed_manifest)
-    resolved_manifest["type"] = "DeclarativeSource"
-    stream_manifest = transformer.propagate_types_and_parameters(
-        "", resolved_manifest["conditions_stream"], {}
-    )
-
-    conditional_streams_manifest = resolved_manifest["streams"][1]
-    conditional_streams_manifest["streams"] = [stream_manifest]
-
-    sandbox_config = {
-        **input_config,
-        "is_sandbox": True,
-    }
-
-    streams = factory.create_component(
-        model_type=ConditionalStreamsModel,
-        component_definition=conditional_streams_manifest,
-        config=sandbox_config,
-    )
-
-    assert len(streams) == 1
-    conditional_stream = streams[0]
-    assert isinstance(conditional_stream, DeclarativeStream)
-    assert conditional_stream.name == "sometimes"
-    assert conditional_stream.primary_key == "id"
-
-    assert isinstance(conditional_stream.retriever, SimpleRetriever)
-    assert conditional_stream.retriever.name == "sometimes"
-    assert conditional_stream.retriever.primary_key == "id"