airbyte_cdk.connector_builder.connector_builder_handler

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5
  6from dataclasses import asdict, dataclass, field
  7from typing import Any, ClassVar, Dict, List, Mapping
  8
  9from airbyte_cdk.connector_builder.test_reader import TestReader
 10from airbyte_cdk.models import (
 11    AirbyteMessage,
 12    AirbyteRecordMessage,
 13    AirbyteStateMessage,
 14    ConfiguredAirbyteCatalog,
 15    Type,
 16)
 17from airbyte_cdk.models import Type as MessageType
 18from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
 19from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 20from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
 21    ModelToComponentFactory,
 22)
 23from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 24from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 25from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 26
 27DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE = 5
 28DEFAULT_MAXIMUM_NUMBER_OF_SLICES = 5
 29DEFAULT_MAXIMUM_RECORDS = 100
 30DEFAULT_MAXIMUM_STREAMS = 100
 31
 32MAX_PAGES_PER_SLICE_KEY = "max_pages_per_slice"
 33MAX_SLICES_KEY = "max_slices"
 34MAX_RECORDS_KEY = "max_records"
 35MAX_STREAMS_KEY = "max_streams"
 36
 37
 38@dataclass
 39class TestLimits:
 40    __test__: ClassVar[bool] = False  # Tell Pytest this is not a Pytest class, despite its name
 41
 42    max_records: int = field(default=DEFAULT_MAXIMUM_RECORDS)
 43    max_pages_per_slice: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
 44    max_slices: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
 45    max_streams: int = field(default=DEFAULT_MAXIMUM_STREAMS)
 46
 47
 48def get_limits(config: Mapping[str, Any]) -> TestLimits:
 49    command_config = config.get("__test_read_config", {})
 50    max_pages_per_slice = (
 51        command_config.get(MAX_PAGES_PER_SLICE_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE
 52    )
 53    max_slices = command_config.get(MAX_SLICES_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_SLICES
 54    max_records = command_config.get(MAX_RECORDS_KEY) or DEFAULT_MAXIMUM_RECORDS
 55    max_streams = command_config.get(MAX_STREAMS_KEY) or DEFAULT_MAXIMUM_STREAMS
 56    return TestLimits(max_records, max_pages_per_slice, max_slices, max_streams)
 57
 58
 59def should_migrate_manifest(config: Mapping[str, Any]) -> bool:
 60    """
 61    Determines whether the manifest should be migrated,
 62    based on the presence of the "__should_migrate" key in the config.
 63
 64    This flag is set by the UI.
 65    """
 66    return config.get("__should_migrate", False)
 67
 68
 69def should_normalize_manifest(config: Mapping[str, Any]) -> bool:
 70    """
 71    Check if the manifest should be normalized.
 72    :param config: The configuration to check
 73    :return: True if the manifest should be normalized, False otherwise.
 74    """
 75    return config.get("__should_normalize", False)
 76
 77
 78def create_source(config: Mapping[str, Any], limits: TestLimits) -> ManifestDeclarativeSource:
 79    manifest = config["__injected_declarative_manifest"]
 80    return ManifestDeclarativeSource(
 81        config=config,
 82        emit_connector_builder_messages=True,
 83        source_config=manifest,
 84        migrate_manifest=should_migrate_manifest(config),
 85        normalize_manifest=should_normalize_manifest(config),
 86        component_factory=ModelToComponentFactory(
 87            emit_connector_builder_messages=True,
 88            limit_pages_fetched_per_slice=limits.max_pages_per_slice,
 89            limit_slices_fetched=limits.max_slices,
 90            disable_retries=True,
 91            disable_cache=True,
 92        ),
 93    )
 94
 95
 96def read_stream(
 97    source: DeclarativeSource,
 98    config: Mapping[str, Any],
 99    configured_catalog: ConfiguredAirbyteCatalog,
100    state: List[AirbyteStateMessage],
101    limits: TestLimits,
102) -> AirbyteMessage:
103    try:
104        test_read_handler = TestReader(
105            limits.max_pages_per_slice, limits.max_slices, limits.max_records
106        )
107        # The connector builder only supports a single stream
108        stream_name = configured_catalog.streams[0].stream.name
109
110        stream_read = test_read_handler.run_test_read(
111            source, config, configured_catalog, state, limits.max_records
112        )
113
114        return AirbyteMessage(
115            type=MessageType.RECORD,
116            record=AirbyteRecordMessage(
117                data=asdict(stream_read), stream=stream_name, emitted_at=_emitted_at()
118            ),
119        )
120    except Exception as exc:
121        error = AirbyteTracedException.from_exception(
122            exc,
123            message=filter_secrets(
124                f"Error reading stream with config={config} and catalog={configured_catalog}: {str(exc)}"
125            ),
126        )
127        return error.as_airbyte_message()
128
129
130def resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
131    try:
132        return AirbyteMessage(
133            type=Type.RECORD,
134            record=AirbyteRecordMessage(
135                data={"manifest": source.resolved_manifest},
136                emitted_at=_emitted_at(),
137                stream="resolve_manifest",
138            ),
139        )
140    except Exception as exc:
141        error = AirbyteTracedException.from_exception(
142            exc, message=f"Error resolving manifest: {str(exc)}"
143        )
144        return error.as_airbyte_message()
145
146
147def full_resolve_manifest(source: ManifestDeclarativeSource, limits: TestLimits) -> AirbyteMessage:
148    try:
149        manifest = {**source.resolved_manifest}
150        streams = manifest.get("streams", [])
151        for stream in streams:
152            stream["dynamic_stream_name"] = None
153
154        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
155        for stream in source.dynamic_streams:
156            generated_streams = mapped_streams.setdefault(stream["dynamic_stream_name"], [])
157
158            if len(generated_streams) < limits.max_streams:
159                generated_streams += [stream]
160
161        for generated_streams_list in mapped_streams.values():
162            streams.extend(generated_streams_list)
163
164        manifest["streams"] = streams
165        return AirbyteMessage(
166            type=Type.RECORD,
167            record=AirbyteRecordMessage(
168                data={"manifest": manifest},
169                emitted_at=_emitted_at(),
170                stream="full_resolve_manifest",
171            ),
172        )
173    except AirbyteTracedException as exc:
174        return exc.as_airbyte_message()
175    except Exception as exc:
176        error = AirbyteTracedException.from_exception(
177            exc, message=f"Error full resolving manifest: {str(exc)}"
178        )
179        return error.as_airbyte_message()
180
181
182def _emitted_at() -> int:
183    return ab_datetime_now().to_epoch_millis()