airbyte_cdk.sources.declarative.concurrent_declarative_source
1# Copyright (c) 2025 Airbyte, Inc., all rights reserved. 2 3import json 4import logging 5import pkgutil 6from copy import deepcopy 7from dataclasses import dataclass, field 8from queue import Queue 9from types import ModuleType 10from typing import ( 11 Any, 12 ClassVar, 13 Dict, 14 Iterator, 15 List, 16 Mapping, 17 Optional, 18 Set, 19) 20 21import orjson 22import yaml 23from airbyte_protocol_dataclasses.models import AirbyteStreamStatus, Level, StreamDescriptor 24from jsonschema.exceptions import ValidationError 25from jsonschema.validators import validate 26 27from airbyte_cdk.config_observation import create_connector_config_control_message 28from airbyte_cdk.connector_builder.models import ( 29 LogMessage as ConnectorBuilderLogMessage, 30) 31from airbyte_cdk.manifest_migrations.migration_handler import ( 32 ManifestMigrationHandler, 33) 34from airbyte_cdk.models import ( 35 AirbyteCatalog, 36 AirbyteConnectionStatus, 37 AirbyteMessage, 38 AirbyteStateMessage, 39 ConfiguredAirbyteCatalog, 40 ConnectorSpecification, 41 FailureType, 42 Status, 43) 44from airbyte_cdk.models.airbyte_protocol_serializers import AirbyteMessageSerializer 45from airbyte_cdk.sources import Source 46from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource 47from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager 48from airbyte_cdk.sources.declarative.checks import COMPONENTS_CHECKER_TYPE_MAPPING 49from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker 50from airbyte_cdk.sources.declarative.concurrency_level import ConcurrencyLevel 51from airbyte_cdk.sources.declarative.interpolation import InterpolatedBoolean 52from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( 53 ConcurrencyLevel as ConcurrencyLevelModel, 54) 55from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( 56 DeclarativeStream as DeclarativeStreamModel, 57) 58from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( 59 Spec as SpecModel, 60) 61from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( 62 StateDelegatingStream as StateDelegatingStreamModel, 63) 64from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import ( 65 get_registered_components_module, 66) 67from airbyte_cdk.sources.declarative.parsers.manifest_component_transformer import ( 68 ManifestComponentTransformer, 69) 70from airbyte_cdk.sources.declarative.parsers.manifest_normalizer import ( 71 ManifestNormalizer, 72) 73from airbyte_cdk.sources.declarative.parsers.manifest_reference_resolver import ( 74 ManifestReferenceResolver, 75) 76from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ( 77 ModelToComponentFactory, 78) 79from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import ( 80 GroupingPartitionRouter, 81) 82from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import ( 83 SubstreamPartitionRouter, 84) 85from airbyte_cdk.sources.declarative.resolvers import COMPONENTS_RESOLVER_TYPE_MAPPING 86from airbyte_cdk.sources.declarative.spec.spec import Spec 87from airbyte_cdk.sources.declarative.types import Config, ConnectionDefinition 88from airbyte_cdk.sources.message.concurrent_repository import ConcurrentMessageRepository 89from airbyte_cdk.sources.message.repository import InMemoryMessageRepository 90from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream 91from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream 92from airbyte_cdk.sources.streams.concurrent.partitions.types import QueueItem 93from airbyte_cdk.sources.utils.slice_logger import ( 94 AlwaysLogSliceLogger, 95 DebugSliceLogger, 96 SliceLogger, 97) 98from airbyte_cdk.utils.stream_status_utils import as_airbyte_message 99from airbyte_cdk.utils.traced_exception import AirbyteTracedException 100 101 102@dataclass 103class TestLimits: 104 __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name 105 106 DEFAULT_MAX_PAGES_PER_SLICE: ClassVar[int] = 5 107 DEFAULT_MAX_SLICES: ClassVar[int] = 5 108 DEFAULT_MAX_RECORDS: ClassVar[int] = 100 109 DEFAULT_MAX_STREAMS: ClassVar[int] = 100 110 111 max_records: int = field(default=DEFAULT_MAX_RECORDS) 112 max_pages_per_slice: int = field(default=DEFAULT_MAX_PAGES_PER_SLICE) 113 max_slices: int = field(default=DEFAULT_MAX_SLICES) 114 max_streams: int = field(default=DEFAULT_MAX_STREAMS) 115 116 117def _get_declarative_component_schema() -> Dict[str, Any]: 118 try: 119 raw_component_schema = pkgutil.get_data( 120 "airbyte_cdk", "sources/declarative/declarative_component_schema.yaml" 121 ) 122 if raw_component_schema is not None: 123 declarative_component_schema = yaml.load(raw_component_schema, Loader=yaml.SafeLoader) 124 return declarative_component_schema # type: ignore 125 else: 126 raise RuntimeError( 127 "Failed to read manifest component json schema required for deduplication" 128 ) 129 except FileNotFoundError as e: 130 raise FileNotFoundError( 131 f"Failed to read manifest component json schema required for deduplication: {e}" 132 ) 133 134 135class ConcurrentDeclarativeSource(Source): 136 # 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 137 # because it has hit the limit of futures but not partition reader is consuming them. 138 _LOWEST_SAFE_CONCURRENCY_LEVEL = 2 139 140 def __init__( 141 self, 142 catalog: Optional[ConfiguredAirbyteCatalog] = None, 143 config: Optional[Mapping[str, Any]] = None, 144 state: Optional[List[AirbyteStateMessage]] = None, 145 *, 146 source_config: ConnectionDefinition, 147 debug: bool = False, 148 emit_connector_builder_messages: bool = False, 149 migrate_manifest: bool = False, 150 normalize_manifest: bool = False, 151 limits: Optional[TestLimits] = None, 152 config_path: Optional[str] = None, 153 custom_components_trusted: bool = True, 154 **kwargs: Any, 155 ) -> None: 156 self.logger = logging.getLogger(f"airbyte.{self.name}") 157 158 self._limits = limits 159 160 # todo: We could remove state from initialization. Now that streams are grouped during the read(), a source 161 # no longer needs to store the original incoming state. But maybe there's an edge case? 162 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 163 164 # We set a maxsize to for the main thread to process record items when the queue size grows. This assumes that there are less 165 # threads generating partitions that than are max number of workers. If it weren't the case, we could have threads only generating 166 # partitions which would fill the queue. This number is arbitrarily set to 10_000 but will probably need to be changed given more 167 # information and might even need to be configurable depending on the source 168 queue: Queue[QueueItem] = Queue(maxsize=10_000) 169 message_repository = InMemoryMessageRepository( 170 Level.DEBUG if emit_connector_builder_messages else Level.INFO 171 ) 172 173 # To reduce the complexity of the concurrent framework, we are not enabling RFR with synthetic 174 # cursors. We do this by no longer automatically instantiating RFR cursors when converting 175 # the declarative models into runtime components. Concurrent sources will continue to checkpoint 176 # incremental streams running in full refresh. 177 component_factory = ModelToComponentFactory( 178 custom_components_trusted=custom_components_trusted, 179 emit_connector_builder_messages=emit_connector_builder_messages, 180 message_repository=ConcurrentMessageRepository(queue, message_repository), 181 configured_catalog=catalog, 182 connector_state_manager=self._connector_state_manager, 183 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), 184 limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None, 185 limit_slices_fetched=limits.max_slices if limits else None, 186 disable_retries=True if limits else False, 187 disable_cache=True if limits else False, 188 ) 189 190 self._should_normalize = normalize_manifest 191 self._should_migrate = migrate_manifest 192 self._declarative_component_schema = _get_declarative_component_schema() 193 # If custom components are needed, locate and/or register them. 194 self.components_module: ModuleType | None = get_registered_components_module(config=config) 195 # set additional attributes 196 self._debug = debug 197 self._emit_connector_builder_messages = emit_connector_builder_messages 198 self._constructor = ( 199 component_factory 200 if component_factory 201 else ModelToComponentFactory( 202 custom_components_trusted=custom_components_trusted, 203 emit_connector_builder_messages=emit_connector_builder_messages, 204 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), 205 ) 206 ) 207 208 self._message_repository = self._constructor.get_message_repository() 209 self._slice_logger: SliceLogger = ( 210 AlwaysLogSliceLogger() if emit_connector_builder_messages else DebugSliceLogger() 211 ) 212 213 # resolve all components in the manifest 214 self._source_config = self._pre_process_manifest(dict(source_config)) 215 # validate resolved manifest against the declarative component schema 216 self._validate_source() 217 # apply additional post-processing to the manifest 218 self._post_process_manifest() 219 220 spec: Optional[Mapping[str, Any]] = self._source_config.get("spec") 221 self._spec_component: Optional[Spec] = ( 222 self._constructor.create_component(SpecModel, spec, dict()) if spec else None 223 ) 224 self._config = self._migrate_and_transform_config(config_path, config) or {} 225 226 concurrency_level_from_manifest = self._source_config.get("concurrency_level") 227 if concurrency_level_from_manifest: 228 concurrency_level_component = self._constructor.create_component( 229 model_type=ConcurrencyLevelModel, 230 component_definition=concurrency_level_from_manifest, 231 config=config or {}, 232 ) 233 if not isinstance(concurrency_level_component, ConcurrencyLevel): 234 raise ValueError( 235 f"Expected to generate a ConcurrencyLevel component, but received {concurrency_level_component.__class__}" 236 ) 237 238 concurrency_level = concurrency_level_component.get_concurrency_level() 239 initial_number_of_partitions_to_generate = max( 240 concurrency_level // 2, 1 241 ) # 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 242 else: 243 concurrency_level = self._LOWEST_SAFE_CONCURRENCY_LEVEL 244 initial_number_of_partitions_to_generate = self._LOWEST_SAFE_CONCURRENCY_LEVEL // 2 245 246 self._concurrent_source = ConcurrentSource.create( 247 num_workers=concurrency_level, 248 initial_number_of_partitions_to_generate=initial_number_of_partitions_to_generate, 249 logger=self.logger, 250 slice_logger=self._slice_logger, 251 queue=queue, 252 message_repository=self._message_repository, 253 ) 254 255 def _pre_process_manifest(self, manifest: Dict[str, Any]) -> Dict[str, Any]: 256 """ 257 Preprocesses the provided manifest dictionary by resolving any manifest references. 258 259 This method modifies the input manifest in place, resolving references using the 260 ManifestReferenceResolver to ensure all references within the manifest are properly handled. 261 262 Args: 263 manifest (Dict[str, Any]): The manifest dictionary to preprocess and resolve references in. 264 265 Returns: 266 None 267 """ 268 # 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 269 manifest = self._fix_source_type(manifest) 270 # Resolve references in the manifest 271 resolved_manifest = ManifestReferenceResolver().preprocess_manifest(manifest) 272 # Propagate types and parameters throughout the manifest 273 propagated_manifest = ManifestComponentTransformer().propagate_types_and_parameters( 274 "", resolved_manifest, {} 275 ) 276 277 return propagated_manifest 278 279 def _fix_source_type(self, manifest: Dict[str, Any]) -> Dict[str, Any]: 280 """ 281 Fix the source type in the manifest. This is necessary because the source type is not always set in the manifest. 282 """ 283 if "type" not in manifest: 284 manifest["type"] = "DeclarativeSource" 285 286 return manifest 287 288 def _post_process_manifest(self) -> None: 289 """ 290 Post-processes the manifest after validation. 291 This method is responsible for any additional modifications or transformations needed 292 after the manifest has been validated and before it is used in the source. 293 """ 294 # apply manifest migration, if required 295 self._migrate_manifest() 296 # apply manifest normalization, if required 297 self._normalize_manifest() 298 299 def _migrate_manifest(self) -> None: 300 """ 301 This method is used to migrate the manifest. It should be called after the manifest has been validated. 302 The migration is done in place, so the original manifest is modified. 303 304 The original manifest is returned if any error occurs during migration. 305 """ 306 if self._should_migrate: 307 manifest_migrator = ManifestMigrationHandler(self._source_config) 308 self._source_config = manifest_migrator.apply_migrations() 309 # validate migrated manifest against the declarative component schema 310 self._validate_source() 311 312 def _normalize_manifest(self) -> None: 313 """ 314 This method is used to normalize the manifest. It should be called after the manifest has been validated. 315 316 Connector Builder UI rendering requires the manifest to be in a specific format. 317 - references have been resolved 318 - the commonly used definitions are extracted to the `definitions.linked.*` 319 """ 320 if self._should_normalize: 321 normalizer = ManifestNormalizer(self._source_config, self._declarative_component_schema) 322 self._source_config = normalizer.normalize() 323 324 def _validate_source(self) -> None: 325 """ 326 Validates the connector manifest against the declarative component schema 327 """ 328 329 try: 330 validate(self._source_config, self._declarative_component_schema) 331 except ValidationError as e: 332 raise ValidationError( 333 "Validation against json schema defined in declarative_component_schema.yaml schema failed" 334 ) from e 335 336 def _migrate_and_transform_config( 337 self, 338 config_path: Optional[str], 339 config: Optional[Config], 340 ) -> Optional[Config]: 341 if not config: 342 return None 343 if not self._spec_component: 344 return config 345 mutable_config = dict(config) 346 self._spec_component.migrate_config(mutable_config) 347 if mutable_config != config: 348 if config_path: 349 with open(config_path, "w") as f: 350 json.dump(mutable_config, f) 351 control_message = create_connector_config_control_message(mutable_config) 352 print(orjson.dumps(AirbyteMessageSerializer.dump(control_message)).decode()) 353 self._spec_component.transform_config(mutable_config) 354 return mutable_config 355 356 def configure(self, config: Mapping[str, Any], temp_dir: str) -> Mapping[str, Any]: 357 config = self._config or config 358 return super().configure(config, temp_dir) 359 360 @property 361 def resolved_manifest(self) -> Mapping[str, Any]: 362 """ 363 Returns the resolved manifest configuration for the source. 364 365 This property provides access to the internal source configuration as a mapping, 366 which contains all settings and parameters required to define the source's behavior. 367 368 Returns: 369 Mapping[str, Any]: The resolved source configuration manifest. 370 """ 371 return self._source_config 372 373 def deprecation_warnings(self) -> List[ConnectorBuilderLogMessage]: 374 return self._constructor.get_model_deprecations() 375 376 def read( 377 self, 378 logger: logging.Logger, 379 config: Mapping[str, Any], 380 catalog: ConfiguredAirbyteCatalog, 381 state: Optional[List[AirbyteStateMessage]] = None, 382 ) -> Iterator[AirbyteMessage]: 383 selected_concurrent_streams = self._select_streams( 384 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 385 configured_catalog=catalog, 386 ) 387 388 # It would appear that passing in an empty set of streams causes an infinite loop in ConcurrentReadProcessor. 389 # This is also evident in concurrent_source_adapter.py so I'll leave this out of scope to fix for now 390 if len(selected_concurrent_streams) > 0: 391 yield from self._concurrent_source.read(selected_concurrent_streams) 392 393 def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog: 394 return AirbyteCatalog( 395 streams=[stream.as_airbyte_stream() for stream in self.streams(config=self._config)] 396 ) 397 398 # todo: add PR comment about whether we can change the signature to List[AbstractStream] 399 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 400 """ 401 The `streams` method is used as part of the AbstractSource in the following cases: 402 * ConcurrentDeclarativeSource.check -> ManifestDeclarativeSource.check -> AbstractSource.check -> DeclarativeSource.check_connection -> CheckStream.check_connection -> streams 403 * 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`) 404 Note that `super.streams(config)` is also called when splitting the streams between concurrent or not in `_group_streams`. 405 406 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. 407 """ 408 409 if self._spec_component: 410 self._spec_component.validate_config(self._config) 411 412 api_budget_model = self._source_config.get("api_budget") 413 if api_budget_model: 414 self._constructor.set_api_budget(api_budget_model, self._config) 415 416 stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams 417 418 prepared_configs = self._initialize_cache_for_parent_streams(deepcopy(stream_configs)) 419 420 source_streams = [ 421 self._constructor.create_component( 422 ( 423 StateDelegatingStreamModel 424 if stream_config.get("type") == StateDelegatingStreamModel.__name__ 425 else DeclarativeStreamModel 426 ), 427 stream_config, 428 self._config, 429 emit_connector_builder_messages=self._emit_connector_builder_messages, 430 ) 431 for stream_config in prepared_configs 432 ] 433 434 self._apply_stream_groups(source_streams) 435 436 return source_streams 437 438 def _apply_stream_groups(self, streams: List[AbstractStream]) -> None: 439 """Set block_simultaneous_read on streams based on the manifest's stream_groups config. 440 441 Iterates over the resolved manifest's stream_groups and matches group membership 442 against actual created stream instances by name. Validates that no stream shares a 443 group with any of its parent streams, which would cause a deadlock. 444 """ 445 stream_groups = self._source_config.get("stream_groups", {}) 446 if not stream_groups: 447 return 448 449 # Build stream_name -> group_name mapping from the resolved manifest 450 stream_name_to_group: Dict[str, str] = {} 451 for group_name, group_config in stream_groups.items(): 452 for stream_ref in group_config.get("streams", []): 453 if isinstance(stream_ref, dict): 454 stream_name = stream_ref.get("name", "") 455 if stream_name: 456 stream_name_to_group[stream_name] = group_name 457 458 # Validate no stream shares a group with any of its ancestor streams 459 stream_name_to_instance: Dict[str, AbstractStream] = {s.name: s for s in streams} 460 461 def _collect_all_ancestor_names(stream_name: str) -> Set[str]: 462 """Recursively collect all ancestor stream names.""" 463 ancestors: Set[str] = set() 464 inst = stream_name_to_instance.get(stream_name) 465 if not isinstance(inst, DefaultStream): 466 return ancestors 467 router = inst.get_partition_router() 468 if isinstance(router, GroupingPartitionRouter): 469 router = router.underlying_partition_router 470 if not isinstance(router, SubstreamPartitionRouter): 471 return ancestors 472 for parent_config in router.parent_stream_configs: 473 parent_name = parent_config.stream.name 474 ancestors.add(parent_name) 475 ancestors.update(_collect_all_ancestor_names(parent_name)) 476 return ancestors 477 478 for stream in streams: 479 if not isinstance(stream, DefaultStream) or stream.name not in stream_name_to_group: 480 continue 481 group_name = stream_name_to_group[stream.name] 482 for ancestor_name in _collect_all_ancestor_names(stream.name): 483 if stream_name_to_group.get(ancestor_name) == group_name: 484 raise ValueError( 485 f"Stream '{stream.name}' and its parent stream '{ancestor_name}' " 486 f"are both in group '{group_name}'. " 487 f"A child stream must not share a group with its parent to avoid deadlock." 488 ) 489 490 # Apply group to matching stream instances 491 for stream in streams: 492 if isinstance(stream, DefaultStream) and stream.name in stream_name_to_group: 493 stream.block_simultaneous_read = stream_name_to_group[stream.name] 494 495 @staticmethod 496 def _initialize_cache_for_parent_streams( 497 stream_configs: List[Dict[str, Any]], 498 ) -> List[Dict[str, Any]]: 499 """Enable caching for parent streams unless explicitly disabled. 500 501 Caching is enabled by default for parent streams to optimize performance when the same 502 parent data is needed by multiple child streams. However, explicit `use_cache: false` 503 settings are respected for streams that cannot use caching (e.g., scroll-based pagination 504 APIs where caching causes duplicate records). 505 """ 506 parent_streams = set() 507 508 def _set_cache_if_not_disabled(requester: Dict[str, Any]) -> None: 509 """Set use_cache to True only if not explicitly disabled.""" 510 if requester.get("use_cache") is not False: 511 requester["use_cache"] = True 512 513 def update_with_cache_parent_configs( 514 parent_configs: list[dict[str, Any]], 515 ) -> None: 516 for parent_config in parent_configs: 517 parent_streams.add(parent_config["stream"]["name"]) 518 if parent_config["stream"]["type"] == "StateDelegatingStream": 519 _set_cache_if_not_disabled( 520 parent_config["stream"]["full_refresh_stream"]["retriever"]["requester"] 521 ) 522 _set_cache_if_not_disabled( 523 parent_config["stream"]["incremental_stream"]["retriever"]["requester"] 524 ) 525 else: 526 _set_cache_if_not_disabled(parent_config["stream"]["retriever"]["requester"]) 527 528 for stream_config in stream_configs: 529 if stream_config.get("incremental_sync", {}).get("parent_stream"): 530 parent_streams.add(stream_config["incremental_sync"]["parent_stream"]["name"]) 531 _set_cache_if_not_disabled( 532 stream_config["incremental_sync"]["parent_stream"]["retriever"]["requester"] 533 ) 534 535 elif stream_config.get("retriever", {}).get("partition_router", {}): 536 partition_router = stream_config["retriever"]["partition_router"] 537 538 if isinstance(partition_router, dict) and partition_router.get( 539 "parent_stream_configs" 540 ): 541 update_with_cache_parent_configs(partition_router["parent_stream_configs"]) 542 elif isinstance(partition_router, list): 543 for router in partition_router: 544 if router.get("parent_stream_configs"): 545 update_with_cache_parent_configs(router["parent_stream_configs"]) 546 547 for stream_config in stream_configs: 548 if stream_config["name"] in parent_streams: 549 if stream_config["type"] == "StateDelegatingStream": 550 _set_cache_if_not_disabled( 551 stream_config["full_refresh_stream"]["retriever"]["requester"] 552 ) 553 _set_cache_if_not_disabled( 554 stream_config["incremental_stream"]["retriever"]["requester"] 555 ) 556 else: 557 _set_cache_if_not_disabled(stream_config["retriever"]["requester"]) 558 return stream_configs 559 560 def spec(self, logger: logging.Logger) -> ConnectorSpecification: 561 """ 562 Returns the connector specification (spec) as defined in the Airbyte Protocol. The spec is an object describing the possible 563 configurations (e.g: username and password) which can be configured when running this connector. For low-code connectors, this 564 will first attempt to load the spec from the manifest's spec block, otherwise it will load it from "spec.yaml" or "spec.json" 565 in the project root. 566 """ 567 return ( 568 self._spec_component.generate_spec() if self._spec_component else super().spec(logger) 569 ) 570 571 def check(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteConnectionStatus: 572 check = self._source_config.get("check") 573 if not check: 574 raise ValueError(f"Missing 'check' component definition within the manifest.") 575 576 if "type" not in check: 577 check["type"] = "CheckStream" 578 connection_checker = self._constructor.create_component( 579 COMPONENTS_CHECKER_TYPE_MAPPING[check["type"]], 580 check, 581 dict(), 582 emit_connector_builder_messages=self._emit_connector_builder_messages, 583 ) 584 if not isinstance(connection_checker, ConnectionChecker): 585 raise ValueError( 586 f"Expected to generate a ConnectionChecker component, but received {connection_checker.__class__}" 587 ) 588 589 check_succeeded, error = connection_checker.check_connection(self, logger, self._config) 590 if not check_succeeded: 591 return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error)) 592 return AirbyteConnectionStatus(status=Status.SUCCEEDED) 593 594 @property 595 def dynamic_streams(self) -> List[Dict[str, Any]]: 596 return self._dynamic_stream_configs( 597 manifest=self._source_config, 598 with_dynamic_stream_name=True, 599 ) 600 601 def _stream_configs(self, manifest: Mapping[str, Any]) -> List[Dict[str, Any]]: 602 # This has a warning flag for static, but after we finish part 4 we'll replace manifest with self._source_config 603 stream_configs = [] 604 for current_stream_config in manifest.get("streams", []): 605 if ( 606 "type" in current_stream_config 607 and current_stream_config["type"] == "ConditionalStreams" 608 ): 609 interpolated_boolean = InterpolatedBoolean( 610 condition=current_stream_config.get("condition"), 611 parameters={}, 612 ) 613 614 if interpolated_boolean.eval(config=self._config): 615 stream_configs.extend(current_stream_config.get("streams", [])) 616 else: 617 if "type" not in current_stream_config: 618 current_stream_config["type"] = "DeclarativeStream" 619 stream_configs.append(current_stream_config) 620 return stream_configs 621 622 def _dynamic_stream_configs( 623 self, 624 manifest: Mapping[str, Any], 625 with_dynamic_stream_name: Optional[bool] = None, 626 ) -> List[Dict[str, Any]]: 627 dynamic_stream_definitions: List[Dict[str, Any]] = manifest.get("dynamic_streams", []) 628 dynamic_stream_configs: List[Dict[str, Any]] = [] 629 seen_dynamic_streams: Set[str] = set() 630 631 for dynamic_definition_index, dynamic_definition in enumerate(dynamic_stream_definitions): 632 components_resolver_config = dynamic_definition["components_resolver"] 633 634 if not components_resolver_config: 635 raise ValueError( 636 f"Missing 'components_resolver' in dynamic definition: {dynamic_definition}" 637 ) 638 639 resolver_type = components_resolver_config.get("type") 640 if not resolver_type: 641 raise ValueError( 642 f"Missing 'type' in components resolver configuration: {components_resolver_config}" 643 ) 644 645 if resolver_type not in COMPONENTS_RESOLVER_TYPE_MAPPING: 646 raise ValueError( 647 f"Invalid components resolver type '{resolver_type}'. " 648 f"Expected one of {list(COMPONENTS_RESOLVER_TYPE_MAPPING.keys())}." 649 ) 650 651 if "retriever" in components_resolver_config: 652 components_resolver_config["retriever"]["requester"]["use_cache"] = True 653 654 # Create a resolver for dynamic components based on type 655 if resolver_type == "HttpComponentsResolver": 656 components_resolver = self._constructor.create_component( 657 model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type], 658 component_definition=components_resolver_config, 659 config=self._config, 660 stream_name=dynamic_definition.get("name"), 661 ) 662 else: 663 components_resolver = self._constructor.create_component( 664 model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type], 665 component_definition=components_resolver_config, 666 config=self._config, 667 ) 668 669 stream_template_config = dynamic_definition["stream_template"] 670 671 for dynamic_stream in components_resolver.resolve_components( 672 stream_template_config=stream_template_config 673 ): 674 # Get the use_parent_parameters configuration from the dynamic definition 675 # Default to True for backward compatibility, since connectors were already using it by default when this param was added 676 use_parent_parameters = dynamic_definition.get("use_parent_parameters", True) 677 678 dynamic_stream = { 679 **ManifestComponentTransformer().propagate_types_and_parameters( 680 "", dynamic_stream, {}, use_parent_parameters=use_parent_parameters 681 ) 682 } 683 684 if "type" not in dynamic_stream: 685 dynamic_stream["type"] = "DeclarativeStream" 686 687 # Ensure that each stream is created with a unique name 688 name = dynamic_stream.get("name") 689 690 if with_dynamic_stream_name: 691 dynamic_stream["dynamic_stream_name"] = dynamic_definition.get( 692 "name", f"dynamic_stream_{dynamic_definition_index}" 693 ) 694 695 if not isinstance(name, str): 696 raise ValueError( 697 f"Expected stream name {name} to be a string, got {type(name)}." 698 ) 699 700 if name in seen_dynamic_streams: 701 error_message = f"Dynamic streams list contains a duplicate name: {name}. Please contact Airbyte Support." 702 failure_type = FailureType.system_error 703 704 if resolver_type == "ConfigComponentsResolver": 705 error_message = f"Dynamic streams list contains a duplicate name: {name}. Please check your configuration." 706 failure_type = FailureType.config_error 707 708 raise AirbyteTracedException( 709 message=error_message, 710 internal_message=error_message, 711 failure_type=failure_type, 712 ) 713 714 seen_dynamic_streams.add(name) 715 dynamic_stream_configs.append(dynamic_stream) 716 717 return dynamic_stream_configs 718 719 def _select_streams( 720 self, streams: List[AbstractStream], configured_catalog: ConfiguredAirbyteCatalog 721 ) -> List[AbstractStream]: 722 stream_name_to_instance: Mapping[str, AbstractStream] = {s.name: s for s in streams} 723 abstract_streams: List[AbstractStream] = [] 724 for configured_stream in configured_catalog.streams: 725 stream_instance = stream_name_to_instance.get(configured_stream.stream.name) 726 if stream_instance: 727 abstract_streams.append(stream_instance) 728 else: 729 # Previous behavior in the legacy synchronous CDK was to also raise an error TRACE message if 730 # the source was configured with raise_exception_on_missing_stream=True. This was used on very 731 # few sources like facebook-marketing and google-ads. We decided not to port this feature over, 732 # but we can do so if we feel it necessary. With the current behavior,we should still result 733 # in a partial failure since missing streams will be marked as INCOMPLETE. 734 self._message_repository.emit_message( 735 as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE) 736 ) 737 return abstract_streams
@dataclass
class
TestLimits:
103@dataclass 104class TestLimits: 105 __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name 106 107 DEFAULT_MAX_PAGES_PER_SLICE: ClassVar[int] = 5 108 DEFAULT_MAX_SLICES: ClassVar[int] = 5 109 DEFAULT_MAX_RECORDS: ClassVar[int] = 100 110 DEFAULT_MAX_STREAMS: ClassVar[int] = 100 111 112 max_records: int = field(default=DEFAULT_MAX_RECORDS) 113 max_pages_per_slice: int = field(default=DEFAULT_MAX_PAGES_PER_SLICE) 114 max_slices: int = field(default=DEFAULT_MAX_SLICES) 115 max_streams: int = field(default=DEFAULT_MAX_STREAMS)
class
ConcurrentDeclarativeSource(airbyte_cdk.connector.DefaultConnectorMixin, airbyte_cdk.sources.source.BaseSource[typing.Mapping[str, typing.Any], typing.List[airbyte_cdk.models.airbyte_protocol.AirbyteStateMessage], airbyte_protocol_dataclasses.models.airbyte_protocol.ConfiguredAirbyteCatalog], abc.ABC):
136class ConcurrentDeclarativeSource(Source): 137 # 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 138 # because it has hit the limit of futures but not partition reader is consuming them. 139 _LOWEST_SAFE_CONCURRENCY_LEVEL = 2 140 141 def __init__( 142 self, 143 catalog: Optional[ConfiguredAirbyteCatalog] = None, 144 config: Optional[Mapping[str, Any]] = None, 145 state: Optional[List[AirbyteStateMessage]] = None, 146 *, 147 source_config: ConnectionDefinition, 148 debug: bool = False, 149 emit_connector_builder_messages: bool = False, 150 migrate_manifest: bool = False, 151 normalize_manifest: bool = False, 152 limits: Optional[TestLimits] = None, 153 config_path: Optional[str] = None, 154 custom_components_trusted: bool = True, 155 **kwargs: Any, 156 ) -> None: 157 self.logger = logging.getLogger(f"airbyte.{self.name}") 158 159 self._limits = limits 160 161 # todo: We could remove state from initialization. Now that streams are grouped during the read(), a source 162 # no longer needs to store the original incoming state. But maybe there's an edge case? 163 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 164 165 # We set a maxsize to for the main thread to process record items when the queue size grows. This assumes that there are less 166 # threads generating partitions that than are max number of workers. If it weren't the case, we could have threads only generating 167 # partitions which would fill the queue. This number is arbitrarily set to 10_000 but will probably need to be changed given more 168 # information and might even need to be configurable depending on the source 169 queue: Queue[QueueItem] = Queue(maxsize=10_000) 170 message_repository = InMemoryMessageRepository( 171 Level.DEBUG if emit_connector_builder_messages else Level.INFO 172 ) 173 174 # To reduce the complexity of the concurrent framework, we are not enabling RFR with synthetic 175 # cursors. We do this by no longer automatically instantiating RFR cursors when converting 176 # the declarative models into runtime components. Concurrent sources will continue to checkpoint 177 # incremental streams running in full refresh. 178 component_factory = ModelToComponentFactory( 179 custom_components_trusted=custom_components_trusted, 180 emit_connector_builder_messages=emit_connector_builder_messages, 181 message_repository=ConcurrentMessageRepository(queue, message_repository), 182 configured_catalog=catalog, 183 connector_state_manager=self._connector_state_manager, 184 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), 185 limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None, 186 limit_slices_fetched=limits.max_slices if limits else None, 187 disable_retries=True if limits else False, 188 disable_cache=True if limits else False, 189 ) 190 191 self._should_normalize = normalize_manifest 192 self._should_migrate = migrate_manifest 193 self._declarative_component_schema = _get_declarative_component_schema() 194 # If custom components are needed, locate and/or register them. 195 self.components_module: ModuleType | None = get_registered_components_module(config=config) 196 # set additional attributes 197 self._debug = debug 198 self._emit_connector_builder_messages = emit_connector_builder_messages 199 self._constructor = ( 200 component_factory 201 if component_factory 202 else ModelToComponentFactory( 203 custom_components_trusted=custom_components_trusted, 204 emit_connector_builder_messages=emit_connector_builder_messages, 205 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), 206 ) 207 ) 208 209 self._message_repository = self._constructor.get_message_repository() 210 self._slice_logger: SliceLogger = ( 211 AlwaysLogSliceLogger() if emit_connector_builder_messages else DebugSliceLogger() 212 ) 213 214 # resolve all components in the manifest 215 self._source_config = self._pre_process_manifest(dict(source_config)) 216 # validate resolved manifest against the declarative component schema 217 self._validate_source() 218 # apply additional post-processing to the manifest 219 self._post_process_manifest() 220 221 spec: Optional[Mapping[str, Any]] = self._source_config.get("spec") 222 self._spec_component: Optional[Spec] = ( 223 self._constructor.create_component(SpecModel, spec, dict()) if spec else None 224 ) 225 self._config = self._migrate_and_transform_config(config_path, config) or {} 226 227 concurrency_level_from_manifest = self._source_config.get("concurrency_level") 228 if concurrency_level_from_manifest: 229 concurrency_level_component = self._constructor.create_component( 230 model_type=ConcurrencyLevelModel, 231 component_definition=concurrency_level_from_manifest, 232 config=config or {}, 233 ) 234 if not isinstance(concurrency_level_component, ConcurrencyLevel): 235 raise ValueError( 236 f"Expected to generate a ConcurrencyLevel component, but received {concurrency_level_component.__class__}" 237 ) 238 239 concurrency_level = concurrency_level_component.get_concurrency_level() 240 initial_number_of_partitions_to_generate = max( 241 concurrency_level // 2, 1 242 ) # 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 243 else: 244 concurrency_level = self._LOWEST_SAFE_CONCURRENCY_LEVEL 245 initial_number_of_partitions_to_generate = self._LOWEST_SAFE_CONCURRENCY_LEVEL // 2 246 247 self._concurrent_source = ConcurrentSource.create( 248 num_workers=concurrency_level, 249 initial_number_of_partitions_to_generate=initial_number_of_partitions_to_generate, 250 logger=self.logger, 251 slice_logger=self._slice_logger, 252 queue=queue, 253 message_repository=self._message_repository, 254 ) 255 256 def _pre_process_manifest(self, manifest: Dict[str, Any]) -> Dict[str, Any]: 257 """ 258 Preprocesses the provided manifest dictionary by resolving any manifest references. 259 260 This method modifies the input manifest in place, resolving references using the 261 ManifestReferenceResolver to ensure all references within the manifest are properly handled. 262 263 Args: 264 manifest (Dict[str, Any]): The manifest dictionary to preprocess and resolve references in. 265 266 Returns: 267 None 268 """ 269 # 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 270 manifest = self._fix_source_type(manifest) 271 # Resolve references in the manifest 272 resolved_manifest = ManifestReferenceResolver().preprocess_manifest(manifest) 273 # Propagate types and parameters throughout the manifest 274 propagated_manifest = ManifestComponentTransformer().propagate_types_and_parameters( 275 "", resolved_manifest, {} 276 ) 277 278 return propagated_manifest 279 280 def _fix_source_type(self, manifest: Dict[str, Any]) -> Dict[str, Any]: 281 """ 282 Fix the source type in the manifest. This is necessary because the source type is not always set in the manifest. 283 """ 284 if "type" not in manifest: 285 manifest["type"] = "DeclarativeSource" 286 287 return manifest 288 289 def _post_process_manifest(self) -> None: 290 """ 291 Post-processes the manifest after validation. 292 This method is responsible for any additional modifications or transformations needed 293 after the manifest has been validated and before it is used in the source. 294 """ 295 # apply manifest migration, if required 296 self._migrate_manifest() 297 # apply manifest normalization, if required 298 self._normalize_manifest() 299 300 def _migrate_manifest(self) -> None: 301 """ 302 This method is used to migrate the manifest. It should be called after the manifest has been validated. 303 The migration is done in place, so the original manifest is modified. 304 305 The original manifest is returned if any error occurs during migration. 306 """ 307 if self._should_migrate: 308 manifest_migrator = ManifestMigrationHandler(self._source_config) 309 self._source_config = manifest_migrator.apply_migrations() 310 # validate migrated manifest against the declarative component schema 311 self._validate_source() 312 313 def _normalize_manifest(self) -> None: 314 """ 315 This method is used to normalize the manifest. It should be called after the manifest has been validated. 316 317 Connector Builder UI rendering requires the manifest to be in a specific format. 318 - references have been resolved 319 - the commonly used definitions are extracted to the `definitions.linked.*` 320 """ 321 if self._should_normalize: 322 normalizer = ManifestNormalizer(self._source_config, self._declarative_component_schema) 323 self._source_config = normalizer.normalize() 324 325 def _validate_source(self) -> None: 326 """ 327 Validates the connector manifest against the declarative component schema 328 """ 329 330 try: 331 validate(self._source_config, self._declarative_component_schema) 332 except ValidationError as e: 333 raise ValidationError( 334 "Validation against json schema defined in declarative_component_schema.yaml schema failed" 335 ) from e 336 337 def _migrate_and_transform_config( 338 self, 339 config_path: Optional[str], 340 config: Optional[Config], 341 ) -> Optional[Config]: 342 if not config: 343 return None 344 if not self._spec_component: 345 return config 346 mutable_config = dict(config) 347 self._spec_component.migrate_config(mutable_config) 348 if mutable_config != config: 349 if config_path: 350 with open(config_path, "w") as f: 351 json.dump(mutable_config, f) 352 control_message = create_connector_config_control_message(mutable_config) 353 print(orjson.dumps(AirbyteMessageSerializer.dump(control_message)).decode()) 354 self._spec_component.transform_config(mutable_config) 355 return mutable_config 356 357 def configure(self, config: Mapping[str, Any], temp_dir: str) -> Mapping[str, Any]: 358 config = self._config or config 359 return super().configure(config, temp_dir) 360 361 @property 362 def resolved_manifest(self) -> Mapping[str, Any]: 363 """ 364 Returns the resolved manifest configuration for the source. 365 366 This property provides access to the internal source configuration as a mapping, 367 which contains all settings and parameters required to define the source's behavior. 368 369 Returns: 370 Mapping[str, Any]: The resolved source configuration manifest. 371 """ 372 return self._source_config 373 374 def deprecation_warnings(self) -> List[ConnectorBuilderLogMessage]: 375 return self._constructor.get_model_deprecations() 376 377 def read( 378 self, 379 logger: logging.Logger, 380 config: Mapping[str, Any], 381 catalog: ConfiguredAirbyteCatalog, 382 state: Optional[List[AirbyteStateMessage]] = None, 383 ) -> Iterator[AirbyteMessage]: 384 selected_concurrent_streams = self._select_streams( 385 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 386 configured_catalog=catalog, 387 ) 388 389 # It would appear that passing in an empty set of streams causes an infinite loop in ConcurrentReadProcessor. 390 # This is also evident in concurrent_source_adapter.py so I'll leave this out of scope to fix for now 391 if len(selected_concurrent_streams) > 0: 392 yield from self._concurrent_source.read(selected_concurrent_streams) 393 394 def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog: 395 return AirbyteCatalog( 396 streams=[stream.as_airbyte_stream() for stream in self.streams(config=self._config)] 397 ) 398 399 # todo: add PR comment about whether we can change the signature to List[AbstractStream] 400 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 401 """ 402 The `streams` method is used as part of the AbstractSource in the following cases: 403 * ConcurrentDeclarativeSource.check -> ManifestDeclarativeSource.check -> AbstractSource.check -> DeclarativeSource.check_connection -> CheckStream.check_connection -> streams 404 * 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`) 405 Note that `super.streams(config)` is also called when splitting the streams between concurrent or not in `_group_streams`. 406 407 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. 408 """ 409 410 if self._spec_component: 411 self._spec_component.validate_config(self._config) 412 413 api_budget_model = self._source_config.get("api_budget") 414 if api_budget_model: 415 self._constructor.set_api_budget(api_budget_model, self._config) 416 417 stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams 418 419 prepared_configs = self._initialize_cache_for_parent_streams(deepcopy(stream_configs)) 420 421 source_streams = [ 422 self._constructor.create_component( 423 ( 424 StateDelegatingStreamModel 425 if stream_config.get("type") == StateDelegatingStreamModel.__name__ 426 else DeclarativeStreamModel 427 ), 428 stream_config, 429 self._config, 430 emit_connector_builder_messages=self._emit_connector_builder_messages, 431 ) 432 for stream_config in prepared_configs 433 ] 434 435 self._apply_stream_groups(source_streams) 436 437 return source_streams 438 439 def _apply_stream_groups(self, streams: List[AbstractStream]) -> None: 440 """Set block_simultaneous_read on streams based on the manifest's stream_groups config. 441 442 Iterates over the resolved manifest's stream_groups and matches group membership 443 against actual created stream instances by name. Validates that no stream shares a 444 group with any of its parent streams, which would cause a deadlock. 445 """ 446 stream_groups = self._source_config.get("stream_groups", {}) 447 if not stream_groups: 448 return 449 450 # Build stream_name -> group_name mapping from the resolved manifest 451 stream_name_to_group: Dict[str, str] = {} 452 for group_name, group_config in stream_groups.items(): 453 for stream_ref in group_config.get("streams", []): 454 if isinstance(stream_ref, dict): 455 stream_name = stream_ref.get("name", "") 456 if stream_name: 457 stream_name_to_group[stream_name] = group_name 458 459 # Validate no stream shares a group with any of its ancestor streams 460 stream_name_to_instance: Dict[str, AbstractStream] = {s.name: s for s in streams} 461 462 def _collect_all_ancestor_names(stream_name: str) -> Set[str]: 463 """Recursively collect all ancestor stream names.""" 464 ancestors: Set[str] = set() 465 inst = stream_name_to_instance.get(stream_name) 466 if not isinstance(inst, DefaultStream): 467 return ancestors 468 router = inst.get_partition_router() 469 if isinstance(router, GroupingPartitionRouter): 470 router = router.underlying_partition_router 471 if not isinstance(router, SubstreamPartitionRouter): 472 return ancestors 473 for parent_config in router.parent_stream_configs: 474 parent_name = parent_config.stream.name 475 ancestors.add(parent_name) 476 ancestors.update(_collect_all_ancestor_names(parent_name)) 477 return ancestors 478 479 for stream in streams: 480 if not isinstance(stream, DefaultStream) or stream.name not in stream_name_to_group: 481 continue 482 group_name = stream_name_to_group[stream.name] 483 for ancestor_name in _collect_all_ancestor_names(stream.name): 484 if stream_name_to_group.get(ancestor_name) == group_name: 485 raise ValueError( 486 f"Stream '{stream.name}' and its parent stream '{ancestor_name}' " 487 f"are both in group '{group_name}'. " 488 f"A child stream must not share a group with its parent to avoid deadlock." 489 ) 490 491 # Apply group to matching stream instances 492 for stream in streams: 493 if isinstance(stream, DefaultStream) and stream.name in stream_name_to_group: 494 stream.block_simultaneous_read = stream_name_to_group[stream.name] 495 496 @staticmethod 497 def _initialize_cache_for_parent_streams( 498 stream_configs: List[Dict[str, Any]], 499 ) -> List[Dict[str, Any]]: 500 """Enable caching for parent streams unless explicitly disabled. 501 502 Caching is enabled by default for parent streams to optimize performance when the same 503 parent data is needed by multiple child streams. However, explicit `use_cache: false` 504 settings are respected for streams that cannot use caching (e.g., scroll-based pagination 505 APIs where caching causes duplicate records). 506 """ 507 parent_streams = set() 508 509 def _set_cache_if_not_disabled(requester: Dict[str, Any]) -> None: 510 """Set use_cache to True only if not explicitly disabled.""" 511 if requester.get("use_cache") is not False: 512 requester["use_cache"] = True 513 514 def update_with_cache_parent_configs( 515 parent_configs: list[dict[str, Any]], 516 ) -> None: 517 for parent_config in parent_configs: 518 parent_streams.add(parent_config["stream"]["name"]) 519 if parent_config["stream"]["type"] == "StateDelegatingStream": 520 _set_cache_if_not_disabled( 521 parent_config["stream"]["full_refresh_stream"]["retriever"]["requester"] 522 ) 523 _set_cache_if_not_disabled( 524 parent_config["stream"]["incremental_stream"]["retriever"]["requester"] 525 ) 526 else: 527 _set_cache_if_not_disabled(parent_config["stream"]["retriever"]["requester"]) 528 529 for stream_config in stream_configs: 530 if stream_config.get("incremental_sync", {}).get("parent_stream"): 531 parent_streams.add(stream_config["incremental_sync"]["parent_stream"]["name"]) 532 _set_cache_if_not_disabled( 533 stream_config["incremental_sync"]["parent_stream"]["retriever"]["requester"] 534 ) 535 536 elif stream_config.get("retriever", {}).get("partition_router", {}): 537 partition_router = stream_config["retriever"]["partition_router"] 538 539 if isinstance(partition_router, dict) and partition_router.get( 540 "parent_stream_configs" 541 ): 542 update_with_cache_parent_configs(partition_router["parent_stream_configs"]) 543 elif isinstance(partition_router, list): 544 for router in partition_router: 545 if router.get("parent_stream_configs"): 546 update_with_cache_parent_configs(router["parent_stream_configs"]) 547 548 for stream_config in stream_configs: 549 if stream_config["name"] in parent_streams: 550 if stream_config["type"] == "StateDelegatingStream": 551 _set_cache_if_not_disabled( 552 stream_config["full_refresh_stream"]["retriever"]["requester"] 553 ) 554 _set_cache_if_not_disabled( 555 stream_config["incremental_stream"]["retriever"]["requester"] 556 ) 557 else: 558 _set_cache_if_not_disabled(stream_config["retriever"]["requester"]) 559 return stream_configs 560 561 def spec(self, logger: logging.Logger) -> ConnectorSpecification: 562 """ 563 Returns the connector specification (spec) as defined in the Airbyte Protocol. The spec is an object describing the possible 564 configurations (e.g: username and password) which can be configured when running this connector. For low-code connectors, this 565 will first attempt to load the spec from the manifest's spec block, otherwise it will load it from "spec.yaml" or "spec.json" 566 in the project root. 567 """ 568 return ( 569 self._spec_component.generate_spec() if self._spec_component else super().spec(logger) 570 ) 571 572 def check(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteConnectionStatus: 573 check = self._source_config.get("check") 574 if not check: 575 raise ValueError(f"Missing 'check' component definition within the manifest.") 576 577 if "type" not in check: 578 check["type"] = "CheckStream" 579 connection_checker = self._constructor.create_component( 580 COMPONENTS_CHECKER_TYPE_MAPPING[check["type"]], 581 check, 582 dict(), 583 emit_connector_builder_messages=self._emit_connector_builder_messages, 584 ) 585 if not isinstance(connection_checker, ConnectionChecker): 586 raise ValueError( 587 f"Expected to generate a ConnectionChecker component, but received {connection_checker.__class__}" 588 ) 589 590 check_succeeded, error = connection_checker.check_connection(self, logger, self._config) 591 if not check_succeeded: 592 return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error)) 593 return AirbyteConnectionStatus(status=Status.SUCCEEDED) 594 595 @property 596 def dynamic_streams(self) -> List[Dict[str, Any]]: 597 return self._dynamic_stream_configs( 598 manifest=self._source_config, 599 with_dynamic_stream_name=True, 600 ) 601 602 def _stream_configs(self, manifest: Mapping[str, Any]) -> List[Dict[str, Any]]: 603 # This has a warning flag for static, but after we finish part 4 we'll replace manifest with self._source_config 604 stream_configs = [] 605 for current_stream_config in manifest.get("streams", []): 606 if ( 607 "type" in current_stream_config 608 and current_stream_config["type"] == "ConditionalStreams" 609 ): 610 interpolated_boolean = InterpolatedBoolean( 611 condition=current_stream_config.get("condition"), 612 parameters={}, 613 ) 614 615 if interpolated_boolean.eval(config=self._config): 616 stream_configs.extend(current_stream_config.get("streams", [])) 617 else: 618 if "type" not in current_stream_config: 619 current_stream_config["type"] = "DeclarativeStream" 620 stream_configs.append(current_stream_config) 621 return stream_configs 622 623 def _dynamic_stream_configs( 624 self, 625 manifest: Mapping[str, Any], 626 with_dynamic_stream_name: Optional[bool] = None, 627 ) -> List[Dict[str, Any]]: 628 dynamic_stream_definitions: List[Dict[str, Any]] = manifest.get("dynamic_streams", []) 629 dynamic_stream_configs: List[Dict[str, Any]] = [] 630 seen_dynamic_streams: Set[str] = set() 631 632 for dynamic_definition_index, dynamic_definition in enumerate(dynamic_stream_definitions): 633 components_resolver_config = dynamic_definition["components_resolver"] 634 635 if not components_resolver_config: 636 raise ValueError( 637 f"Missing 'components_resolver' in dynamic definition: {dynamic_definition}" 638 ) 639 640 resolver_type = components_resolver_config.get("type") 641 if not resolver_type: 642 raise ValueError( 643 f"Missing 'type' in components resolver configuration: {components_resolver_config}" 644 ) 645 646 if resolver_type not in COMPONENTS_RESOLVER_TYPE_MAPPING: 647 raise ValueError( 648 f"Invalid components resolver type '{resolver_type}'. " 649 f"Expected one of {list(COMPONENTS_RESOLVER_TYPE_MAPPING.keys())}." 650 ) 651 652 if "retriever" in components_resolver_config: 653 components_resolver_config["retriever"]["requester"]["use_cache"] = True 654 655 # Create a resolver for dynamic components based on type 656 if resolver_type == "HttpComponentsResolver": 657 components_resolver = self._constructor.create_component( 658 model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type], 659 component_definition=components_resolver_config, 660 config=self._config, 661 stream_name=dynamic_definition.get("name"), 662 ) 663 else: 664 components_resolver = self._constructor.create_component( 665 model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type], 666 component_definition=components_resolver_config, 667 config=self._config, 668 ) 669 670 stream_template_config = dynamic_definition["stream_template"] 671 672 for dynamic_stream in components_resolver.resolve_components( 673 stream_template_config=stream_template_config 674 ): 675 # Get the use_parent_parameters configuration from the dynamic definition 676 # Default to True for backward compatibility, since connectors were already using it by default when this param was added 677 use_parent_parameters = dynamic_definition.get("use_parent_parameters", True) 678 679 dynamic_stream = { 680 **ManifestComponentTransformer().propagate_types_and_parameters( 681 "", dynamic_stream, {}, use_parent_parameters=use_parent_parameters 682 ) 683 } 684 685 if "type" not in dynamic_stream: 686 dynamic_stream["type"] = "DeclarativeStream" 687 688 # Ensure that each stream is created with a unique name 689 name = dynamic_stream.get("name") 690 691 if with_dynamic_stream_name: 692 dynamic_stream["dynamic_stream_name"] = dynamic_definition.get( 693 "name", f"dynamic_stream_{dynamic_definition_index}" 694 ) 695 696 if not isinstance(name, str): 697 raise ValueError( 698 f"Expected stream name {name} to be a string, got {type(name)}." 699 ) 700 701 if name in seen_dynamic_streams: 702 error_message = f"Dynamic streams list contains a duplicate name: {name}. Please contact Airbyte Support." 703 failure_type = FailureType.system_error 704 705 if resolver_type == "ConfigComponentsResolver": 706 error_message = f"Dynamic streams list contains a duplicate name: {name}. Please check your configuration." 707 failure_type = FailureType.config_error 708 709 raise AirbyteTracedException( 710 message=error_message, 711 internal_message=error_message, 712 failure_type=failure_type, 713 ) 714 715 seen_dynamic_streams.add(name) 716 dynamic_stream_configs.append(dynamic_stream) 717 718 return dynamic_stream_configs 719 720 def _select_streams( 721 self, streams: List[AbstractStream], configured_catalog: ConfiguredAirbyteCatalog 722 ) -> List[AbstractStream]: 723 stream_name_to_instance: Mapping[str, AbstractStream] = {s.name: s for s in streams} 724 abstract_streams: List[AbstractStream] = [] 725 for configured_stream in configured_catalog.streams: 726 stream_instance = stream_name_to_instance.get(configured_stream.stream.name) 727 if stream_instance: 728 abstract_streams.append(stream_instance) 729 else: 730 # Previous behavior in the legacy synchronous CDK was to also raise an error TRACE message if 731 # the source was configured with raise_exception_on_missing_stream=True. This was used on very 732 # few sources like facebook-marketing and google-ads. We decided not to port this feature over, 733 # but we can do so if we feel it necessary. With the current behavior,we should still result 734 # in a partial failure since missing streams will be marked as INCOMPLETE. 735 self._message_repository.emit_message( 736 as_airbyte_message(configured_stream.stream, AirbyteStreamStatus.INCOMPLETE) 737 ) 738 return abstract_streams
Helper class that provides a standard way to create an ABC using inheritance.
ConcurrentDeclarativeSource( catalog: Optional[airbyte_protocol_dataclasses.models.airbyte_protocol.ConfiguredAirbyteCatalog] = None, config: Optional[Mapping[str, Any]] = None, state: Optional[List[airbyte_cdk.models.airbyte_protocol.AirbyteStateMessage]] = None, *, source_config: Mapping[str, Any], debug: bool = False, emit_connector_builder_messages: bool = False, migrate_manifest: bool = False, normalize_manifest: bool = False, limits: Optional[TestLimits] = None, config_path: Optional[str] = None, custom_components_trusted: bool = True, **kwargs: Any)
141 def __init__( 142 self, 143 catalog: Optional[ConfiguredAirbyteCatalog] = None, 144 config: Optional[Mapping[str, Any]] = None, 145 state: Optional[List[AirbyteStateMessage]] = None, 146 *, 147 source_config: ConnectionDefinition, 148 debug: bool = False, 149 emit_connector_builder_messages: bool = False, 150 migrate_manifest: bool = False, 151 normalize_manifest: bool = False, 152 limits: Optional[TestLimits] = None, 153 config_path: Optional[str] = None, 154 custom_components_trusted: bool = True, 155 **kwargs: Any, 156 ) -> None: 157 self.logger = logging.getLogger(f"airbyte.{self.name}") 158 159 self._limits = limits 160 161 # todo: We could remove state from initialization. Now that streams are grouped during the read(), a source 162 # no longer needs to store the original incoming state. But maybe there's an edge case? 163 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 164 165 # We set a maxsize to for the main thread to process record items when the queue size grows. This assumes that there are less 166 # threads generating partitions that than are max number of workers. If it weren't the case, we could have threads only generating 167 # partitions which would fill the queue. This number is arbitrarily set to 10_000 but will probably need to be changed given more 168 # information and might even need to be configurable depending on the source 169 queue: Queue[QueueItem] = Queue(maxsize=10_000) 170 message_repository = InMemoryMessageRepository( 171 Level.DEBUG if emit_connector_builder_messages else Level.INFO 172 ) 173 174 # To reduce the complexity of the concurrent framework, we are not enabling RFR with synthetic 175 # cursors. We do this by no longer automatically instantiating RFR cursors when converting 176 # the declarative models into runtime components. Concurrent sources will continue to checkpoint 177 # incremental streams running in full refresh. 178 component_factory = ModelToComponentFactory( 179 custom_components_trusted=custom_components_trusted, 180 emit_connector_builder_messages=emit_connector_builder_messages, 181 message_repository=ConcurrentMessageRepository(queue, message_repository), 182 configured_catalog=catalog, 183 connector_state_manager=self._connector_state_manager, 184 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), 185 limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None, 186 limit_slices_fetched=limits.max_slices if limits else None, 187 disable_retries=True if limits else False, 188 disable_cache=True if limits else False, 189 ) 190 191 self._should_normalize = normalize_manifest 192 self._should_migrate = migrate_manifest 193 self._declarative_component_schema = _get_declarative_component_schema() 194 # If custom components are needed, locate and/or register them. 195 self.components_module: ModuleType | None = get_registered_components_module(config=config) 196 # set additional attributes 197 self._debug = debug 198 self._emit_connector_builder_messages = emit_connector_builder_messages 199 self._constructor = ( 200 component_factory 201 if component_factory 202 else ModelToComponentFactory( 203 custom_components_trusted=custom_components_trusted, 204 emit_connector_builder_messages=emit_connector_builder_messages, 205 max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"), 206 ) 207 ) 208 209 self._message_repository = self._constructor.get_message_repository() 210 self._slice_logger: SliceLogger = ( 211 AlwaysLogSliceLogger() if emit_connector_builder_messages else DebugSliceLogger() 212 ) 213 214 # resolve all components in the manifest 215 self._source_config = self._pre_process_manifest(dict(source_config)) 216 # validate resolved manifest against the declarative component schema 217 self._validate_source() 218 # apply additional post-processing to the manifest 219 self._post_process_manifest() 220 221 spec: Optional[Mapping[str, Any]] = self._source_config.get("spec") 222 self._spec_component: Optional[Spec] = ( 223 self._constructor.create_component(SpecModel, spec, dict()) if spec else None 224 ) 225 self._config = self._migrate_and_transform_config(config_path, config) or {} 226 227 concurrency_level_from_manifest = self._source_config.get("concurrency_level") 228 if concurrency_level_from_manifest: 229 concurrency_level_component = self._constructor.create_component( 230 model_type=ConcurrencyLevelModel, 231 component_definition=concurrency_level_from_manifest, 232 config=config or {}, 233 ) 234 if not isinstance(concurrency_level_component, ConcurrencyLevel): 235 raise ValueError( 236 f"Expected to generate a ConcurrencyLevel component, but received {concurrency_level_component.__class__}" 237 ) 238 239 concurrency_level = concurrency_level_component.get_concurrency_level() 240 initial_number_of_partitions_to_generate = max( 241 concurrency_level // 2, 1 242 ) # 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 243 else: 244 concurrency_level = self._LOWEST_SAFE_CONCURRENCY_LEVEL 245 initial_number_of_partitions_to_generate = self._LOWEST_SAFE_CONCURRENCY_LEVEL // 2 246 247 self._concurrent_source = ConcurrentSource.create( 248 num_workers=concurrency_level, 249 initial_number_of_partitions_to_generate=initial_number_of_partitions_to_generate, 250 logger=self.logger, 251 slice_logger=self._slice_logger, 252 queue=queue, 253 message_repository=self._message_repository, 254 )
resolved_manifest: Mapping[str, Any]
361 @property 362 def resolved_manifest(self) -> Mapping[str, Any]: 363 """ 364 Returns the resolved manifest configuration for the source. 365 366 This property provides access to the internal source configuration as a mapping, 367 which contains all settings and parameters required to define the source's behavior. 368 369 Returns: 370 Mapping[str, Any]: The resolved source configuration manifest. 371 """ 372 return self._source_config
Returns the resolved manifest configuration for the source.
This property provides access to the internal source configuration as a mapping, which contains all settings and parameters required to define the source's behavior.
Returns:
Mapping[str, Any]: The resolved source configuration manifest.
def
read( self, logger: logging.Logger, config: Mapping[str, Any], catalog: airbyte_protocol_dataclasses.models.airbyte_protocol.ConfiguredAirbyteCatalog, state: Optional[List[airbyte_cdk.models.airbyte_protocol.AirbyteStateMessage]] = None) -> Iterator[airbyte_cdk.AirbyteMessage]:
377 def read( 378 self, 379 logger: logging.Logger, 380 config: Mapping[str, Any], 381 catalog: ConfiguredAirbyteCatalog, 382 state: Optional[List[AirbyteStateMessage]] = None, 383 ) -> Iterator[AirbyteMessage]: 384 selected_concurrent_streams = self._select_streams( 385 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 386 configured_catalog=catalog, 387 ) 388 389 # It would appear that passing in an empty set of streams causes an infinite loop in ConcurrentReadProcessor. 390 # This is also evident in concurrent_source_adapter.py so I'll leave this out of scope to fix for now 391 if len(selected_concurrent_streams) > 0: 392 yield from self._concurrent_source.read(selected_concurrent_streams)
Returns a generator of the AirbyteMessages generated by reading the source with the given configuration, catalog, and state.
def
discover( self, logger: logging.Logger, config: Mapping[str, Any]) -> airbyte_protocol_dataclasses.models.airbyte_protocol.AirbyteCatalog:
394 def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog: 395 return AirbyteCatalog( 396 streams=[stream.as_airbyte_stream() for stream in self.streams(config=self._config)] 397 )
Returns an AirbyteCatalog representing the available streams and fields in this integration. For example, given valid credentials to a Postgres database, returns an Airbyte catalog where each postgres table is a stream, and each table column is a field.
def
streams( self, config: Mapping[str, Any]) -> List[airbyte_cdk.sources.streams.concurrent.abstract_stream.AbstractStream]:
400 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 401 """ 402 The `streams` method is used as part of the AbstractSource in the following cases: 403 * ConcurrentDeclarativeSource.check -> ManifestDeclarativeSource.check -> AbstractSource.check -> DeclarativeSource.check_connection -> CheckStream.check_connection -> streams 404 * 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`) 405 Note that `super.streams(config)` is also called when splitting the streams between concurrent or not in `_group_streams`. 406 407 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. 408 """ 409 410 if self._spec_component: 411 self._spec_component.validate_config(self._config) 412 413 api_budget_model = self._source_config.get("api_budget") 414 if api_budget_model: 415 self._constructor.set_api_budget(api_budget_model, self._config) 416 417 stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams 418 419 prepared_configs = self._initialize_cache_for_parent_streams(deepcopy(stream_configs)) 420 421 source_streams = [ 422 self._constructor.create_component( 423 ( 424 StateDelegatingStreamModel 425 if stream_config.get("type") == StateDelegatingStreamModel.__name__ 426 else DeclarativeStreamModel 427 ), 428 stream_config, 429 self._config, 430 emit_connector_builder_messages=self._emit_connector_builder_messages, 431 ) 432 for stream_config in prepared_configs 433 ] 434 435 self._apply_stream_groups(source_streams) 436 437 return source_streams
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
- 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) Note thatsuper.streams(config)is also called when splitting the streams between concurrent or not in_group_streams.
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.
561 def spec(self, logger: logging.Logger) -> ConnectorSpecification: 562 """ 563 Returns the connector specification (spec) as defined in the Airbyte Protocol. The spec is an object describing the possible 564 configurations (e.g: username and password) which can be configured when running this connector. For low-code connectors, this 565 will first attempt to load the spec from the manifest's spec block, otherwise it will load it from "spec.yaml" or "spec.json" 566 in the project root. 567 """ 568 return ( 569 self._spec_component.generate_spec() if self._spec_component else super().spec(logger) 570 )
Returns the connector specification (spec) as defined in the Airbyte Protocol. The spec is an object describing the possible configurations (e.g: username and password) which can be configured when running this connector. For low-code connectors, this will first attempt to load the spec from the manifest's spec block, otherwise it will load it from "spec.yaml" or "spec.json" in the project root.
def
check( self, logger: logging.Logger, config: Mapping[str, Any]) -> airbyte_protocol_dataclasses.models.airbyte_protocol.AirbyteConnectionStatus:
572 def check(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteConnectionStatus: 573 check = self._source_config.get("check") 574 if not check: 575 raise ValueError(f"Missing 'check' component definition within the manifest.") 576 577 if "type" not in check: 578 check["type"] = "CheckStream" 579 connection_checker = self._constructor.create_component( 580 COMPONENTS_CHECKER_TYPE_MAPPING[check["type"]], 581 check, 582 dict(), 583 emit_connector_builder_messages=self._emit_connector_builder_messages, 584 ) 585 if not isinstance(connection_checker, ConnectionChecker): 586 raise ValueError( 587 f"Expected to generate a ConnectionChecker component, but received {connection_checker.__class__}" 588 ) 589 590 check_succeeded, error = connection_checker.check_connection(self, logger, self._config) 591 if not check_succeeded: 592 return AirbyteConnectionStatus(status=Status.FAILED, message=repr(error)) 593 return AirbyteConnectionStatus(status=Status.SUCCEEDED)
Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect to the Stripe API.