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,
112            config,
113            configured_catalog,
114            stream_name,
115            state,
116            limits.max_records,
117        )
118
119        return AirbyteMessage(
120            type=MessageType.RECORD,
121            record=AirbyteRecordMessage(
122                data=asdict(stream_read), stream=stream_name, emitted_at=_emitted_at()
123            ),
124        )
125    except Exception as exc:
126        error = AirbyteTracedException.from_exception(
127            exc,
128            message=filter_secrets(
129                f"Error reading stream with config={config} and catalog={configured_catalog}: {str(exc)}"
130            ),
131        )
132        return error.as_airbyte_message()
133
134
135def resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
136    try:
137        return AirbyteMessage(
138            type=Type.RECORD,
139            record=AirbyteRecordMessage(
140                data={"manifest": source.resolved_manifest},
141                emitted_at=_emitted_at(),
142                stream="resolve_manifest",
143            ),
144        )
145    except Exception as exc:
146        error = AirbyteTracedException.from_exception(
147            exc, message=f"Error resolving manifest: {str(exc)}"
148        )
149        return error.as_airbyte_message()
150
151
152def full_resolve_manifest(source: ManifestDeclarativeSource, limits: TestLimits) -> AirbyteMessage:
153    try:
154        manifest = {**source.resolved_manifest}
155        streams = manifest.get("streams", [])
156        for stream in streams:
157            stream["dynamic_stream_name"] = None
158
159        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
160        for stream in source.dynamic_streams:
161            generated_streams = mapped_streams.setdefault(stream["dynamic_stream_name"], [])
162
163            if len(generated_streams) < limits.max_streams:
164                generated_streams += [stream]
165
166        for generated_streams_list in mapped_streams.values():
167            streams.extend(generated_streams_list)
168
169        manifest["streams"] = streams
170        return AirbyteMessage(
171            type=Type.RECORD,
172            record=AirbyteRecordMessage(
173                data={"manifest": manifest},
174                emitted_at=_emitted_at(),
175                stream="full_resolve_manifest",
176            ),
177        )
178    except AirbyteTracedException as exc:
179        return exc.as_airbyte_message()
180    except Exception as exc:
181        error = AirbyteTracedException.from_exception(
182            exc, message=f"Error full resolving manifest: {str(exc)}"
183        )
184        return error.as_airbyte_message()
185
186
187def _emitted_at() -> int:
188    return ab_datetime_now().to_epoch_millis()