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

airbyte_cdk/utils/mapping_helpers.py

@@ -0,0 +1,162 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+import copy
+from typing import Any, Dict, List, Mapping, Optional, Union
+
+from airbyte_cdk.sources.declarative.requesters.request_option import (
+    RequestOption,
+    RequestOptionType,
+)
+from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
+
+
+def _merge_mappings(
+    target: Dict[str, Any],
+    source: Mapping[str, Any],
+    path: Optional[List[str]] = None,
+    allow_same_value_merge: bool = False,
+) -> None:
+    """
+    Recursively merge two dictionaries, raising an error if there are any conflicts.
+    For body_json requests (allow_same_value_merge=True), a conflict occurs only when the same path has different values.
+    For other request types (allow_same_value_merge=False), any duplicate key is a conflict, regardless of value.
+
+    Args:
+        target: The dictionary to merge into
+        source: The dictionary to merge from
+        path: The current path in the nested structure (for error messages)
+        allow_same_value_merge: Whether to allow merging the same value into the same key. Set to false by default, should only be true for body_json injections
+    """
+    path = path or []
+    for key, source_value in source.items():
+        current_path = path + [str(key)]
+
+        if key in target:
+            target_value = target[key]
+            if isinstance(target_value, dict) and isinstance(source_value, dict):
+                # Only body_json supports nested_structures
+                if not allow_same_value_merge:
+                    raise ValueError(
+                        f"Request body collision, duplicate keys detected at key path: {'.'.join(current_path)}. Please ensure that all keys in the request are unique."
+                    )
+                # If both are dictionaries, recursively merge them
+                _merge_mappings(target_value, source_value, current_path, allow_same_value_merge)
+
+            elif not allow_same_value_merge or target_value != source_value:
+                # If same key has different values, that's a conflict
+                raise ValueError(
+                    f"Request body collision, duplicate keys detected at key path: {'.'.join(current_path)}. Please ensure that all keys in the request are unique."
+                )
+        else:
+            # No conflict, just copy the value (using deepcopy for nested structures)
+            target[key] = copy.deepcopy(source_value)
+
+
+def combine_mappings(
+    mappings: List[Optional[Union[Mapping[str, Any], str]]],
+    allow_same_value_merge: bool = False,
+) -> Union[Mapping[str, Any], str]:
+    """
+    Combine multiple mappings into a single mapping.
+
+    For body_json requests (allow_same_value_merge=True):
+        - Supports nested structures (e.g., {"data": {"user": {"id": 1}}})
+        - Allows duplicate keys if their values match
+        - Raises error if same path has different values
+
+    For other request types (allow_same_value_merge=False):
+        - Only supports flat structures
+        - Any duplicate key raises an error, regardless of value
+
+    Args:
+        mappings: List of mappings to combine
+        allow_same_value_merge: Whether to allow duplicate keys with matching values.
+                              Should only be True for body_json requests.
+
+    Returns:
+        A single mapping combining all inputs, or a string if there is exactly one
+        string mapping and no other non-empty mappings.
+
+    Raises:
+        ValueError: If there are:
+            - Multiple string mappings
+            - Both a string mapping and non-empty dictionary mappings
+            - Conflicting keys/paths based on allow_same_value_merge setting
+    """
+    if not mappings:
+        return {}
+
+    # Count how many string options we have, ignoring None values
+    string_options = sum(isinstance(mapping, str) for mapping in mappings if mapping is not None)
+    if string_options > 1:
+        raise ValueError("Cannot combine multiple string options")
+
+    # Filter out None values and empty mappings
+    non_empty_mappings = [
+        m for m in mappings if m is not None and not (isinstance(m, Mapping) and not m)
+    ]
+
+    # If there is only one string option and no other non-empty mappings, return it
+    if string_options == 1:
+        if len(non_empty_mappings) > 1:
+            raise ValueError("Cannot combine multiple options if one is a string")
+        return next(m for m in non_empty_mappings if isinstance(m, str))
+
+    # Start with an empty result and merge each mapping into it
+    result: Dict[str, Any] = {}
+    for mapping in non_empty_mappings:
+        if mapping and isinstance(mapping, Mapping):
+            _merge_mappings(result, mapping, allow_same_value_merge=allow_same_value_merge)
+
+    return result
+
+
+def _validate_component_request_option_paths(
+    config: Config, *request_options: Optional[RequestOption]
+) -> None:
+    """
+    Validates that a component with multiple request options does not have conflicting paths.
+    Uses dummy values for validation since actual values might not be available at init time.
+    """
+    grouped_options: Dict[RequestOptionType, List[RequestOption]] = {}
+    for option in request_options:
+        if option:
+            grouped_options.setdefault(option.inject_into, []).append(option)
+
+    for inject_type, options in grouped_options.items():
+        if len(options) <= 1:
+            continue
+
+        option_dicts: List[Optional[Union[Mapping[str, Any], str]]] = []
+        for i, option in enumerate(options):
+            option_dict: Dict[str, Any] = {}
+            # Use indexed dummy values to ensure we catch conflicts
+            option.inject_into_request(option_dict, f"dummy_value_{i}", config)
+            option_dicts.append(option_dict)
+
+        try:
+            combine_mappings(
+                option_dicts, allow_same_value_merge=(inject_type == RequestOptionType.body_json)
+            )
+        except ValueError as error:
+            raise ValueError(error)
+
+
+def get_interpolation_context(
+    stream_state: Optional[StreamState] = None,
+    stream_slice: Optional[StreamSlice] = None,
+    next_page_token: Optional[Mapping[str, Any]] = None,
+) -> Mapping[str, Any]:
+    return {
+        "stream_slice": stream_slice,
+        "next_page_token": next_page_token,
+        # update the context with extra fields, if passed.
+        **(
+            stream_slice.extra_fields
+            if stream_slice is not None and hasattr(stream_slice, "extra_fields")
+            else {}
+        ),
+    }

airbyte_cdk/utils/message_utils.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from airbyte_cdk.models import AirbyteMessage, Type
+from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
+
+
+def get_stream_descriptor(message: AirbyteMessage) -> HashableStreamDescriptor:
+    match message.type:
+        case Type.RECORD:
+            return HashableStreamDescriptor(
+                name=message.record.stream,  # type: ignore[union-attr] # record has `stream`
+                namespace=message.record.namespace,  # type: ignore[union-attr] # record has `namespace`
+            )
+        case Type.STATE:
+            if not message.state.stream or not message.state.stream.stream_descriptor:  # type: ignore[union-attr] # state has `stream`
+                raise ValueError(
+                    "State message was not in per-stream state format, which is required for record counts."
+                )
+            return HashableStreamDescriptor(
+                name=message.state.stream.stream_descriptor.name,  # type: ignore[union-attr] # state has `stream`
+                namespace=message.state.stream.stream_descriptor.namespace,  # type: ignore[union-attr] # state has `stream`
+            )
+        case _:
+            raise NotImplementedError(
+                f"get_stream_descriptor is not implemented for message type '{message.type}'."
+            )

airbyte_cdk/utils/oneof_option_config.py

@@ -0,0 +1,33 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Any, Dict
+
+
+class OneOfOptionConfig:
+    """
+    Base class to configure a Pydantic model that's used as a oneOf option in a parent model in a way that's compatible with all Airbyte consumers.
+
+    Inherit from this class in the nested Config class in a model and set title and description (these show up in the UI) and discriminator (this is making sure it's marked as required in the schema).
+
+    Usage:
+
+        ```python
+        class OptionModel(BaseModel):
+            mode: Literal["option_a"] = Field("option_a", const=True)
+            option_a_field: str = Field(...)
+
+            class Config(OneOfOptionConfig):
+                title = "Option A"
+                description = "Option A description"
+                discriminator = "mode"
+        ```
+    """
+
+    @staticmethod
+    def schema_extra(schema: Dict[str, Any], model: Any) -> None:
+        if hasattr(model.Config, "description"):
+            schema["description"] = model.Config.description
+        if hasattr(model.Config, "discriminator"):
+            schema.setdefault("required", []).append(model.Config.discriminator)

airbyte_cdk/utils/print_buffer.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import sys
+import time
+from io import StringIO
+from threading import RLock
+from types import TracebackType
+from typing import Optional
+
+
+class PrintBuffer:
+    """
+    A class to buffer print statements and flush them at a specified interval.
+
+    The PrintBuffer class is designed to capture and buffer output that would
+    normally be printed to the standard output (stdout). This can be useful for
+    scenarios where you want to minimize the number of I/O operations by grouping
+    multiple print statements together and flushing them as a single operation.
+
+    Attributes:
+        buffer (StringIO): A buffer to store the messages before flushing.
+        flush_interval (float): The time interval (in seconds) after which the buffer is flushed.
+        last_flush_time (float): The last time the buffer was flushed.
+        lock (RLock): A reentrant lock to ensure thread-safe operations.
+
+    Methods:
+        write(message: str) -> None:
+            Writes a message to the buffer and flushes if the interval has passed.
+
+        flush() -> None:
+            Flushes the buffer content to the standard output.
+
+        __enter__() -> "PrintBuffer":
+            Enters the runtime context related to this object, redirecting stdout and stderr.
+
+        __exit__(exc_type, exc_val, exc_tb) -> None:
+            Exits the runtime context and restores the original stdout and stderr.
+    """
+
+    def __init__(self, flush_interval: float = 0.1):
+        self.buffer = StringIO()
+        self.flush_interval = flush_interval
+        self.last_flush_time = time.monotonic()
+        self.lock = RLock()
+
+    def write(self, message: str) -> None:
+        with self.lock:
+            self.buffer.write(message)
+            current_time = time.monotonic()
+            if (current_time - self.last_flush_time) >= self.flush_interval:
+                self.flush()
+                self.last_flush_time = current_time
+
+    def flush(self) -> None:
+        with self.lock:
+            combined_message = self.buffer.getvalue()
+            sys.__stdout__.write(combined_message)  # type: ignore[union-attr]
+            self.buffer = StringIO()
+
+    def __enter__(self) -> "PrintBuffer":
+        self.old_stdout, self.old_stderr = sys.stdout, sys.stderr
+        # Used to disable buffering during the pytest session, because it is not compatible with capsys
+        if "pytest" not in str(type(sys.stdout)).lower():
+            sys.stdout = self
+            sys.stderr = self
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[BaseException],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self.flush()
+        sys.stdout, sys.stderr = self.old_stdout, self.old_stderr

airbyte_cdk/utils/schema_inferrer.py

@@ -0,0 +1,270 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from collections import defaultdict
+from typing import Any, Dict, List, Mapping, Optional
+
+from genson import SchemaBuilder, SchemaNode
+from genson.schema.strategies.object import Object
+from genson.schema.strategies.scalar import Number
+
+from airbyte_cdk.models import AirbyteRecordMessage
+
+# schema keywords
+_TYPE = "type"
+_NULL_TYPE = "null"
+_OBJECT_TYPE = "object"
+_ANY_OF = "anyOf"
+_ITEMS = "items"
+_PROPERTIES = "properties"
+_REQUIRED = "required"
+
+
+class NoRequiredObj(Object):
+    """
+    This class has Object behaviour, but it does not generate "required[]" fields
+    every time it parses object. So we don't add unnecessary extra field.
+
+    The logic is that even reading all the data from a source, it does not mean that there can be another record added with those fields as
+    optional. Hence, we make everything nullable.
+    """
+
+    def to_schema(self) -> Mapping[str, Any]:
+        schema: Dict[str, Any] = super(NoRequiredObj, self).to_schema()
+        schema.pop("required", None)
+        return schema
+
+
+class IntegerToNumber(Number):
+    """
+    This class has the regular Number behaviour, but it will never emit an integer type.
+    """
+
+    def __init__(self, node_class: SchemaNode):
+        super().__init__(node_class)
+        self._type = "number"
+
+
+class NoRequiredSchemaBuilder(SchemaBuilder):
+    EXTRA_STRATEGIES = (NoRequiredObj, IntegerToNumber)
+
+
+# This type is inferred from the genson lib, but there is no alias provided for it - creating it here for type safety
+InferredSchema = Dict[str, Any]
+
+
+class SchemaValidationException(Exception):
+    @classmethod
+    def merge_exceptions(
+        cls, exceptions: List["SchemaValidationException"]
+    ) -> "SchemaValidationException":
+        # We assume the schema is the same for all SchemaValidationException
+        return SchemaValidationException(
+            exceptions[0].schema,
+            [x for exception in exceptions for x in exception._validation_errors],
+        )
+
+    def __init__(self, schema: InferredSchema, validation_errors: List[Exception]):
+        self._schema = schema
+        self._validation_errors = validation_errors
+
+    @property
+    def schema(self) -> InferredSchema:
+        return self._schema
+
+    @property
+    def validation_errors(self) -> List[str]:
+        return list(map(lambda error: str(error), self._validation_errors))
+
+
+class SchemaInferrer:
+    """
+    This class is used to infer a JSON schema which fits all the records passed into it
+    throughout its lifecycle via the accumulate method.
+
+    Instances of this class are stateful, meaning they build their inferred schemas
+    from every record passed into the accumulate method.
+
+    """
+
+    stream_to_builder: Dict[str, SchemaBuilder]
+
+    def __init__(
+        self, pk: Optional[List[List[str]]] = None, cursor_field: Optional[List[List[str]]] = None
+    ) -> None:
+        self.stream_to_builder = defaultdict(NoRequiredSchemaBuilder)
+        self._pk = [] if pk is None else pk
+        self._cursor_field = [] if cursor_field is None else cursor_field
+
+    def accumulate(self, record: AirbyteRecordMessage) -> None:
+        """Uses the input record to add to the inferred schemas maintained by this object"""
+        self.stream_to_builder[record.stream].add_object(record.data)
+
+    def _null_type_in_any_of(self, node: InferredSchema) -> bool:
+        if _ANY_OF in node:
+            return {_TYPE: _NULL_TYPE} in node[_ANY_OF]
+        else:
+            return False
+
+    def _remove_type_from_any_of(self, node: InferredSchema) -> None:
+        if _ANY_OF in node:
+            node.pop(_TYPE, None)
+
+    def _clean_any_of(self, node: InferredSchema) -> None:
+        if len(node[_ANY_OF]) == 2 and self._null_type_in_any_of(node):
+            real_type = (
+                node[_ANY_OF][1] if node[_ANY_OF][0][_TYPE] == _NULL_TYPE else node[_ANY_OF][0]
+            )
+            node.update(real_type)
+            node[_TYPE] = [node[_TYPE], _NULL_TYPE]
+            node.pop(_ANY_OF)
+        # populate `type` for `anyOf` if it's not present to pass all other checks
+        elif len(node[_ANY_OF]) == 2 and not self._null_type_in_any_of(node):
+            node[_TYPE] = [_NULL_TYPE]
+
+    def _clean_properties(self, node: InferredSchema) -> None:
+        for key, value in list(node[_PROPERTIES].items()):
+            if isinstance(value, dict) and value.get(_TYPE, None) == _NULL_TYPE:
+                node[_PROPERTIES].pop(key)
+            else:
+                self._clean(value)
+
+    def _ensure_null_type_on_top(self, node: InferredSchema) -> None:
+        if isinstance(node[_TYPE], list):
+            if _NULL_TYPE in node[_TYPE]:
+                # we want to make sure null is always at the end as it makes schemas more readable
+                node[_TYPE].remove(_NULL_TYPE)
+            node[_TYPE].append(_NULL_TYPE)
+        else:
+            node[_TYPE] = [node[_TYPE], _NULL_TYPE]
+
+    def _clean(self, node: InferredSchema) -> InferredSchema:
+        """
+        Recursively cleans up a produced schema:
+        - remove anyOf if one of them is just a null value
+        - remove properties of type "null"
+        """
+
+        if isinstance(node, dict):
+            if _ANY_OF in node:
+                self._clean_any_of(node)
+
+            if _PROPERTIES in node and isinstance(node[_PROPERTIES], dict):
+                self._clean_properties(node)
+
+            if _ITEMS in node:
+                self._clean(node[_ITEMS])
+
+            # this check needs to follow the "anyOf" cleaning as it might populate `type`
+            self._ensure_null_type_on_top(node)
+
+        # remove added `type: ["null"]` for `anyOf` nested node
+        self._remove_type_from_any_of(node)
+
+        return node
+
+    def _add_required_properties(self, node: InferredSchema) -> InferredSchema:
+        """
+        This method takes properties that should be marked as required (self._pk and self._cursor_field) and travel the schema to mark every
+        node as required.
+        """
+        # Removing nullable for the root as when we call `_clean`, we make everything nullable
+        node[_TYPE] = _OBJECT_TYPE
+
+        exceptions = []
+        for field in [x for x in [self._pk, self._cursor_field] if x]:
+            try:
+                self._add_fields_as_required(node, field)
+            except SchemaValidationException as exception:
+                exceptions.append(exception)
+
+        if exceptions:
+            raise SchemaValidationException.merge_exceptions(exceptions)
+
+        return node
+
+    def _add_fields_as_required(self, node: InferredSchema, composite_key: List[List[str]]) -> None:
+        """
+        Take a list of nested keys (this list represents a composite key) and travel the schema to mark every node as required.
+        """
+        errors: List[Exception] = []
+
+        for path in composite_key:
+            try:
+                self._add_field_as_required(node, path)
+            except ValueError as exception:
+                errors.append(exception)
+
+        if errors:
+            raise SchemaValidationException(node, errors)
+
+    def _add_field_as_required(
+        self, node: InferredSchema, path: List[str], traveled_path: Optional[List[str]] = None
+    ) -> None:
+        """
+        Take a nested key and travel the schema to mark every node as required.
+        """
+        self._remove_null_from_type(node)
+        if self._is_leaf(path):
+            return
+
+        if not traveled_path:
+            traveled_path = []
+
+        if _PROPERTIES not in node:
+            # This validation is only relevant when `traveled_path` is empty
+            raise ValueError(
+                f"Path {traveled_path} does not refer to an object but is `{node}` and hence {path} can't be marked as required."
+            )
+
+        next_node = path[0]
+        if next_node not in node[_PROPERTIES]:
+            raise ValueError(
+                f"Path {traveled_path} does not have field `{next_node}` in the schema and hence can't be marked as required."
+            )
+
+        if _TYPE not in node:
+            # We do not expect this case to happen but we added a specific error message just in case
+            raise ValueError(
+                f"Unknown schema error: {traveled_path} is expected to have a type but did not. Schema inferrence is probably broken"
+            )
+
+        if node[_TYPE] not in [
+            _OBJECT_TYPE,
+            [_NULL_TYPE, _OBJECT_TYPE],
+            [_OBJECT_TYPE, _NULL_TYPE],
+        ]:
+            raise ValueError(
+                f"Path {traveled_path} is expected to be an object but was of type `{node['properties'][next_node]['type']}`"
+            )
+
+        if _REQUIRED not in node or not node[_REQUIRED]:
+            node[_REQUIRED] = [next_node]
+        elif next_node not in node[_REQUIRED]:
+            node[_REQUIRED].append(next_node)
+
+        traveled_path.append(next_node)
+        self._add_field_as_required(node[_PROPERTIES][next_node], path[1:], traveled_path)
+
+    def _is_leaf(self, path: List[str]) -> bool:
+        return len(path) == 0
+
+    def _remove_null_from_type(self, node: InferredSchema) -> None:
+        if isinstance(node[_TYPE], list):
+            if _NULL_TYPE in node[_TYPE]:
+                node[_TYPE].remove(_NULL_TYPE)
+            if len(node[_TYPE]) == 1:
+                node[_TYPE] = node[_TYPE][0]
+
+    def get_stream_schema(self, stream_name: str) -> Optional[InferredSchema]:
+        """
+        Returns the inferred JSON schema for the specified stream. Might be `None` if there were no records for the given stream name.
+        """
+        return (
+            self._add_required_properties(
+                self._clean(self.stream_to_builder[stream_name].to_schema())
+            )
+            if stream_name in self.stream_to_builder
+            else None
+        )

airbyte_cdk/utils/slice_hasher.py

@@ -0,0 +1,37 @@
+import hashlib
+import json
+from typing import Any, Final, Mapping, Optional
+
+
+class SliceEncoder(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        if hasattr(obj, "__json_serializable__"):
+            return obj.__json_serializable__()
+
+        # Let the base class default method raise the TypeError
+        return super().default(obj)
+
+
+class SliceHasher:
+    _ENCODING: Final = "utf-8"
+
+    @classmethod
+    def hash(
+        cls,
+        stream_name: str = "<stream name not provided>",
+        stream_slice: Optional[Mapping[str, Any]] = None,
+    ) -> int:
+        """
+        Note that streams partition with the same slicing value but with different names might collapse if stream name is not provided
+        """
+        if stream_slice:
+            try:
+                s = json.dumps(stream_slice, sort_keys=True, cls=SliceEncoder)
+                hash_input = f"{stream_name}:{s}".encode(cls._ENCODING)
+            except TypeError as e:
+                raise ValueError(f"Failed to serialize stream slice: {e}")
+        else:
+            hash_input = stream_name.encode(cls._ENCODING)
+
+        # Use last 8 bytes as 64-bit integer for better distribution
+        return int.from_bytes(hashlib.sha256(hash_input).digest()[-8:], "big")

airbyte_cdk/utils/spec_schema_transformations.py

@@ -0,0 +1,26 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import json
+import re
+from typing import Any
+
+from jsonschema import RefResolver
+
+
+def resolve_refs(schema: dict[str, Any]) -> dict[str, Any]:
+    """
+    For spec schemas generated using Pydantic models, the resulting JSON schema can contain refs between object
+    relationships.
+    """
+    json_schema_ref_resolver = RefResolver.from_schema(schema)
+    str_schema = json.dumps(schema)
+    for ref_block in re.findall(r'{"\$ref": "#\/definitions\/.+?(?="})"}', str_schema):
+        ref = json.loads(ref_block)["$ref"]
+        str_schema = str_schema.replace(
+            ref_block, json.dumps(json_schema_ref_resolver.resolve(ref)[1])
+        )
+    pyschema: dict[str, Any] = json.loads(str_schema)
+    del pyschema["definitions"]
+    return pyschema

airbyte_cdk/utils/stream_status_utils.py

@@ -0,0 +1,43 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+from datetime import datetime
+from typing import List, Optional, Union
+
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteStream,
+    AirbyteStreamStatus,
+    AirbyteStreamStatusReason,
+    AirbyteStreamStatusTraceMessage,
+    AirbyteTraceMessage,
+    StreamDescriptor,
+    TraceType,
+)
+from airbyte_cdk.models import Type as MessageType
+
+
+def as_airbyte_message(
+    stream: Union[AirbyteStream, StreamDescriptor],
+    current_status: AirbyteStreamStatus,
+    reasons: Optional[List[AirbyteStreamStatusReason]] = None,
+) -> AirbyteMessage:
+    """
+    Builds an AirbyteStreamStatusTraceMessage for the provided stream
+    """
+
+    now_millis = datetime.now().timestamp() * 1000.0
+
+    trace_message = AirbyteTraceMessage(
+        type=TraceType.STREAM_STATUS,
+        emitted_at=now_millis,
+        stream_status=AirbyteStreamStatusTraceMessage(
+            stream_descriptor=StreamDescriptor(name=stream.name, namespace=stream.namespace),
+            status=current_status,
+            reasons=reasons,
+        ),
+    )
+
+    return AirbyteMessage(type=MessageType.TRACE, trace=trace_message)