51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
#
|
|
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
#
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Generic, Iterable, Optional, TypeVar
|
|
|
|
from airbyte_cdk.connector import TConfig
|
|
from airbyte_cdk.sources.embedded.catalog import create_configured_catalog, get_stream, get_stream_names
|
|
from airbyte_cdk.sources.embedded.runner import SourceRunner
|
|
from airbyte_cdk.sources.embedded.tools import get_defined_id
|
|
from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit
|
|
from airbyte_protocol.models import AirbyteRecordMessage, AirbyteStateMessage, SyncMode, Type
|
|
|
|
TOutput = TypeVar("TOutput")
|
|
|
|
|
|
class BaseEmbeddedIntegration(ABC, Generic[TConfig, TOutput]):
|
|
def __init__(self, runner: SourceRunner[TConfig], config: TConfig):
|
|
check_config_against_spec_or_exit(config, runner.spec())
|
|
|
|
self.source = runner
|
|
self.config = config
|
|
|
|
self.last_state: Optional[AirbyteStateMessage] = None
|
|
|
|
@abstractmethod
|
|
def _handle_record(self, record: AirbyteRecordMessage, id: Optional[str]) -> Optional[TOutput]:
|
|
"""
|
|
Turn an Airbyte record into the appropriate output type for the integration.
|
|
"""
|
|
pass
|
|
|
|
def _load_data(self, stream_name: str, state: Optional[AirbyteStateMessage] = None) -> Iterable[TOutput]:
|
|
catalog = self.source.discover(self.config)
|
|
stream = get_stream(catalog, stream_name)
|
|
if not stream:
|
|
raise ValueError(f"Stream {stream_name} not found, the following streams are available: {', '.join(get_stream_names(catalog))}")
|
|
if SyncMode.incremental not in stream.supported_sync_modes:
|
|
configured_catalog = create_configured_catalog(stream, sync_mode=SyncMode.full_refresh)
|
|
else:
|
|
configured_catalog = create_configured_catalog(stream, sync_mode=SyncMode.incremental)
|
|
|
|
for message in self.source.read(self.config, configured_catalog, state):
|
|
if message.type == Type.RECORD:
|
|
output = self._handle_record(message.record, get_defined_id(stream, message.record.data))
|
|
if output:
|
|
yield output
|
|
elif message.type is Type.STATE and message.state:
|
|
self.last_state = message.state
|