airbyte_cdk.connector_builder.connector_builder_handler

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5
  6from dataclasses import asdict
  7from typing import Any, Dict, List, Mapping, Optional
  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.concurrent_declarative_source import (
 19    ConcurrentDeclarativeSource,
 20    TestLimits,
 21)
 22from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 23from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 24from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 25
 26MAX_PAGES_PER_SLICE_KEY = "max_pages_per_slice"
 27MAX_SLICES_KEY = "max_slices"
 28MAX_RECORDS_KEY = "max_records"
 29MAX_STREAMS_KEY = "max_streams"
 30
 31
 32def get_limits(config: Mapping[str, Any]) -> TestLimits:
 33    command_config = config.get("__test_read_config", {})
 34    return TestLimits(
 35        max_records=command_config.get(MAX_RECORDS_KEY, TestLimits.DEFAULT_MAX_RECORDS),
 36        max_pages_per_slice=command_config.get(
 37            MAX_PAGES_PER_SLICE_KEY, TestLimits.DEFAULT_MAX_PAGES_PER_SLICE
 38        ),
 39        max_slices=command_config.get(MAX_SLICES_KEY, TestLimits.DEFAULT_MAX_SLICES),
 40        max_streams=command_config.get(MAX_STREAMS_KEY, TestLimits.DEFAULT_MAX_STREAMS),
 41    )
 42
 43
 44def should_migrate_manifest(config: Mapping[str, Any]) -> bool:
 45    """
 46    Determines whether the manifest should be migrated,
 47    based on the presence of the "__should_migrate" key in the config.
 48
 49    This flag is set by the UI.
 50    """
 51    return config.get("__should_migrate", False)
 52
 53
 54def should_normalize_manifest(config: Mapping[str, Any]) -> bool:
 55    """
 56    Check if the manifest should be normalized.
 57    :param config: The configuration to check
 58    :return: True if the manifest should be normalized, False otherwise.
 59    """
 60    return config.get("__should_normalize", False)
 61
 62
 63def create_source(
 64    config: Mapping[str, Any],
 65    limits: TestLimits | None = None,
 66    catalog: ConfiguredAirbyteCatalog | None = None,
 67    state: List[AirbyteStateMessage] | None = None,
 68) -> ConcurrentDeclarativeSource:
 69    manifest = config["__injected_declarative_manifest"]
 70
 71    # We enforce a concurrency level of 1 so that the stream is processed on a single thread
 72    # to retain ordering for the grouping of the builder message responses.
 73    if "concurrency_level" in manifest:
 74        manifest["concurrency_level"]["default_concurrency"] = 1
 75    else:
 76        manifest["concurrency_level"] = {"type": "ConcurrencyLevel", "default_concurrency": 1}
 77
 78    return ConcurrentDeclarativeSource(
 79        catalog=catalog,
 80        config=config,
 81        state=state,
 82        source_config=manifest,
 83        emit_connector_builder_messages=True,
 84        migrate_manifest=should_migrate_manifest(config),
 85        normalize_manifest=should_normalize_manifest(config),
 86        limits=limits,
 87        # The manifest is supplied by the Connector Builder caller rather than bundled in a
 88        # connector image, so any custom components it references are untrusted code.
 89        custom_components_trusted=False,
 90    )
 91
 92
 93def read_stream(
 94    source: ConcurrentDeclarativeSource,
 95    config: Mapping[str, Any],
 96    configured_catalog: ConfiguredAirbyteCatalog,
 97    state: List[AirbyteStateMessage],
 98    limits: TestLimits,
 99) -> AirbyteMessage:
100    try:
101        test_read_handler = TestReader(
102            limits.max_pages_per_slice, limits.max_slices, limits.max_records
103        )
104        # The connector builder only supports a single stream
105        stream_name = configured_catalog.streams[0].stream.name
106
107        stream_read = test_read_handler.run_test_read(
108            source,
109            config,
110            configured_catalog,
111            stream_name,
112            state,
113            limits.max_records,
114        )
115
116        return AirbyteMessage(
117            type=MessageType.RECORD,
118            record=AirbyteRecordMessage(
119                data=asdict(stream_read), stream=stream_name, emitted_at=_emitted_at()
120            ),
121        )
122    except Exception as exc:
123        error = AirbyteTracedException.from_exception(
124            exc,
125            message=filter_secrets(
126                f"Error reading stream with config={config} and catalog={configured_catalog}: {str(exc)}"
127            ),
128        )
129        return error.as_airbyte_message()
130
131
132def resolve_manifest(
133    source: ConcurrentDeclarativeSource,
134) -> AirbyteMessage:
135    try:
136        return AirbyteMessage(
137            type=Type.RECORD,
138            record=AirbyteRecordMessage(
139                data={"manifest": source.resolved_manifest},
140                emitted_at=_emitted_at(),
141                stream="resolve_manifest",
142            ),
143        )
144    except Exception as exc:
145        error = AirbyteTracedException.from_exception(
146            exc, message=f"Error resolving manifest: {str(exc)}"
147        )
148        return error.as_airbyte_message()
149
150
151def full_resolve_manifest(
152    source: ConcurrentDeclarativeSource, limits: TestLimits
153) -> AirbyteMessage:
154    try:
155        manifest = {**source.resolved_manifest}
156        streams = manifest.get("streams", [])
157        for stream in streams:
158            stream["dynamic_stream_name"] = None
159
160        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
161        for stream in source.dynamic_streams:
162            generated_streams = mapped_streams.setdefault(stream["dynamic_stream_name"], [])
163
164            if len(generated_streams) < limits.max_streams:
165                generated_streams += [stream]
166
167        for generated_streams_list in mapped_streams.values():
168            streams.extend(generated_streams_list)
169
170        manifest["streams"] = streams
171        return AirbyteMessage(
172            type=Type.RECORD,
173            record=AirbyteRecordMessage(
174                data={"manifest": manifest},
175                emitted_at=_emitted_at(),
176                stream="full_resolve_manifest",
177            ),
178        )
179    except AirbyteTracedException as exc:
180        return exc.as_airbyte_message()
181    except Exception as exc:
182        error = AirbyteTracedException.from_exception(
183            exc, message=f"Error full resolving manifest: {str(exc)}"
184        )
185        return error.as_airbyte_message()
186
187
188def _emitted_at() -> int:
189    return ab_datetime_now().to_epoch_millis()