bb5741a0c0
* New connector_builder module for handling requests from the Connector Builder. Also implements `resolve_manifest` handler * Automated Commit - Formatting Changes * Rename ConnectorBuilderSource to ConnectorBuilderHandler * Update source_declarative_manifest README * Reorganize * read records * paste unit tests from connector builder server * compiles but tests fail * first test passes * Second test passes * 3rd test passes * one more test * another test * one more test * test * return StreamRead * test * test * rename * test * test * test * main seems to work * Update * Update * Update * Update * update * error message * rename * update * Update * CR improvements * fix test_source_declarative_manifest * fix tests * Update * Update * Update * Update * rename * rename * rename * format * Give connector_builder its own main.py * Update * reset * delete dead code * remove debug print * update test * Update * set right stream * Add --catalog argument * Remove unneeded preparse * Update README * handle error * tests pass * more explicit test * reset * format * fix merge * raise exception * fix * black format * raise with config * update * fix flake * __test_read_config is optional * fix * Automated Commit - Formatting Changes * fix * exclude_unset --------- Co-authored-by: Catherine Noll <noll.catherine@gmail.com> Co-authored-by: clnoll <clnoll@users.noreply.github.com> Co-authored-by: girarda <girarda@users.noreply.github.com>
63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
#
|
|
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
#
|
|
|
|
import dataclasses
|
|
from datetime import datetime
|
|
from typing import Any, Mapping
|
|
|
|
from airbyte_cdk.models import AirbyteMessage, AirbyteRecordMessage, ConfiguredAirbyteCatalog
|
|
from airbyte_cdk.models import Type
|
|
from airbyte_cdk.models import Type as MessageType
|
|
from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
|
|
from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
|
|
from airbyte_cdk.utils.traced_exception import AirbyteTracedException
|
|
from connector_builder.message_grouper import MessageGrouper
|
|
|
|
|
|
def list_streams() -> AirbyteMessage:
|
|
raise NotImplementedError
|
|
|
|
|
|
DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE = 5
|
|
DEFAULT_MAXIMUM_NUMBER_OF_SLICES = 5
|
|
DEFAULT_MAX_RECORDS = 100
|
|
|
|
|
|
def read_stream(source: DeclarativeSource, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog) -> AirbyteMessage:
|
|
try:
|
|
command_config = config.get("__test_read_config", {})
|
|
max_pages_per_slice = command_config.get("max_pages_per_slice", DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
|
|
max_slices = command_config.get("max_slices", DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
|
|
max_records = command_config.get("max_records", DEFAULT_MAX_RECORDS)
|
|
handler = MessageGrouper(max_pages_per_slice, max_slices)
|
|
stream_name = configured_catalog.streams[0].stream.name # The connector builder only supports a single stream
|
|
stream_read = handler.get_message_groups(source, config, configured_catalog, max_records)
|
|
return AirbyteMessage(type=MessageType.RECORD, record=AirbyteRecordMessage(
|
|
data=dataclasses.asdict(stream_read),
|
|
stream=stream_name,
|
|
emitted_at=_emitted_at()
|
|
))
|
|
except Exception as exc:
|
|
error = AirbyteTracedException.from_exception(exc, message=f"Error reading stream with config={config} and catalog={configured_catalog}")
|
|
return error.as_airbyte_message()
|
|
|
|
|
|
def resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
|
|
try:
|
|
return AirbyteMessage(
|
|
type=Type.RECORD,
|
|
record=AirbyteRecordMessage(
|
|
data={"manifest": source.resolved_manifest},
|
|
emitted_at=_emitted_at(),
|
|
stream="resolve_manifest",
|
|
),
|
|
)
|
|
except Exception as exc:
|
|
error = AirbyteTracedException.from_exception(exc, message="Error resolving manifest.")
|
|
return error.as_airbyte_message()
|
|
|
|
|
|
def _emitted_at():
|
|
return int(datetime.now().timestamp()) * 1000
|