airbyte_cdk.sources.declarative.concurrent_declarative_source

  1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
  2
  3import json
  4import logging
  5import pkgutil
  6from contextlib import contextmanager
  7from copy import deepcopy
  8from dataclasses import dataclass, field
  9from queue import Queue
 10from types import ModuleType
 11from typing import (
 12    Any,
 13    ClassVar,
 14    Dict,
 15    Iterator,
 16    List,
 17    Mapping,
 18    Optional,
 19    Set,
 20)
 21
 22import orjson
 23import yaml
 24from airbyte_protocol_dataclasses.models import AirbyteStreamStatus, Level, StreamDescriptor
 25from jsonschema.exceptions import ValidationError
 26from jsonschema.validators import validate
 27
 28from airbyte_cdk.config_observation import create_connector_config_control_message
 29from airbyte_cdk.connector_builder.models import (
 30    LogMessage as ConnectorBuilderLogMessage,
 31)
 32from airbyte_cdk.manifest_migrations.migration_handler import (
 33    ManifestMigrationHandler,
 34)
 35from airbyte_cdk.models import (
 36    AirbyteCatalog,
 37    AirbyteConnectionStatus,
 38    AirbyteMessage,
 39    AirbyteStateMessage,
 40    ConfiguredAirbyteCatalog,
 41    ConnectorSpecification,
 42    FailureType,
 43    Status,
 44)
 45from airbyte_cdk.models.airbyte_protocol_serializers import AirbyteMessageSerializer
 46from airbyte_cdk.sources import Source
 47from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
 48from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 49from airbyte_cdk.sources.declarative.checks import COMPONENTS_CHECKER_TYPE_MAPPING
 50from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 51from airbyte_cdk.sources.declarative.concurrency_level import ConcurrencyLevel
 52from airbyte_cdk.sources.declarative.interpolation import InterpolatedBoolean
 53from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
 54    ConcurrencyLevel as ConcurrencyLevelModel,
 55)
 56from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
 57    DeclarativeStream as DeclarativeStreamModel,
 58)
 59from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
 60    Spec as SpecModel,
 61)
 62from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
 63    StateDelegatingStream as StateDelegatingStreamModel,
 64)
 65from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
 66    get_registered_components_module,
 67)
 68from airbyte_cdk.sources.declarative.parsers.manifest_component_transformer import (
 69    ManifestComponentTransformer,
 70)
 71from airbyte_cdk.sources.declarative.parsers.manifest_normalizer import (
 72    ManifestNormalizer,
 73)
 74from airbyte_cdk.sources.declarative.parsers.manifest_reference_resolver import (
 75    ManifestReferenceResolver,
 76)
 77from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
 78    ModelToComponentFactory,
 79)
 80from airbyte_cdk.sources.declarative.partition_routers.cartesian_product_stream_slicer import (
 81    CartesianProductStreamSlicer,
 82)
 83from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
 84    GroupingPartitionRouter,
 85)
 86from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
 87    SubstreamPartitionRouter,
 88)
 89from airbyte_cdk.sources.declarative.partition_routers.union_partition_router import (
 90    UnionPartitionRouter,
 91)
 92from airbyte_cdk.sources.declarative.resolvers import COMPONENTS_RESOLVER_TYPE_MAPPING
 93from airbyte_cdk.sources.declarative.spec.spec import Spec
 94from airbyte_cdk.sources.declarative.types import Config, ConnectionDefinition
 95from airbyte_cdk.sources.message.concurrent_repository import ConcurrentMessageRepository
 96from airbyte_cdk.sources.message.repository import InMemoryMessageRepository
 97from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 98from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 99from airbyte_cdk.sources.streams.concurrent.partitions.types import QueueItem
100from airbyte_cdk.sources.utils.slice_logger import (
101    AlwaysLogSliceLogger,
102    DebugSliceLogger,
103    SliceLogger,
104)
105from airbyte_cdk.utils.airbyte_secrets_utils import add_to_secrets, get_secrets
106from airbyte_cdk.utils.stream_status_utils import as_airbyte_message
107from airbyte_cdk.utils.traced_exception import AirbyteTracedException
108
109
110@dataclass
111class TestLimits:
112    __test__: ClassVar[bool] = False  # Tell Pytest this is not a Pytest class, despite its name
113
114    DEFAULT_MAX_PAGES_PER_SLICE: ClassVar[int] = 5
115    DEFAULT_MAX_SLICES: ClassVar[int] = 5
116    DEFAULT_MAX_RECORDS: ClassVar[int] = 100
117    DEFAULT_MAX_STREAMS: ClassVar[int] = 100
118
119    max_records: int = field(default=DEFAULT_MAX_RECORDS)
120    max_pages_per_slice: int = field(default=DEFAULT_MAX_PAGES_PER_SLICE)
121    max_slices: int = field(default=DEFAULT_MAX_SLICES)
122    max_streams: int = field(default=DEFAULT_MAX_STREAMS)
123
124
125def _get_declarative_component_schema() -> Dict[str, Any]:
126    try:
127        raw_component_schema = pkgutil.get_data(
128            "airbyte_cdk", "sources/declarative/declarative_component_schema.yaml"
129        )
130        if raw_component_schema is not None:
131            declarative_component_schema = yaml.load(raw_component_schema, Loader=yaml.SafeLoader)
132            return declarative_component_schema  # type: ignore
133        else:
134            raise RuntimeError(
135                "Failed to read manifest component json schema required for deduplication"
136            )
137    except FileNotFoundError as e:
138        raise FileNotFoundError(
139            f"Failed to read manifest component json schema required for deduplication: {e}"
140        )
141
142
143class ConcurrentDeclarativeSource(Source):
144    # By default, we defer to a value of 2. A value lower than could cause a PartitionEnqueuer to be stuck in a state of deadlock
145    # because it has hit the limit of futures but not partition reader is consuming them.
146    _LOWEST_SAFE_CONCURRENCY_LEVEL = 2
147
148    # Component types that hold the connector config and can write it back. `refresh_token_updater` is
149    # declared on `OAuthAuthenticator` alone; `CustomAuthenticator` is the type the transformer injects
150    # for a `class_name` component.
151    _CONFIG_PERSISTING_AUTHENTICATOR_TYPES = frozenset(
152        {"OAuthAuthenticator", "CustomAuthenticator"}
153    )
154
155    def __init__(
156        self,
157        catalog: Optional[ConfiguredAirbyteCatalog] = None,
158        config: Optional[Mapping[str, Any]] = None,
159        state: Optional[List[AirbyteStateMessage]] = None,
160        *,
161        source_config: ConnectionDefinition,
162        debug: bool = False,
163        emit_connector_builder_messages: bool = False,
164        migrate_manifest: bool = False,
165        normalize_manifest: bool = False,
166        limits: Optional[TestLimits] = None,
167        config_path: Optional[str] = None,
168        custom_components_trusted: bool = True,
169        **kwargs: Any,
170    ) -> None:
171        self.logger = logging.getLogger(f"airbyte.{self.name}")
172
173        self._limits = limits
174
175        # todo: We could remove state from initialization. Now that streams are grouped during the read(), a source
176        #  no longer needs to store the original incoming state. But maybe there's an edge case?
177        self._connector_state_manager = ConnectorStateManager(state=state)  # type: ignore  # state is always in the form of List[AirbyteStateMessage]. The ConnectorStateManager should use generics, but this can be done later
178
179        # We set a maxsize to for the main thread to process record items when the queue size grows. This assumes that there are less
180        # threads generating partitions that than are max number of workers. If it weren't the case, we could have threads only generating
181        # partitions which would fill the queue. This number is arbitrarily set to 10_000 but will probably need to be changed given more
182        # information and might even need to be configurable depending on the source
183        queue: Queue[QueueItem] = Queue(maxsize=10_000)
184        message_repository = InMemoryMessageRepository(
185            Level.DEBUG if emit_connector_builder_messages else Level.INFO
186        )
187
188        # To reduce the complexity of the concurrent framework, we are not enabling RFR with synthetic
189        # cursors. We do this by no longer automatically instantiating RFR cursors when converting
190        # the declarative models into runtime components. Concurrent sources will continue to checkpoint
191        # incremental streams running in full refresh.
192        component_factory = ModelToComponentFactory(
193            custom_components_trusted=custom_components_trusted,
194            emit_connector_builder_messages=emit_connector_builder_messages,
195            message_repository=ConcurrentMessageRepository(queue, message_repository),
196            configured_catalog=catalog,
197            connector_state_manager=self._connector_state_manager,
198            max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"),
199            limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None,
200            limit_slices_fetched=limits.max_slices if limits else None,
201            disable_retries=True if limits else False,
202            disable_cache=True if limits else False,
203        )
204
205        self._should_normalize = normalize_manifest
206        self._should_migrate = migrate_manifest
207        self._declarative_component_schema = _get_declarative_component_schema()
208        # If custom components are needed, locate and/or register them.
209        self.components_module: ModuleType | None = get_registered_components_module(config=config)
210        # set additional attributes
211        self._debug = debug
212        self._emit_connector_builder_messages = emit_connector_builder_messages
213        self._constructor = (
214            component_factory
215            if component_factory
216            else ModelToComponentFactory(
217                custom_components_trusted=custom_components_trusted,
218                emit_connector_builder_messages=emit_connector_builder_messages,
219                max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"),
220            )
221        )
222
223        self._message_repository = self._constructor.get_message_repository()
224        self._slice_logger: SliceLogger = (
225            AlwaysLogSliceLogger() if emit_connector_builder_messages else DebugSliceLogger()
226        )
227
228        # resolve all components in the manifest
229        self._source_config = self._pre_process_manifest(dict(source_config))
230        # validate resolved manifest against the declarative component schema
231        self._validate_source()
232        # apply additional post-processing to the manifest
233        self._post_process_manifest()
234
235        spec: Optional[Mapping[str, Any]] = self._source_config.get("spec")
236        self._spec_component: Optional[Spec] = (
237            self._constructor.create_component(SpecModel, spec, dict()) if spec else None
238        )
239        self._config = self._migrate_and_transform_config(config_path, config) or {}
240        # `check` may temporarily overlay values onto `self._config` (see
241        # `_config_overridden_for_check`). The manifest's `config_validations` express intent about the
242        # config as supplied, so they must run against it rather than against a check-time overlay.
243        self._config_for_validation = self._config
244
245        concurrency_level_from_manifest = self._source_config.get("concurrency_level")
246        if concurrency_level_from_manifest:
247            concurrency_level_component = self._constructor.create_component(
248                model_type=ConcurrencyLevelModel,
249                component_definition=concurrency_level_from_manifest,
250                config=config or {},
251            )
252            if not isinstance(concurrency_level_component, ConcurrencyLevel):
253                raise ValueError(
254                    f"Expected to generate a ConcurrencyLevel component, but received {concurrency_level_component.__class__}"
255                )
256
257            concurrency_level = concurrency_level_component.get_concurrency_level()
258            initial_number_of_partitions_to_generate = max(
259                concurrency_level // 2, 1
260            )  # Partition_generation iterates using range based on this value. If this is floored to zero we end up in a dead lock during start up
261        else:
262            concurrency_level = self._LOWEST_SAFE_CONCURRENCY_LEVEL
263            initial_number_of_partitions_to_generate = self._LOWEST_SAFE_CONCURRENCY_LEVEL // 2
264
265        self._concurrent_source = ConcurrentSource.create(
266            num_workers=concurrency_level,
267            initial_number_of_partitions_to_generate=initial_number_of_partitions_to_generate,
268            logger=self.logger,
269            slice_logger=self._slice_logger,
270            queue=queue,
271            message_repository=self._message_repository,
272        )
273
274    def _pre_process_manifest(self, manifest: Dict[str, Any]) -> Dict[str, Any]:
275        """
276        Preprocesses the provided manifest dictionary by resolving any manifest references.
277
278        This method modifies the input manifest in place, resolving references using the
279        ManifestReferenceResolver to ensure all references within the manifest are properly handled.
280
281        Args:
282            manifest (Dict[str, Any]): The manifest dictionary to preprocess and resolve references in.
283
284        Returns:
285            None
286        """
287        # For ease of use we don't require the type to be specified at the top level manifest, but it should be included during processing
288        manifest = self._fix_source_type(manifest)
289        # Resolve references in the manifest
290        resolved_manifest = ManifestReferenceResolver().preprocess_manifest(manifest)
291        # Propagate types and parameters throughout the manifest
292        propagated_manifest = ManifestComponentTransformer().propagate_types_and_parameters(
293            "", resolved_manifest, {}
294        )
295
296        return propagated_manifest
297
298    def _fix_source_type(self, manifest: Dict[str, Any]) -> Dict[str, Any]:
299        """
300        Fix the source type in the manifest. This is necessary because the source type is not always set in the manifest.
301        """
302        if "type" not in manifest:
303            manifest["type"] = "DeclarativeSource"
304
305        return manifest
306
307    def _post_process_manifest(self) -> None:
308        """
309        Post-processes the manifest after validation.
310        This method is responsible for any additional modifications or transformations needed
311        after the manifest has been validated and before it is used in the source.
312        """
313        # apply manifest migration, if required
314        self._migrate_manifest()
315        # apply manifest normalization, if required
316        self._normalize_manifest()
317
318    def _migrate_manifest(self) -> None:
319        """
320        This method is used to migrate the manifest. It should be called after the manifest has been validated.
321        The migration is done in place, so the original manifest is modified.
322
323        The original manifest is returned if any error occurs during migration.
324        """
325        if self._should_migrate:
326            manifest_migrator = ManifestMigrationHandler(self._source_config)
327            self._source_config = manifest_migrator.apply_migrations()
328            # validate migrated manifest against the declarative component schema
329            self._validate_source()
330
331    def _normalize_manifest(self) -> None:
332        """
333        This method is used to normalize the manifest. It should be called after the manifest has been validated.
334
335        Connector Builder UI rendering requires the manifest to be in a specific format.
336         - references have been resolved
337         - the commonly used definitions are extracted to the `definitions.linked.*`
338        """
339        if self._should_normalize:
340            normalizer = ManifestNormalizer(self._source_config, self._declarative_component_schema)
341            self._source_config = normalizer.normalize()
342
343    def _validate_source(self) -> None:
344        """
345        Validates the connector manifest against the declarative component schema
346        """
347
348        try:
349            validate(self._source_config, self._declarative_component_schema)
350        except ValidationError as e:
351            raise ValidationError(
352                "Validation against json schema defined in declarative_component_schema.yaml schema failed"
353            ) from e
354
355    def _migrate_and_transform_config(
356        self,
357        config_path: Optional[str],
358        config: Optional[Config],
359    ) -> Optional[Config]:
360        if not config:
361            return None
362        if not self._spec_component:
363            return config
364        mutable_config = dict(config)
365        self._spec_component.migrate_config(mutable_config)
366        if mutable_config != config:
367            if config_path:
368                with open(config_path, "w") as f:
369                    json.dump(mutable_config, f)
370            control_message = create_connector_config_control_message(mutable_config)
371            print(orjson.dumps(AirbyteMessageSerializer.dump(control_message)).decode())
372        self._spec_component.transform_config(mutable_config)
373        return mutable_config
374
375    def configure(self, config: Mapping[str, Any], temp_dir: str) -> Mapping[str, Any]:
376        config = self._config or config
377        return super().configure(config, temp_dir)
378
379    @property
380    def resolved_manifest(self) -> Mapping[str, Any]:
381        """
382        Returns the resolved manifest configuration for the source.
383
384        This property provides access to the internal source configuration as a mapping,
385        which contains all settings and parameters required to define the source's behavior.
386
387        Returns:
388            Mapping[str, Any]: The resolved source configuration manifest.
389        """
390        return self._source_config
391
392    def deprecation_warnings(self) -> List[ConnectorBuilderLogMessage]:
393        return self._constructor.get_model_deprecations()
394
395    def read(
396        self,
397        logger: logging.Logger,
398        config: Mapping[str, Any],
399        catalog: ConfiguredAirbyteCatalog,
400        state: Optional[List[AirbyteStateMessage]] = None,
401    ) -> Iterator[AirbyteMessage]:
402        selected_concurrent_streams = self._select_streams(
403            streams=self.streams(config=self._config),  # type: ignore  # We are migrating away from the DeclarativeStream implementation and streams() only returns the concurrent-compatible AbstractStream. To preserve compatibility, we retain the existing method interface
404            configured_catalog=catalog,
405        )
406
407        # It would appear that passing in an empty set of streams causes an infinite loop in ConcurrentReadProcessor.
408        # This is also evident in concurrent_source_adapter.py so I'll leave this out of scope to fix for now
409        if len(selected_concurrent_streams) > 0:
410            yield from self._concurrent_source.read(selected_concurrent_streams)
411
412    def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog:
413        return AirbyteCatalog(
414            streams=[stream.as_airbyte_stream() for stream in self.streams(config=self._config)]
415        )
416
417    # todo: add PR comment about whether we can change the signature to List[AbstractStream]
418    def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]:  # type: ignore  # we are migrating away from the AbstractSource and are expecting that this will only be called by ConcurrentDeclarativeSource or the Connector Builder
419        """
420        The `streams` method is used as part of the AbstractSource in the following cases:
421        * ConcurrentDeclarativeSource.check -> ManifestDeclarativeSource.check -> AbstractSource.check -> DeclarativeSource.check_connection -> CheckStream.check_connection -> streams
422        * ConcurrentDeclarativeSource.read -> AbstractSource.read -> streams (note that we filter for a specific catalog which excludes concurrent streams so not all streams actually read from all the streams returned by `streams`)
423        Note that `super.streams(config)` is also called when splitting the streams between concurrent or not in `_group_streams`.
424
425        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.
426        """
427
428        if self._spec_component:
429            self._spec_component.validate_config(self._config_for_validation)
430
431        api_budget_model = self._source_config.get("api_budget")
432        if api_budget_model:
433            self._constructor.set_api_budget(api_budget_model, self._config)
434
435        stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams
436
437        prepared_configs = self._initialize_cache_for_parent_streams(deepcopy(stream_configs))
438
439        source_streams = [
440            self._constructor.create_component(
441                (
442                    StateDelegatingStreamModel
443                    if stream_config.get("type") == StateDelegatingStreamModel.__name__
444                    else DeclarativeStreamModel
445                ),
446                stream_config,
447                self._config,
448                emit_connector_builder_messages=self._emit_connector_builder_messages,
449            )
450            for stream_config in prepared_configs
451        ]
452
453        self._apply_stream_groups(source_streams)
454
455        return source_streams
456
457    def _apply_stream_groups(self, streams: List[AbstractStream]) -> None:
458        """Set block_simultaneous_read on streams based on the manifest's stream_groups config.
459
460        Iterates over the resolved manifest's stream_groups and matches group membership
461        against actual created stream instances by name. Validates that no stream shares a
462        group with any of its parent streams, which would cause a deadlock.
463        """
464        stream_groups = self._source_config.get("stream_groups", {})
465        if not stream_groups:
466            return
467
468        # Build stream_name -> group_name mapping from the resolved manifest
469        stream_name_to_group: Dict[str, str] = {}
470        for group_name, group_config in stream_groups.items():
471            for stream_ref in group_config.get("streams", []):
472                if isinstance(stream_ref, dict):
473                    stream_name = stream_ref.get("name", "")
474                    if stream_name:
475                        stream_name_to_group[stream_name] = group_name
476
477        # Validate no stream shares a group with any of its ancestor streams
478        stream_name_to_instance: Dict[str, AbstractStream] = {s.name: s for s in streams}
479
480        def _collect_all_ancestor_names(stream_name: str) -> Set[str]:
481            """Recursively collect all ancestor stream names."""
482            ancestors: Set[str] = set()
483            inst = stream_name_to_instance.get(stream_name)
484            if not isinstance(inst, DefaultStream):
485                return ancestors
486            partition_router = inst.get_partition_router()
487            routers = [partition_router] if partition_router is not None else []
488            while routers:
489                router = routers.pop()
490                if isinstance(router, GroupingPartitionRouter):
491                    routers.append(router.underlying_partition_router)
492                elif isinstance(router, UnionPartitionRouter):
493                    routers.extend(router.partition_routers)
494                elif isinstance(router, CartesianProductStreamSlicer):
495                    routers.extend(router.stream_slicers)
496                elif isinstance(router, SubstreamPartitionRouter):
497                    for parent_config in router.parent_stream_configs:
498                        parent_name = parent_config.stream.name
499                        ancestors.add(parent_name)
500                        ancestors.update(_collect_all_ancestor_names(parent_name))
501            return ancestors
502
503        for stream in streams:
504            if not isinstance(stream, DefaultStream) or stream.name not in stream_name_to_group:
505                continue
506            group_name = stream_name_to_group[stream.name]
507            for ancestor_name in _collect_all_ancestor_names(stream.name):
508                if stream_name_to_group.get(ancestor_name) == group_name:
509                    raise ValueError(
510                        f"Stream '{stream.name}' and its parent stream '{ancestor_name}' "
511                        f"are both in group '{group_name}'. "
512                        f"A child stream must not share a group with its parent to avoid deadlock."
513                    )
514
515        # Apply group to matching stream instances
516        for stream in streams:
517            if isinstance(stream, DefaultStream) and stream.name in stream_name_to_group:
518                stream.block_simultaneous_read = stream_name_to_group[stream.name]
519
520    @staticmethod
521    def _initialize_cache_for_parent_streams(
522        stream_configs: List[Dict[str, Any]],
523    ) -> List[Dict[str, Any]]:
524        """Enable caching for parent streams unless explicitly disabled.
525
526        Caching is enabled by default for parent streams to optimize performance when the same
527        parent data is needed by multiple child streams. However, explicit `use_cache: false`
528        settings are respected for streams that cannot use caching (e.g., scroll-based pagination
529        APIs where caching causes duplicate records).
530        """
531        parent_streams = set()
532
533        def _set_cache_if_not_disabled(requester: Dict[str, Any]) -> None:
534            """Set use_cache to True only if not explicitly disabled."""
535            if requester.get("use_cache") is not False:
536                requester["use_cache"] = True
537
538        def update_with_cache_parent_configs(
539            parent_configs: list[dict[str, Any]],
540        ) -> None:
541            for parent_config in parent_configs:
542                parent_streams.add(parent_config["stream"]["name"])
543                if parent_config["stream"]["type"] == "StateDelegatingStream":
544                    _set_cache_if_not_disabled(
545                        parent_config["stream"]["full_refresh_stream"]["retriever"]["requester"]
546                    )
547                    _set_cache_if_not_disabled(
548                        parent_config["stream"]["incremental_stream"]["retriever"]["requester"]
549                    )
550                else:
551                    _set_cache_if_not_disabled(parent_config["stream"]["retriever"]["requester"])
552
553        for stream_config in stream_configs:
554            if stream_config.get("incremental_sync", {}).get("parent_stream"):
555                parent_streams.add(stream_config["incremental_sync"]["parent_stream"]["name"])
556                _set_cache_if_not_disabled(
557                    stream_config["incremental_sync"]["parent_stream"]["retriever"]["requester"]
558                )
559
560            elif stream_config.get("retriever", {}).get("partition_router", {}):
561                partition_router = stream_config["retriever"]["partition_router"]
562
563                routers = (
564                    list(partition_router)
565                    if isinstance(partition_router, list)
566                    else [partition_router]
567                )
568                while routers:
569                    router = routers.pop()
570                    if not isinstance(router, dict):
571                        continue
572                    if router.get("parent_stream_configs"):
573                        update_with_cache_parent_configs(router["parent_stream_configs"])
574                    # Descend into composed partition routers (e.g. UnionPartitionRouter's
575                    # partition_routers or GroupingPartitionRouter's underlying_partition_router)
576                    # so nested parent streams also get caching enabled.
577                    routers.extend(router.get("partition_routers") or [])
578                    if router.get("underlying_partition_router"):
579                        routers.append(router["underlying_partition_router"])
580
581        for stream_config in stream_configs:
582            if stream_config["name"] in parent_streams:
583                if stream_config["type"] == "StateDelegatingStream":
584                    _set_cache_if_not_disabled(
585                        stream_config["full_refresh_stream"]["retriever"]["requester"]
586                    )
587                    _set_cache_if_not_disabled(
588                        stream_config["incremental_stream"]["retriever"]["requester"]
589                    )
590                else:
591                    _set_cache_if_not_disabled(stream_config["retriever"]["requester"])
592        return stream_configs
593
594    def spec(self, logger: logging.Logger) -> ConnectorSpecification:
595        """
596        Returns the connector specification (spec) as defined in the Airbyte Protocol. The spec is an object describing the possible
597        configurations (e.g: username and password) which can be configured when running this connector. For low-code connectors, this
598        will first attempt to load the spec from the manifest's spec block, otherwise it will load it from "spec.yaml" or "spec.json"
599        in the project root.
600        """
601        return (
602            self._spec_component.generate_spec() if self._spec_component else super().spec(logger)
603        )
604
605    def check(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
606        check = self._source_config.get("check")
607        if not check:
608            raise ValueError(f"Missing 'check' component definition within the manifest.")
609
610        if "type" not in check:
611            check["type"] = "CheckStream"
612        connection_checker = self._constructor.create_component(
613            COMPONENTS_CHECKER_TYPE_MAPPING[check["type"]],
614            check,
615            dict(),
616            emit_connector_builder_messages=self._emit_connector_builder_messages,
617        )
618        if not isinstance(connection_checker, ConnectionChecker):
619            raise ValueError(
620                f"Expected to generate a ConnectionChecker component, but received {connection_checker.__class__}"
621            )
622
623        with self._config_overridden_for_check(check.get("config_overrides")):
624            check_succeeded, error = connection_checker.check_connection(self, logger, self._config)
625        if not check_succeeded:
626            return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error))
627        return AirbyteConnectionStatus(status=Status.SUCCEEDED)
628
629    @contextmanager
630    def _config_overridden_for_check(
631        self, config_overrides: Optional[Mapping[str, Any]]
632    ) -> Iterator[None]:
633        """Overlay the check component's `config_overrides` onto the config for the duration of a check.
634
635        Rebinding `self._config` reaches the whole component tree, because `streams()` interpolates from
636        it and ignores its own `config` argument. Not thread safe; `check` is one command per process.
637
638        Semantics:
639          - values are applied verbatim, never interpolated;
640          - the merge is one level deep, so an object-valued override replaces rather than merges;
641          - it happens after `_migrate_and_transform_config`, so derived fields are not recomputed;
642          - `config_validations` run against `self._config_for_validation`, not the overlay.
643
644        The copy is shallow, so a component writing into a nested path writes through to the config the
645        source was constructed with and the restore does not undo it. Only top-level writes are
646        discarded, which is why `_raise_if_config_is_persisted` refuses config-persisting manifests.
647
648        Known limitation: `$parameters` declared on a check component still propagate into object-valued
649        overrides. See https://github.com/airbytehq/airbyte-internal-issues/issues/16995.
650        """
651        if not config_overrides:
652            yield
653            return
654
655        self._raise_on_non_string_override_keys(config_overrides)
656        self._raise_on_reserved_override_keys(config_overrides)
657        self._raise_if_config_is_persisted(config_overrides)
658        self._warn_on_unknown_override_keys(config_overrides)
659        self._register_override_secrets(config_overrides)
660        # Keys only. An override may name a secret field, so values must not reach the logs.
661        self.logger.info(
662            f"Overriding config keys for the check operation: {', '.join(sorted(config_overrides))}"
663        )
664
665        unmodified_config = self._config
666        self._config = {**self._config, **config_overrides}
667        try:
668            yield
669        finally:
670            self._config = unmodified_config
671
672    @staticmethod
673    def _raise_on_reserved_override_keys(config_overrides: Mapping[str, Any]) -> None:
674        """Refuse keys in the platform's reserved `__airbyte` namespace.
675
676        `CheckStream` reads `__airbyte_check_stream_names` out of the config this overlay writes to, so
677        an override there would change which streams a check tests. The whole prefix is refused.
678        """
679        reserved = sorted(key for key in config_overrides if key.startswith("__airbyte"))
680        if reserved:
681            raise AirbyteTracedException(
682                message=(
683                    f"This connector's manifest is invalid: its check component overrides the "
684                    f"reserved config key(s) {reserved}. Keys prefixed with `__airbyte` belong to "
685                    "the platform, not to the connector's spec. This is a bug in the connector "
686                    "rather than in this connection's settings."
687                ),
688                internal_message=(
689                    f"config_overrides rejected: reserved __airbyte key(s) {reserved}. "
690                    "`CheckStream` reads `__airbyte_check_stream_names` out of the config this "
691                    "overlay writes to, so an override there would change which streams a check "
692                    "tests. Use `stream_names` instead."
693                ),
694                failure_type=FailureType.system_error,
695            )
696
697    @staticmethod
698    def _raise_on_non_string_override_keys(config_overrides: Mapping[str, Any]) -> None:
699        """Reject non-string keys, which YAML allows and every consumer below assumes away."""
700        non_strings = [key for key in config_overrides if not isinstance(key, str)]
701        if non_strings:
702            raise AirbyteTracedException(
703                message=(
704                    "This connector's manifest is invalid: its check component has "
705                    f"`config_overrides` key(s) {sorted(map(repr, non_strings))} that are not "
706                    "strings, so they cannot name a field in the connector's spec. This is a bug "
707                    "in the connector rather than in this connection's settings."
708                ),
709                internal_message=(
710                    "config_overrides rejected: non-string key(s) "
711                    f"{sorted(map(repr, non_strings))}. Quote them in the manifest so they name a "
712                    "field in the connector's spec."
713                ),
714                failure_type=FailureType.system_error,
715            )
716
717    def _raise_if_config_is_persisted(self, config_overrides: Mapping[str, Any]) -> None:
718        """Reject `config_overrides` on a manifest whose authenticator writes the config back.
719
720        A `refresh_token_updater` builds a `DeclarativeSingleUseRefreshTokenOauth2Authenticator`, whose
721        `_emit_control_message` emits the entire config it was handed as a `CONNECTOR_CONFIG` message.
722        During a check that is the overlay, and the platform persists it, so a check-only override would
723        become the connection's saved config. The restore cannot recall a message already on stdout.
724        """
725        if not self._manifest_writes_back_config(self._source_config):
726            return
727        raise AirbyteTracedException(
728            message=(
729                "This connector's manifest is invalid: `config_overrides` cannot be used by a "
730                "manifest that declares a `refresh_token_updater`. A token refresh during `check` "
731                "emits the whole config it was handed as a CONNECTOR_CONFIG control message, which "
732                f"the platform persists, so the check-only override(s) {sorted(config_overrides)} "
733                "would be saved as this connection's config and applied to every later sync. This "
734                "is a bug in the connector rather than in this connection's settings."
735            ),
736            internal_message=(
737                "config_overrides rejected: manifest persists config via refresh_token_updater. "
738                "Remove `config_overrides` from the check component, or drop the "
739                "`refresh_token_updater`."
740            ),
741            failure_type=FailureType.system_error,
742        )
743
744    @staticmethod
745    def _manifest_writes_back_config(definition: Any) -> bool:
746        """Whether any component in the manifest emits the connector config back to the platform.
747
748        Walks the whole manifest, since an authenticator can sit under any requester, inside a
749        `SelectiveAuthenticator`, or in a `ConditionalStreams` branch. Callers pass `self._source_config`,
750        which is post-`_pre_process_manifest`, so a `$ref`-ed authenticator is already inlined and
751        carries its own `type`.
752
753        `is not None` rather than truthiness because `refresh_token_updater: {}` takes every default and
754        still builds a single-use authenticator - the factory tests a model instance, which is truthy.
755        The type condition is what separates a component from a data blob that happens to use the name.
756
757        Does not detect a `CustomAuthenticator` that persists config without declaring the field.
758        """
759        if isinstance(definition, Mapping):
760            if (
761                definition.get("refresh_token_updater") is not None
762                and definition.get("type")
763                in ConcurrentDeclarativeSource._CONFIG_PERSISTING_AUTHENTICATOR_TYPES
764            ):
765                return True
766            return any(
767                ConcurrentDeclarativeSource._manifest_writes_back_config(value)
768                for value in definition.values()
769            )
770        if isinstance(definition, list):
771            return any(
772                ConcurrentDeclarativeSource._manifest_writes_back_config(item)
773                for item in definition
774            )
775        return False
776
777    def _warn_on_unknown_override_keys(self, config_overrides: Mapping[str, Any]) -> None:
778        """Warn about override keys the spec does not declare.
779
780        Nothing validates the overlay - the entrypoint validates what the user supplied, and
781        `config_validations` run against `self._config_for_validation` - so a typo is a silent no-op.
782        """
783        if not self._spec_component:
784            return
785        declared = self._declared_config_properties(self._spec_component.connection_specification)
786        if declared is None:
787            return
788        unknown = sorted(key for key in config_overrides if key not in declared)
789        if unknown:
790            self.logger.warning(
791                f"Check-only config override(s) {unknown} are not declared in the connector spec, so "
792                "they will have no effect on any component that reads the config by field name."
793            )
794
795    @staticmethod
796    def _declared_config_properties(
797        connection_specification: Mapping[str, Any],
798    ) -> Optional[Set[str]]:
799        """Every field name a spec declares, or `None` when it does not enumerate them.
800
801        `oneOf` and `allOf` composition put declarations one level down, so reading only the top-level
802        `properties` would report a declared field as undeclared.
803        """
804        if not isinstance(connection_specification, Mapping):
805            return None
806
807        names: Set[str] = set()
808        found_any = False
809        properties = connection_specification.get("properties")
810        if isinstance(properties, Mapping):
811            found_any = True
812            names.update(str(key) for key in properties)
813        for keyword in ("allOf", "anyOf", "oneOf"):
814            for branch in connection_specification.get(keyword) or []:
815                nested = ConcurrentDeclarativeSource._declared_config_properties(branch)
816                if nested is not None:
817                    found_any = True
818                    names.update(nested)
819        return names if found_any else None
820
821    def _register_override_secrets(self, config_overrides: Mapping[str, Any]) -> None:
822        """Register override values landing on an `airbyte_secret` field, so they get redacted.
823
824        The entrypoint builds the secret list from the config the user supplied, so a value substituted
825        here is unknown to `filter_secrets`. Uses the same discovery as the entrypoint, so the overlay is
826        redacted exactly where the user's own value at that path would be.
827        """
828        if not self._spec_component:
829            return
830        for secret in get_secrets(
831            self._spec_component.connection_specification, dict(config_overrides)
832        ):
833            if secret is not None:
834                add_to_secrets(str(secret))
835
836    @property
837    def dynamic_streams(self) -> List[Dict[str, Any]]:
838        return self._dynamic_stream_configs(
839            manifest=self._source_config,
840            with_dynamic_stream_name=True,
841        )
842
843    def _stream_configs(self, manifest: Mapping[str, Any]) -> List[Dict[str, Any]]:
844        # This has a warning flag for static, but after we finish part 4 we'll replace manifest with self._source_config
845        stream_configs = []
846        for current_stream_config in manifest.get("streams", []):
847            if (
848                "type" in current_stream_config
849                and current_stream_config["type"] == "ConditionalStreams"
850            ):
851                interpolated_boolean = InterpolatedBoolean(
852                    condition=current_stream_config.get("condition"),
853                    parameters={},
854                )
855
856                if interpolated_boolean.eval(config=self._config):
857                    stream_configs.extend(current_stream_config.get("streams", []))
858            else:
859                if "type" not in current_stream_config:
860                    current_stream_config["type"] = "DeclarativeStream"
861                stream_configs.append(current_stream_config)
862        return stream_configs
863
864    def _dynamic_stream_configs(
865        self,
866        manifest: Mapping[str, Any],
867        with_dynamic_stream_name: Optional[bool] = None,
868    ) -> List[Dict[str, Any]]:
869        dynamic_stream_definitions: List[Dict[str, Any]] = manifest.get("dynamic_streams", [])
870        dynamic_stream_configs: List[Dict[str, Any]] = []
871        seen_dynamic_streams: Set[str] = set()
872
873        for dynamic_definition_index, dynamic_definition in enumerate(dynamic_stream_definitions):
874            components_resolver_config = dynamic_definition["components_resolver"]
875
876            if not components_resolver_config:
877                raise ValueError(
878                    f"Missing 'components_resolver' in dynamic definition: {dynamic_definition}"
879                )
880
881            resolver_type = components_resolver_config.get("type")
882            if not resolver_type:
883                raise ValueError(
884                    f"Missing 'type' in components resolver configuration: {components_resolver_config}"
885                )
886
887            if resolver_type not in COMPONENTS_RESOLVER_TYPE_MAPPING:
888                raise ValueError(
889                    f"Invalid components resolver type '{resolver_type}'. "
890                    f"Expected one of {list(COMPONENTS_RESOLVER_TYPE_MAPPING.keys())}."
891                )
892
893            if "retriever" in components_resolver_config:
894                components_resolver_config["retriever"]["requester"]["use_cache"] = True
895
896            # Create a resolver for dynamic components based on type
897            if resolver_type == "HttpComponentsResolver":
898                components_resolver = self._constructor.create_component(
899                    model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type],
900                    component_definition=components_resolver_config,
901                    config=self._config,
902                    stream_name=dynamic_definition.get("name"),
903                )
904            else:
905                components_resolver = self._constructor.create_component(
906                    model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type],
907                    component_definition=components_resolver_config,
908                    config=self._config,
909                )
910
911            stream_template_config = dynamic_definition["stream_template"]
912
913            for dynamic_stream in components_resolver.resolve_components(
914                stream_template_config=stream_template_config
915            ):
916                # Get the use_parent_parameters configuration from the dynamic definition
917                # Default to True for backward compatibility, since connectors were already using it by default when this param was added
918                use_parent_parameters = dynamic_definition.get("use_parent_parameters", True)
919
920                dynamic_stream = {
921                    **ManifestComponentTransformer().propagate_types_and_parameters(
922                        "", dynamic_stream, {}, use_parent_parameters=use_parent_parameters
923                    )
924                }
925
926                if "type" not in dynamic_stream:
927                    dynamic_stream["type"] = "DeclarativeStream"
928
929                # Ensure that each stream is created with a unique name
930                name = dynamic_stream.get("name")
931
932                if with_dynamic_stream_name:
933                    dynamic_stream["dynamic_stream_name"] = dynamic_definition.get(
934                        "name", f"dynamic_stream_{dynamic_definition_index}"
935                    )
936
937                if not isinstance(name, str):
938                    raise ValueError(
939                        f"Expected stream name {name} to be a string, got {type(name)}."
940                    )
941
942                if name in seen_dynamic_streams:
943                    error_message = f"Dynamic streams list contains a duplicate name: {name}. Please contact Airbyte Support."
944                    failure_type = FailureType.system_error
945
946                    if resolver_type == "ConfigComponentsResolver":
947                        error_message = f"Dynamic streams list contains a duplicate name: {name}. Please check your configuration."
948                        failure_type = FailureType.config_error
949
950                    raise AirbyteTracedException(
951                        message=error_message,
952                        internal_message=error_message,
953                        failure_type=failure_type,
954                    )
955
956                seen_dynamic_streams.add(name)
957                dynamic_stream_configs.append(dynamic_stream)
958
959        return dynamic_stream_configs
960
961    def _select_streams(
962        self, streams: List[AbstractStream], configured_catalog: ConfiguredAirbyteCatalog
963    ) -> List[AbstractStream]:
964        stream_name_to_instance: Mapping[str, AbstractStream] = {s.name: s for s in streams}
965        abstract_streams: List[AbstractStream] = []
966        for configured_stream in configured_catalog.streams:
967            stream_instance = stream_name_to_instance.get(configured_stream.stream.name)
968            if stream_instance:
969                abstract_streams.append(stream_instance)
970            else:
971                # Previous behavior in the legacy synchronous CDK was to also raise an error TRACE message if
972                # the source was configured with raise_exception_on_missing_stream=True. This was used on very
973                # few sources like facebook-marketing and google-ads. We decided not to port this feature over,
974                # but we can do so if we feel it necessary. With the current behavior,we should still result
975                # in a partial failure since missing streams will be marked as INCOMPLETE.
976                self._message_repository.emit_message(
977                    as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE)
978                )
979        return abstract_streams