airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/destinations/vector_db_based/utils.py

@@ -0,0 +1,35 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import itertools
+import traceback
+from typing import Any, Iterable, Iterator, Tuple, Union
+
+from airbyte_cdk.models import AirbyteRecordMessage, AirbyteStream
+
+
+def format_exception(exception: Exception) -> str:
+    return (
+        str(exception)
+        + "\n"
+        + "".join(traceback.TracebackException.from_exception(exception).format())
+    )
+
+
+def create_chunks(iterable: Iterable[Any], batch_size: int) -> Iterator[Tuple[Any, ...]]:
+    """A helper function to break an iterable into chunks of size batch_size."""
+    it = iter(iterable)
+    chunk = tuple(itertools.islice(it, batch_size))
+    while chunk:
+        yield chunk
+        chunk = tuple(itertools.islice(it, batch_size))
+
+
+def create_stream_identifier(stream: Union[AirbyteStream, AirbyteRecordMessage]) -> str:
+    if isinstance(stream, AirbyteStream):
+        return str(stream.name if stream.namespace is None else f"{stream.namespace}_{stream.name}")
+    else:
+        return str(
+            stream.stream if stream.namespace is None else f"{stream.namespace}_{stream.stream}"
+        )

airbyte_cdk/destinations/vector_db_based/writer.py

@@ -0,0 +1,104 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+from collections import defaultdict
+from typing import Dict, Iterable, List, Tuple
+
+from airbyte_cdk.destinations.vector_db_based.config import ProcessingConfigModel
+from airbyte_cdk.destinations.vector_db_based.document_processor import Chunk, DocumentProcessor
+from airbyte_cdk.destinations.vector_db_based.embedder import Document, Embedder
+from airbyte_cdk.destinations.vector_db_based.indexer import Indexer
+from airbyte_cdk.models import AirbyteMessage, ConfiguredAirbyteCatalog, Type
+
+
+class Writer:
+    """
+    The Writer class is orchestrating the document processor, the embedder and the indexer:
+    * Incoming records are passed through the document processor to generate chunks
+    * One the configured batch size is reached, the chunks are passed to the embedder to generate embeddings
+    * The embedder embeds the chunks
+    * The indexer deletes old chunks by the associated record id before indexing the new ones
+
+    The destination connector is responsible to create a writer instance and pass the input messages iterable to the write method.
+    The batch size can be configured by the destination connector to give the freedom of either letting the user configure it or hardcoding it to a sensible value depending on the destination.
+    The omit_raw_text parameter can be used to omit the raw text from the chunks. This can be useful if the raw text is very large and not needed for the destination.
+    """
+
+    def __init__(
+        self,
+        processing_config: ProcessingConfigModel,
+        indexer: Indexer,
+        embedder: Embedder,
+        batch_size: int,
+        omit_raw_text: bool,
+    ) -> None:
+        self.processing_config = processing_config
+        self.indexer = indexer
+        self.embedder = embedder
+        self.batch_size = batch_size
+        self.omit_raw_text = omit_raw_text
+        self._init_batch()
+
+    def _init_batch(self) -> None:
+        self.chunks: Dict[Tuple[str, str], List[Chunk]] = defaultdict(list)
+        self.ids_to_delete: Dict[Tuple[str, str], List[str]] = defaultdict(list)
+        self.number_of_chunks = 0
+
+    def _convert_to_document(self, chunk: Chunk) -> Document:
+        """
+        Convert a chunk to a document for the embedder.
+        """
+        if chunk.page_content is None:
+            raise ValueError("Cannot embed a chunk without page content")
+        return Document(page_content=chunk.page_content, record=chunk.record)
+
+    def _process_batch(self) -> None:
+        for (namespace, stream), ids in self.ids_to_delete.items():
+            self.indexer.delete(ids, namespace, stream)
+
+        for (namespace, stream), chunks in self.chunks.items():
+            embeddings = self.embedder.embed_documents(
+                [self._convert_to_document(chunk) for chunk in chunks]
+            )
+            for i, document in enumerate(chunks):
+                document.embedding = embeddings[i]
+                if self.omit_raw_text:
+                    document.page_content = None
+            self.indexer.index(chunks, namespace, stream)
+
+        self._init_batch()
+
+    def write(
+        self, configured_catalog: ConfiguredAirbyteCatalog, input_messages: Iterable[AirbyteMessage]
+    ) -> Iterable[AirbyteMessage]:
+        self.processor = DocumentProcessor(self.processing_config, configured_catalog)
+        self.indexer.pre_sync(configured_catalog)
+        for message in input_messages:
+            if message.type == Type.STATE:
+                # Emitting a state message indicates that all records which came before it have been written to the destination. So we flush
+                # the queue to ensure writes happen, then output the state message to indicate it's safe to checkpoint state
+                self._process_batch()
+                yield message
+            elif message.type == Type.RECORD:
+                record_chunks, record_id_to_delete = self.processor.process(message.record)
+                self.chunks[
+                    (  # type: ignore [index] # expected "tuple[str, str]", got "tuple[str | Any | None, str | Any]"
+                        message.record.namespace,  # type: ignore [union-attr] # record not None
+                        message.record.stream,  # type: ignore [union-attr] # record not None
+                    )
+                ].extend(record_chunks)
+                if record_id_to_delete is not None:
+                    self.ids_to_delete[
+                        (  # type: ignore [index] # expected "tuple[str, str]", got "tuple[str | Any | None, str | Any]"
+                            message.record.namespace,  # type: ignore [union-attr] # record not None
+                            message.record.stream,  # type: ignore [union-attr] # record not None
+                        )
+                    ].append(record_id_to_delete)
+                self.number_of_chunks += len(record_chunks)
+                if self.number_of_chunks >= self.batch_size:
+                    self._process_batch()
+
+        self._process_batch()
+        yield from self.indexer.post_sync()

airbyte_cdk/entrypoint.py

@@ -0,0 +1,414 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import argparse
+import importlib
+import ipaddress
+import json
+import logging
+import os.path
+import socket
+import sys
+import tempfile
+from collections import defaultdict
+from functools import wraps
+from typing import Any, DefaultDict, Iterable, List, Mapping, Optional
+from urllib.parse import urlparse
+
+import orjson
+import requests
+from requests import PreparedRequest, Response, Session
+
+from airbyte_cdk.connector import TConfig
+from airbyte_cdk.exception_handler import init_uncaught_exception_handler
+from airbyte_cdk.logger import PRINT_BUFFER, init_logger
+from airbyte_cdk.models import (
+    AirbyteConnectionStatus,
+    AirbyteMessage,
+    AirbyteMessageSerializer,
+    AirbyteStateStats,
+    ConnectorSpecification,
+    FailureType,
+    Status,
+    Type,
+)
+from airbyte_cdk.sources import Source
+from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
+from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit, split_config
+
+# from airbyte_cdk.utils import PrintBuffer, is_cloud_environment, message_utils  # add PrintBuffer back once fixed
+from airbyte_cdk.utils import is_cloud_environment, message_utils
+from airbyte_cdk.utils.airbyte_secrets_utils import get_secrets, update_secrets
+from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+logger = init_logger("airbyte")
+
+VALID_URL_SCHEMES = ["https"]
+CLOUD_DEPLOYMENT_MODE = "cloud"
+_HAS_LOGGED_FOR_SERIALIZATION_ERROR = False
+
+
+class AirbyteEntrypoint(object):
+    def __init__(self, source: Source):
+        init_uncaught_exception_handler(logger)
+
+        # Deployment mode is read when instantiating the entrypoint because it is the common path shared by syncs and connector builder test requests
+        if is_cloud_environment():
+            _init_internal_request_filter()
+
+        self.source = source
+        self.logger = logging.getLogger(f"airbyte.{getattr(source, 'name', '')}")
+
+    @staticmethod
+    def parse_args(args: List[str]) -> argparse.Namespace:
+        # set up parent parsers
+        parent_parser = argparse.ArgumentParser(add_help=False)
+        parent_parser.add_argument(
+            "--debug", action="store_true", help="enables detailed debug logs related to the sync"
+        )
+        main_parser = argparse.ArgumentParser()
+        subparsers = main_parser.add_subparsers(title="commands", dest="command")
+
+        # spec
+        subparsers.add_parser(
+            "spec", help="outputs the json configuration specification", parents=[parent_parser]
+        )
+
+        # check
+        check_parser = subparsers.add_parser(
+            "check", help="checks the config can be used to connect", parents=[parent_parser]
+        )
+        required_check_parser = check_parser.add_argument_group("required named arguments")
+        required_check_parser.add_argument(
+            "--config", type=str, required=True, help="path to the json configuration file"
+        )
+
+        # discover
+        discover_parser = subparsers.add_parser(
+            "discover",
+            help="outputs a catalog describing the source's schema",
+            parents=[parent_parser],
+        )
+        required_discover_parser = discover_parser.add_argument_group("required named arguments")
+        required_discover_parser.add_argument(
+            "--config", type=str, required=True, help="path to the json configuration file"
+        )
+
+        # read
+        read_parser = subparsers.add_parser(
+            "read", help="reads the source and outputs messages to STDOUT", parents=[parent_parser]
+        )
+
+        read_parser.add_argument(
+            "--state", type=str, required=False, help="path to the json-encoded state file"
+        )
+        required_read_parser = read_parser.add_argument_group("required named arguments")
+        required_read_parser.add_argument(
+            "--config", type=str, required=True, help="path to the json configuration file"
+        )
+        required_read_parser.add_argument(
+            "--catalog",
+            type=str,
+            required=True,
+            help="path to the catalog used to determine which data to read",
+        )
+
+        return main_parser.parse_args(args)
+
+    def run(self, parsed_args: argparse.Namespace) -> Iterable[str]:
+        cmd = parsed_args.command
+        if not cmd:
+            raise Exception("No command passed")
+
+        if hasattr(parsed_args, "debug") and parsed_args.debug:
+            self.logger.setLevel(logging.DEBUG)
+            logger.setLevel(logging.DEBUG)
+            self.logger.debug("Debug logs enabled")
+        else:
+            self.logger.setLevel(logging.INFO)
+
+        source_spec: ConnectorSpecification = self.source.spec(self.logger)
+        try:
+            with tempfile.TemporaryDirectory(
+                # Cleanup can fail on Windows due to file locks. Ignore if so,
+                # rather than failing the whole process.
+                ignore_cleanup_errors=True,
+            ) as temp_dir:
+                os.environ[ENV_REQUEST_CACHE_PATH] = (
+                    temp_dir  # set this as default directory for request_cache to store *.sqlite files
+                )
+                if cmd == "spec":
+                    message = AirbyteMessage(type=Type.SPEC, spec=source_spec)
+                    yield from [
+                        self.airbyte_message_to_string(queued_message)
+                        for queued_message in self._emit_queued_messages(self.source)
+                    ]
+                    yield self.airbyte_message_to_string(message)
+                else:
+                    raw_config = self.source.read_config(parsed_args.config)
+                    config = self.source.configure(raw_config, temp_dir)
+
+                    yield from [
+                        self.airbyte_message_to_string(queued_message)
+                        for queued_message in self._emit_queued_messages(self.source)
+                    ]
+                    if cmd == "check":
+                        yield from map(
+                            AirbyteEntrypoint.airbyte_message_to_string,
+                            self.check(source_spec, config),
+                        )
+                    elif cmd == "discover":
+                        yield from map(
+                            AirbyteEntrypoint.airbyte_message_to_string,
+                            self.discover(source_spec, config),
+                        )
+                    elif cmd == "read":
+                        config_catalog = self.source.read_catalog(parsed_args.catalog)
+                        state = self.source.read_state(parsed_args.state)
+
+                        yield from map(
+                            AirbyteEntrypoint.airbyte_message_to_string,
+                            self.read(source_spec, config, config_catalog, state),
+                        )
+                    else:
+                        raise Exception("Unexpected command " + cmd)
+        finally:
+            yield from [
+                self.airbyte_message_to_string(queued_message)
+                for queued_message in self._emit_queued_messages(self.source)
+            ]
+
+    def check(
+        self, source_spec: ConnectorSpecification, config: TConfig
+    ) -> Iterable[AirbyteMessage]:
+        self.set_up_secret_filter(config, source_spec.connectionSpecification)
+        try:
+            self.validate_connection(source_spec, config)
+        except AirbyteTracedException as traced_exc:
+            connection_status = traced_exc.as_connection_status_message()
+            # The platform uses the exit code to surface unexpected failures so we raise the exception if the failure type not a config error
+            # If the failure is not exceptional, we'll emit a failed connection status message and return
+            if traced_exc.failure_type != FailureType.config_error:
+                raise traced_exc
+            if connection_status:
+                yield from self._emit_queued_messages(self.source)
+                yield connection_status
+                return
+
+        try:
+            check_result = self.source.check(self.logger, config)
+        except AirbyteTracedException as traced_exc:
+            yield traced_exc.as_airbyte_message()
+            # The platform uses the exit code to surface unexpected failures so we raise the exception if the failure type not a config error
+            # If the failure is not exceptional, we'll emit a failed connection status message and return
+            if traced_exc.failure_type != FailureType.config_error:
+                raise traced_exc
+            else:
+                yield AirbyteMessage(
+                    type=Type.CONNECTION_STATUS,
+                    connectionStatus=AirbyteConnectionStatus(
+                        status=Status.FAILED, message=traced_exc.message
+                    ),
+                )
+                return
+        if check_result.status == Status.SUCCEEDED:
+            self.logger.info("Check succeeded")
+        else:
+            self.logger.error("Check failed")
+
+        yield from self._emit_queued_messages(self.source)
+        yield AirbyteMessage(type=Type.CONNECTION_STATUS, connectionStatus=check_result)
+
+    def discover(
+        self, source_spec: ConnectorSpecification, config: TConfig
+    ) -> Iterable[AirbyteMessage]:
+        self.set_up_secret_filter(config, source_spec.connectionSpecification)
+        if self.source.check_config_against_spec:
+            self.validate_connection(source_spec, config)
+        catalog = self.source.discover(self.logger, config)
+
+        yield from self._emit_queued_messages(self.source)
+        yield AirbyteMessage(type=Type.CATALOG, catalog=catalog)
+
+    def read(
+        self, source_spec: ConnectorSpecification, config: TConfig, catalog: Any, state: list[Any]
+    ) -> Iterable[AirbyteMessage]:
+        self.set_up_secret_filter(config, source_spec.connectionSpecification)
+        if self.source.check_config_against_spec:
+            self.validate_connection(source_spec, config)
+
+        # The Airbyte protocol dictates that counts be expressed as float/double to better protect against integer overflows
+        stream_message_counter: DefaultDict[HashableStreamDescriptor, float] = defaultdict(float)
+        for message in self.source.read(self.logger, config, catalog, state):
+            yield self.handle_record_counts(message, stream_message_counter)
+        for message in self._emit_queued_messages(self.source):
+            yield self.handle_record_counts(message, stream_message_counter)
+
+    @staticmethod
+    def handle_record_counts(
+        message: AirbyteMessage, stream_message_count: DefaultDict[HashableStreamDescriptor, float]
+    ) -> AirbyteMessage:
+        match message.type:
+            case Type.RECORD:
+                if message.record is None:
+                    raise ValueError("Record message must have a record attribute")
+
+                stream_message_count[
+                    HashableStreamDescriptor(
+                        name=message.record.stream,  # type: ignore[union-attr] # record has `stream`
+                        namespace=message.record.namespace,  # type: ignore[union-attr] # record has `namespace`
+                    )
+                ] += 1.0
+            case Type.STATE:
+                if message.state is None:
+                    raise ValueError("State message must have a state attribute")
+
+                stream_descriptor = message_utils.get_stream_descriptor(message)
+
+                # Set record count from the counter onto the state message
+                message.state.sourceStats = message.state.sourceStats or AirbyteStateStats()  # type: ignore[union-attr] # state has `sourceStats`
+                message.state.sourceStats.recordCount = stream_message_count.get(  # type: ignore[union-attr] # state has `sourceStats`
+                    stream_descriptor, 0.0
+                )
+
+                # Reset the counter
+                stream_message_count[stream_descriptor] = 0.0
+        return message
+
+    @staticmethod
+    def validate_connection(source_spec: ConnectorSpecification, config: TConfig) -> None:
+        # Remove internal flags from config before validating so
+        # jsonschema's additionalProperties flag won't fail the validation
+        connector_config, _ = split_config(config)
+        check_config_against_spec_or_exit(connector_config, source_spec)
+
+    @staticmethod
+    def set_up_secret_filter(config: TConfig, connection_specification: Mapping[str, Any]) -> None:
+        # Now that we have the config, we can use it to get a list of ai airbyte_secrets
+        # that we should filter in logging to avoid leaking secrets
+        config_secrets = get_secrets(connection_specification, config)
+        update_secrets(config_secrets)
+
+    @staticmethod
+    def airbyte_message_to_string(airbyte_message: AirbyteMessage) -> str:
+        global _HAS_LOGGED_FOR_SERIALIZATION_ERROR
+        serialized_message = AirbyteMessageSerializer.dump(airbyte_message)
+        try:
+            return orjson.dumps(serialized_message).decode()
+        except Exception as exception:
+            if not _HAS_LOGGED_FOR_SERIALIZATION_ERROR:
+                logger.warning(
+                    f"There was an error during the serialization of an AirbyteMessage: `{exception}`. This might impact the sync performances."
+                )
+                _HAS_LOGGED_FOR_SERIALIZATION_ERROR = True
+            return json.dumps(serialized_message)
+
+    @classmethod
+    def extract_state(cls, args: List[str]) -> Optional[Any]:
+        parsed_args = cls.parse_args(args)
+        if hasattr(parsed_args, "state"):
+            return parsed_args.state
+        return None
+
+    @classmethod
+    def extract_catalog(cls, args: List[str]) -> Optional[Any]:
+        parsed_args = cls.parse_args(args)
+        if hasattr(parsed_args, "catalog"):
+            return parsed_args.catalog
+        return None
+
+    @classmethod
+    def extract_config(cls, args: List[str]) -> Optional[Any]:
+        parsed_args = cls.parse_args(args)
+        if hasattr(parsed_args, "config"):
+            return parsed_args.config
+        return None
+
+    def _emit_queued_messages(self, source: Source) -> Iterable[AirbyteMessage]:
+        if hasattr(source, "message_repository") and source.message_repository:
+            yield from source.message_repository.consume_queue()
+        return
+
+
+def launch(source: Source, args: List[str]) -> None:
+    source_entrypoint = AirbyteEntrypoint(source)
+    parsed_args = source_entrypoint.parse_args(args)
+    # temporarily removes the PrintBuffer because we're seeing weird print behavior for concurrent syncs
+    # Refer to: https://github.com/airbytehq/oncall/issues/6235
+    with PRINT_BUFFER:
+        for message in source_entrypoint.run(parsed_args):
+            # simply printing is creating issues for concurrent CDK as Python uses different two instructions to print: one for the message and
+            # the other for the break line. Adding `\n` to the message ensure that both are printed at the same time
+            print(f"{message}\n", end="")
+
+
+def _init_internal_request_filter() -> None:
+    """
+    Wraps the Python requests library to prevent sending requests to internal URL endpoints.
+    """
+    wrapped_fn = Session.send
+
+    @wraps(wrapped_fn)
+    def filtered_send(self: Any, request: PreparedRequest, **kwargs: Any) -> Response:
+        parsed_url = urlparse(request.url)
+
+        if parsed_url.scheme not in VALID_URL_SCHEMES:
+            raise requests.exceptions.InvalidSchema(
+                "Invalid Protocol Scheme: The endpoint that data is being requested from is using an invalid or insecure "
+                + f"protocol {parsed_url.scheme!r}. Valid protocol schemes: {','.join(VALID_URL_SCHEMES)}"
+            )
+
+        if not parsed_url.hostname:
+            raise requests.exceptions.InvalidURL(
+                "Invalid URL specified: The endpoint that data is being requested from is not a valid URL"
+            )
+
+        try:
+            is_private = _is_private_url(parsed_url.hostname, parsed_url.port)  # type: ignore [arg-type]
+            if is_private:
+                raise AirbyteTracedException(
+                    internal_message=f"Invalid URL endpoint: `{parsed_url.hostname!r}` belongs to a private network",
+                    failure_type=FailureType.config_error,
+                    message="Invalid URL endpoint: The endpoint that data is being requested from belongs to a private network. Source connectors only support requesting data from public API endpoints.",
+                )
+        except socket.gaierror as exception:
+            # This is a special case where the developer specifies an IP address string that is not formatted correctly like trailing
+            # whitespace which will fail the socket IP lookup. This only happens when using IP addresses and not text hostnames.
+            # Knowing that this is a request using the requests library, we will mock the exception without calling the lib
+            raise requests.exceptions.InvalidURL(f"Invalid URL {parsed_url}: {exception}")
+
+        return wrapped_fn(self, request, **kwargs)
+
+    Session.send = filtered_send  # type: ignore [method-assign]
+
+
+def _is_private_url(hostname: str, port: int) -> bool:
+    """
+    Helper method that checks if any of the IP addresses associated with a hostname belong to a private network.
+    """
+    address_info_entries = socket.getaddrinfo(hostname, port)
+    for entry in address_info_entries:
+        # getaddrinfo() returns entries in the form of a 5-tuple where the IP is stored as the sockaddr. For IPv4 this
+        # is a 2-tuple and for IPv6 it is a 4-tuple, but the address is always the first value of the tuple at 0.
+        # See https://docs.python.org/3/library/socket.html#socket.getaddrinfo for more details.
+        ip_address = entry[4][0]
+        if ipaddress.ip_address(ip_address).is_private:
+            return True
+    return False
+
+
+def main() -> None:
+    impl_module = os.environ.get("AIRBYTE_IMPL_MODULE", Source.__module__)
+    impl_class = os.environ.get("AIRBYTE_IMPL_PATH", Source.__name__)
+    module = importlib.import_module(impl_module)
+    impl = getattr(module, impl_class)
+
+    # set up and run entrypoint
+    source = impl()
+
+    if not isinstance(source, Source):
+        raise Exception("Source implementation provided does not implement Source class!")
+
+    launch(source, sys.argv[1:])

airbyte_cdk/exception_handler.py

@@ -0,0 +1,56 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+import sys
+from types import TracebackType
+from typing import Any, List, Mapping, Optional
+
+from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+
+def assemble_uncaught_exception(
+    exception_type: type[BaseException], exception_value: BaseException
+) -> AirbyteTracedException:
+    if issubclass(exception_type, AirbyteTracedException):
+        return exception_value  # type: ignore  # validated as part of the previous line
+    return AirbyteTracedException.from_exception(exception_value)
+
+
+def init_uncaught_exception_handler(logger: logging.Logger) -> None:
+    """
+    Handles uncaught exceptions by emitting an AirbyteTraceMessage and making sure they are not
+    printed to the console without having secrets removed.
+    """
+
+    def hook_fn(
+        exception_type: type[BaseException],
+        exception_value: BaseException,
+        traceback_: Optional[TracebackType],
+    ) -> Any:
+        # For developer ergonomics, we want to see the stack trace in the logs when we do a ctrl-c
+        if issubclass(exception_type, KeyboardInterrupt):
+            sys.__excepthook__(exception_type, exception_value, traceback_)
+            return
+
+        logger.fatal(exception_value, exc_info=exception_value)
+
+        # emit an AirbyteTraceMessage for any exception that gets to this spot
+        traced_exc = assemble_uncaught_exception(exception_type, exception_value)
+
+        traced_exc.emit_message()
+
+    sys.excepthook = hook_fn
+
+
+def generate_failed_streams_error_message(stream_failures: Mapping[str, List[Exception]]) -> str:
+    failures = "\n".join(
+        [
+            f"{stream}: {filter_secrets(exception.__repr__())}"
+            for stream, exceptions in stream_failures.items()
+            for exception in exceptions
+        ]
+    )
+    return f"During the sync, the following streams did not sync successfully: {failures}"