airbyte-cdk 6.34.0.dev0__py3-none-any.whl → 6.34.0.dev2__py3-none-any.whl

airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -2,10 +2,11 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import dataclasses
+
+from dataclasses import asdict, dataclass, field
 from typing import Any, List, Mapping
 
-from airbyte_cdk.connector_builder.message_grouper import MessageGrouper
+from airbyte_cdk.connector_builder.test_reader import TestReader
 from airbyte_cdk.models import (
     AirbyteMessage,
     AirbyteRecordMessage,
@@ -32,11 +33,11 @@ MAX_SLICES_KEY = "max_slices"
 MAX_RECORDS_KEY = "max_records"
 
 
-@dataclasses.dataclass
+@dataclass
 class TestReadLimits:
-    max_records: int = dataclasses.field(default=DEFAULT_MAXIMUM_RECORDS)
-    max_pages_per_slice: int = dataclasses.field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
-    max_slices: int = dataclasses.field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
+    max_records: int = field(default=DEFAULT_MAXIMUM_RECORDS)
+    max_pages_per_slice: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
+    max_slices: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
 
 
 def get_limits(config: Mapping[str, Any]) -> TestReadLimits:
@@ -73,17 +74,20 @@ def read_stream(
     limits: TestReadLimits,
 ) -> AirbyteMessage:
     try:
-        handler = MessageGrouper(limits.max_pages_per_slice, limits.max_slices, limits.max_records)
-        stream_name = configured_catalog.streams[
-            0
-        ].stream.name  # The connector builder only supports a single stream
-        stream_read = handler.get_message_groups(
+        test_read_handler = TestReader(
+            limits.max_pages_per_slice, limits.max_slices, limits.max_records
+        )
+        # The connector builder only supports a single stream
+        stream_name = configured_catalog.streams[0].stream.name
+
+        stream_read = test_read_handler.run_test_read(
             source, config, configured_catalog, state, limits.max_records
         )
+
         return AirbyteMessage(
             type=MessageType.RECORD,
             record=AirbyteRecordMessage(
-                data=dataclasses.asdict(stream_read), stream=stream_name, emitted_at=_emitted_at()
+                data=asdict(stream_read), stream=stream_name, emitted_at=_emitted_at()
             ),
         )
     except Exception as exc:

airbyte_cdk/connector_builder/test_reader/__init__.py

@@ -0,0 +1,7 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from .reader import TestReader
+
+__all__ = ("TestReader",)

airbyte_cdk/connector_builder/test_reader/helpers.py

@@ -0,0 +1,591 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import json
+from copy import deepcopy
+from json import JSONDecodeError
+from typing import Any, Dict, List, Mapping, Optional
+
+from airbyte_cdk.connector_builder.models import (
+    AuxiliaryRequest,
+    HttpRequest,
+    HttpResponse,
+    StreamReadPages,
+    StreamReadSlices,
+)
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    OrchestratorType,
+    TraceType,
+)
+from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.sources.utils.slice_logger import SliceLogger
+from airbyte_cdk.sources.utils.types import JsonType
+from airbyte_cdk.utils.datetime_format_inferrer import DatetimeFormatInferrer
+from airbyte_cdk.utils.schema_inferrer import (
+    SchemaInferrer,
+)
+
+from .types import LOG_MESSAGES_OUTPUT_TYPE
+
+# -------
+# Parsers
+# -------
+
+
+def airbyte_message_to_json(message: AirbyteMessage) -> Optional[Dict[str, JsonType]]:
+    """
+    Converts an AirbyteMessage to a JSON dictionary if its type is LOG.
+
+    This function attempts to parse the 'log' field of the given AirbyteMessage when its type is MessageType.LOG.
+    If the parsed JSON object exists but is not a dictionary, a ValueError is raised. If the message is not of type LOG,
+    the function returns None.
+
+    Parameters:
+        message (AirbyteMessage): The AirbyteMessage instance containing the log data.
+
+    Returns:
+        Optional[Dict[str, JsonType]]: The parsed log message as a dictionary if the message type is LOG, otherwise None.
+
+    Raises:
+        ValueError: If the parsed log message is not a dictionary.
+    """
+    if is_log_message(message):
+        json_object = parse_json(message.log)  # type: ignore
+
+        if json_object is not None and not isinstance(json_object, dict):
+            raise ValueError(
+                f"Expected log message to be a dict, got {json_object} of type {type(json_object)}"
+            )
+
+        return json_object
+    return None
+
+
+def clean_config(config: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Cleans the configuration dictionary by removing all keys that start with a double underscore.
+
+    This function creates a deep copy of the provided configuration dictionary and iterates
+    over its keys, deleting any key that begins with '__'. This is useful for filtering out
+    internal or meta-data fields that are not meant to be part of the final configuration.
+
+    Args:
+        config (Dict[str, Any]): The input configuration dictionary containing various key-value pairs.
+
+    Returns:
+        Dict[str, Any]: A deep copy of the original configuration with keys starting with '__' removed.
+    """
+    cleaned_config = deepcopy(config)
+    for key in config.keys():
+        if key.startswith("__"):
+            del cleaned_config[key]
+    return cleaned_config
+
+
+def create_request_from_log_message(json_http_message: Dict[str, Any]) -> HttpRequest:
+    """
+    Creates an HttpRequest object from the provided JSON-formatted log message.
+
+    This function parses a dictionary that represents a logged HTTP message, extracting the URL, HTTP method,
+    headers, and body from nested dictionary structures. It is assumed that the expected keys and nested keys exist
+    or default values are used.
+
+    Parameters:
+        json_http_message (Dict[str, Any]):
+            A dictionary containing log message details with the following expected structure:
+                {
+                    "url": {
+                        "full": "<full_url>"
+                    },
+                    "http": {
+                        "request": {
+                            "method": "<HTTP_method>",
+                            "headers": <headers>,
+                            "body": {
+                                "content": "<body_content>"
+                            }
+                        }
+                    }
+                }
+
+    Returns:
+        HttpRequest:
+            An HttpRequest instance initialized with:
+                - url: Extracted from json_http_message["url"]["full"], defaults to an empty string if missing.
+                - http_method: Extracted from json_http_message["http"]["request"]["method"], defaults to an empty string if missing.
+                - headers: Extracted from json_http_message["http"]["request"]["headers"].
+                - body: Extracted from json_http_message["http"]["request"]["body"]["content"], defaults to an empty string if missing.
+    """
+    url = json_http_message.get("url", {}).get("full", "")
+    request = json_http_message.get("http", {}).get("request", {})
+    return HttpRequest(
+        url=url,
+        http_method=request.get("method", ""),
+        headers=request.get("headers"),
+        body=request.get("body", {}).get("content", ""),
+    )
+
+
+def create_response_from_log_message(json_http_message: Dict[str, Any]) -> HttpResponse:
+    """
+    Generate an HttpResponse instance from a JSON log message containing HTTP response details.
+
+    Parameters:
+        json_http_message (Dict[str, Any]): A dictionary representing a JSON-encoded HTTP message.
+            It should include an "http" key with a nested "response" dictionary that contains:
+                - "status_code": The HTTP status code.
+                - "body": A dictionary with a "content" key for the response body.
+                - "headers": The HTTP response headers.
+
+    Returns:
+        HttpResponse: An HttpResponse object constructed from the extracted status code, body content, and headers.
+    """
+    response = json_http_message.get("http", {}).get("response", {})
+    body = response.get("body", {}).get("content", "")
+    return HttpResponse(
+        status=response.get("status_code"), body=body, headers=response.get("headers")
+    )
+
+
+def parse_json(log_message: AirbyteLogMessage) -> JsonType:
+    """
+    Parse and extract a JSON object from an Airbyte log message.
+
+    This function attempts to decode the JSON string contained in the message field
+    of the provided AirbyteLogMessage instance. If the decoding process fails due to
+    malformed JSON, the function returns None.
+
+    Args:
+        log_message (AirbyteLogMessage): A log message object containing a JSON-formatted string in its 'message' attribute.
+
+    Returns:
+        JsonType: The parsed JSON object if decoding is successful; otherwise, None.
+    """
+    # TODO: As a temporary stopgap, the CDK emits request/response data as a log message string. Ideally this should come in the
+    # form of a custom message object defined in the Airbyte protocol, but this unblocks us in the immediate while the
+    # protocol change is worked on.
+    try:
+        json_object: JsonType = json.loads(log_message.message)
+        return json_object
+    except JSONDecodeError:
+        return None
+
+
+def parse_slice_description(log_message: str) -> Dict[str, Any]:
+    """
+    Parses a log message containing a JSON payload and returns it as a dictionary.
+
+    The function removes a predefined logging prefix (defined by the constant
+    SliceLogger.SLICE_LOG_PREFIX) from the beginning of the log message and then
+    parses the remaining string as JSON.
+
+    Parameters:
+        log_message (str): The log message string that includes the JSON payload,
+                        prefixed by SliceLogger.SLICE_LOG_PREFIX.
+
+    Returns:
+        Dict[str, Any]: A dictionary resulting from parsing the modified log message.
+
+    Raises:
+        json.JSONDecodeError: If the log message (after prefix removal) is not a valid JSON.
+    """
+    return json.loads(log_message.replace(SliceLogger.SLICE_LOG_PREFIX, "", 1))  # type: ignore
+
+
+# -------
+# Conditions
+# -------
+
+
+def should_close_page(
+    at_least_one_page_in_group: bool,
+    message: AirbyteMessage,
+    json_message: Optional[Dict[str, Any]],
+) -> bool:
+    """
+    Determines whether a page should be closed based on its content and state.
+
+    Args:
+        at_least_one_page_in_group (bool): Indicates if there is at least one page in the group.
+        message (AirbyteMessage): The message object containing details such as type and log information.
+        json_message (Optional[Dict[str, Any]]): A JSON representation of the message that may provide additional context,
+            particularly for HTTP requests.
+
+    Returns:
+        bool: True if all of the following conditions are met:
+            - There is at least one page in the group.
+            - The message type is MessageType.LOG.
+            - Either the JSON message corresponds to a page HTTP request (as determined by _is_page_http_request)
+            or the log message starts with "slice:".
+        Otherwise, returns False.
+    """
+    return (
+        at_least_one_page_in_group
+        and is_log_message(message)
+        and (
+            is_page_http_request(json_message) or message.log.message.startswith("slice:")  # type: ignore[union-attr] # AirbyteMessage with MessageType.LOG has log.message
+        )
+    )
+
+
+def should_process_slice_descriptor(message: AirbyteMessage) -> bool:
+    """
+    Determines whether the given AirbyteMessage should be processed as a slice descriptor.
+
+    This function checks if the message is a log message and if its log content starts with the
+    specific slice log prefix. It is used to filter out messages that represent slice descriptors
+    for further processing.
+
+    Parameters:
+        message (AirbyteMessage): The message to evaluate.
+
+    Returns:
+        bool: True if the message is a log message whose log message starts with the predefined
+            slice log prefix, indicating it is a slice descriptor; otherwise, False.
+    """
+    return is_log_message(message) and message.log.message.startswith(  # type: ignore
+        SliceLogger.SLICE_LOG_PREFIX
+    )
+
+
+def should_close_page_for_slice(at_least_one_page_in_group: bool, message: AirbyteMessage) -> bool:
+    """
+    Determines whether the current slice page should be closed.
+
+    This function checks if there is at least one page in the current group and if further processing
+    of the slice descriptor is required based on the provided Airbyte message.
+
+    Args:
+        at_least_one_page_in_group (bool): Indicates if at least one page already exists in the slice group.
+        message (AirbyteMessage): The message containing the slice descriptor information to be evaluated.
+
+    Returns:
+        bool: True if both conditions are met and the slice page needs to be closed; otherwise, False.
+    """
+    return at_least_one_page_in_group and should_process_slice_descriptor(message)
+
+
+def is_page_http_request(json_message: Optional[Dict[str, Any]]) -> bool:
+    """
+    Determines whether a given JSON message represents a page HTTP request.
+
+    This function checks if the provided JSON message qualifies as a page HTTP request by verifying that:
+    1. The JSON message exists.
+    2. The JSON message is recognized as a valid HTTP log.
+    3. The JSON message is not classified as an auxiliary HTTP request.
+
+    Args:
+        json_message (Optional[Dict[str, Any]]): A dictionary containing the JSON message to be evaluated.
+            If None or empty, the message will not be considered a page HTTP request.
+
+    Returns:
+        bool: True if the JSON message is a valid HTTP log and not an auxiliary HTTP request; otherwise, False.
+    """
+    if not json_message:
+        return False
+    else:
+        return is_http_log(json_message) and not is_auxiliary_http_request(json_message)
+
+
+def is_http_log(message: Dict[str, JsonType]) -> bool:
+    """
+    Determine if the provided log message represents an HTTP log.
+
+    This function inspects the given message dictionary for the presence of the "http" key.
+    If the key exists and its value is truthy, the function interprets the message as an HTTP log.
+
+    Args:
+        message (Dict[str, JsonType]): A dictionary containing log data. It may include an "http" key
+            whose truthy value indicates an HTTP log.
+
+    Returns:
+        bool: True if the message is an HTTP log (i.e., "http" exists and is truthy); otherwise, False.
+    """
+    return bool(message.get("http", False))
+
+
+def is_auxiliary_http_request(message: Optional[Dict[str, Any]]) -> bool:
+    """
+    Determines if the provided message represents an auxiliary HTTP request.
+
+    A auxiliary request is a request that is performed and will not directly lead to record for the specific stream it is being queried.
+
+    A couple of examples are:
+        * OAuth authentication
+        * Substream slice generation
+
+    Parameters:
+        message (Optional[Dict[str, Any]]): A dictionary representing a log message for an HTTP request.
+            The dictionary may contain nested keys indicating whether the request is auxiliary.
+
+    Returns:
+        bool: True if the message is an HTTP log and indicates an auxiliary request; otherwise, False.
+    """
+    if not message:
+        return False
+
+    return is_http_log(message) and message.get("http", {}).get("is_auxiliary", False)
+
+
+def is_log_message(message: AirbyteMessage) -> bool:
+    """
+    Determines whether the provided message is of type LOG.
+
+    Args:
+        message (AirbyteMessage): The message to evaluate.
+
+    Returns:
+        bool: True if the message's type is LOG, otherwise False.
+    """
+    return message.type == MessageType.LOG  # type: ignore
+
+
+def is_trace_with_error(message: AirbyteMessage) -> bool:
+    """
+    Determines whether the provided AirbyteMessage is a TRACE message with an error.
+
+    This function checks if the message's type is TRACE and that its trace component is of type ERROR.
+
+    Parameters:
+        message (AirbyteMessage): The Airbyte message to be evaluated.
+
+    Returns:
+        bool: True if the message is a TRACE message with an error, False otherwise.
+    """
+    return message.type == MessageType.TRACE and message.trace.type == TraceType.ERROR  # type: ignore
+
+
+def is_record_message(message: AirbyteMessage) -> bool:
+    """
+    Determines whether the provided Airbyte message represents a record.
+
+    Parameters:
+        message (AirbyteMessage): The message instance to check. It should include a 'type' attribute that is comparable to MessageType.RECORD.
+
+    Returns:
+        bool: True if the message type is RECORD, otherwise False.
+    """
+    return message.type == MessageType.RECORD  # type: ignore
+
+
+def is_config_update_message(message: AirbyteMessage) -> bool:
+    """
+    Determine whether the provided AirbyteMessage represents a connector configuration update.
+
+    This function evaluates if the message is a control message and if its control type
+    matches that of a connector configuration update (i.e., OrchestratorType.CONNECTOR_CONFIG).
+    It is primarily used to filter messages related to configuration updates in the data pipeline.
+
+    Parameters:
+        message (AirbyteMessage): The message object to be evaluated.
+
+    Returns:
+        bool: True if the message is a connector configuration update message, False otherwise.
+    """
+    return (  # type: ignore
+        message.type == MessageType.CONTROL
+        and message.control.type == OrchestratorType.CONNECTOR_CONFIG  # type: ignore
+    )
+
+
+def is_state_message(message: AirbyteMessage) -> bool:
+    """
+    Determines whether the provided AirbyteMessage is a state message.
+
+    Parameters:
+        message (AirbyteMessage): The message to inspect.
+
+    Returns:
+        bool: True if the message's type is MessageType.STATE, False otherwise.
+    """
+    return message.type == MessageType.STATE  # type: ignore
+
+
+# -------
+# Handlers
+# -------
+
+
+def handle_current_slice(
+    current_slice_pages: List[StreamReadPages],
+    current_slice_descriptor: Optional[Dict[str, Any]] = None,
+    latest_state_message: Optional[Dict[str, Any]] = None,
+) -> StreamReadSlices:
+    """
+    Handles the current slice by packaging its pages, descriptor, and state into a StreamReadSlices instance.
+
+    Args:
+        current_slice_pages (List[StreamReadPages]): The pages to be included in the slice.
+        current_slice_descriptor (Optional[Dict[str, Any]]): Descriptor for the current slice, optional.
+        latest_state_message (Optional[Dict[str, Any]]): The latest state message, optional.
+
+    Returns:
+        StreamReadSlices: An object containing the current slice's pages, descriptor, and state.
+    """
+    return StreamReadSlices(
+        pages=current_slice_pages,
+        slice_descriptor=current_slice_descriptor,
+        state=[latest_state_message] if latest_state_message else [],
+    )
+
+
+def handle_current_page(
+    current_page_request: Optional[HttpRequest],
+    current_page_response: Optional[HttpResponse],
+    current_slice_pages: List[StreamReadPages],
+    current_page_records: List[Mapping[str, Any]],
+) -> tuple[None, None]:
+    """
+    Closes the current page by appending its request, response, and records
+    to the list of pages and clearing the current page records.
+
+    Args:
+        current_page_request (Optional[HttpRequest]): The HTTP request associated with the current page.
+        current_page_response (Optional[HttpResponse]): The HTTP response associated with the current page.
+        current_slice_pages (List[StreamReadPages]): A list to append the current page information.
+        current_page_records (List[Mapping[str, Any]]): The records of the current page to be cleared after processing.
+
+    Returns:
+        tuple[None, None]: A tuple indicating that no values are returned.
+    """
+
+    current_slice_pages.append(
+        StreamReadPages(
+            request=current_page_request,
+            response=current_page_response,
+            records=deepcopy(current_page_records),  # type: ignore [arg-type]
+        )
+    )
+    current_page_records.clear()
+
+    return None, None
+
+
+def handle_auxiliary_request(json_message: Dict[str, JsonType]) -> AuxiliaryRequest:
+    """
+    Parses the provided JSON message and constructs an AuxiliaryRequest object by extracting
+    relevant fields from nested dictionaries.
+
+    This function retrieves and validates the "airbyte_cdk", "stream", and "http" dictionaries
+    from the JSON message. It raises a ValueError if any of these are not of type dict. A title
+    is dynamically created by checking if the stream is a substream and then combining a prefix
+    with the "title" field from the "http" dictionary. The function also uses helper functions
+    to generate the request and response portions of the AuxiliaryRequest.
+
+    Parameters:
+        json_message (Dict[str, JsonType]): A dictionary representing the JSON log message containing
+        auxiliary request details.
+
+    Returns:
+        AuxiliaryRequest: An object containing the formatted title, description, request, and response
+        extracted from the JSON message.
+
+    Raises:
+        ValueError: If any of the "airbyte_cdk", "stream", or "http" fields is not a dictionary.
+    """
+    airbyte_cdk = json_message.get("airbyte_cdk", {})
+
+    if not isinstance(airbyte_cdk, dict):
+        raise ValueError(
+            f"Expected airbyte_cdk to be a dict, got {airbyte_cdk} of type {type(airbyte_cdk)}"
+        )
+
+    stream = airbyte_cdk.get("stream", {})
+
+    if not isinstance(stream, dict):
+        raise ValueError(f"Expected stream to be a dict, got {stream} of type {type(stream)}")
+
+    title_prefix = "Parent stream: " if stream.get("is_substream", False) else ""
+    http = json_message.get("http", {})
+
+    if not isinstance(http, dict):
+        raise ValueError(f"Expected http to be a dict, got {http} of type {type(http)}")
+
+    return AuxiliaryRequest(
+        title=title_prefix + str(http.get("title", None)),
+        description=str(http.get("description", None)),
+        request=create_request_from_log_message(json_message),
+        response=create_response_from_log_message(json_message),
+    )
+
+
+def handle_log_message(
+    message: AirbyteMessage,
+    json_message: Dict[str, JsonType] | None,
+    at_least_one_page_in_group: bool,
+    current_page_request: Optional[HttpRequest],
+    current_page_response: Optional[HttpResponse],
+) -> LOG_MESSAGES_OUTPUT_TYPE:
+    """
+    Process a log message by handling both HTTP-specific and auxiliary log entries.
+
+    Parameters:
+        message (AirbyteMessage): The original log message received.
+        json_message (Dict[str, JsonType] | None): A parsed JSON representation of the log message, if available.
+        at_least_one_page_in_group (bool): Indicates whether at least one page within the group has been processed.
+        current_page_request (Optional[HttpRequest]): The HTTP request object corresponding to the current page, if any.
+        current_page_response (Optional[HttpResponse]): The HTTP response object corresponding to the current page, if any.
+
+    Returns:
+        LOG_MESSAGES_OUTPUT_TYPE: A tuple containing:
+            - A boolean flag that determines whether the group contains at least one page.
+            - An updated HttpRequest for the current page (if applicable).
+            - An updated HttpResponse for the current page (if applicable).
+            - The auxiliary log message, which might be the original HTTP log or another log field.
+
+    Note:
+        If the parsed JSON message indicates an HTTP log and represents an auxiliary HTTP request,
+        the auxiliary log is handled via _handle_auxiliary_request. Otherwise, if the JSON log is a standard HTTP log,
+        the function updates the current page's request and response objects by generating them from the log message.
+    """
+    auxiliary_request = None
+    log_message = None
+
+    if json_message is not None and is_http_log(json_message):
+        if is_auxiliary_http_request(json_message):
+            auxiliary_request = handle_auxiliary_request(json_message)
+        else:
+            at_least_one_page_in_group = True
+            current_page_request = create_request_from_log_message(json_message)
+            current_page_response = create_response_from_log_message(json_message)
+    else:
+        log_message = message.log
+
+    return (
+        at_least_one_page_in_group,
+        current_page_request,
+        current_page_response,
+        auxiliary_request or log_message,
+    )
+
+
+def handle_record_message(
+    message: AirbyteMessage,
+    schema_inferrer: SchemaInferrer,
+    datetime_format_inferrer: DatetimeFormatInferrer,
+    records_count: int,
+    current_page_records: List[Mapping[str, Any]],
+) -> int:
+    """
+    Processes an Airbyte record message by updating the current batch and accumulating schema and datetime format information.
+
+    Parameters:
+        message (AirbyteMessage): The Airbyte message to process. Expected to have a 'type' attribute and, if of type RECORD, a 'record' attribute containing the record data.
+        schema_inferrer (SchemaInferrer): An instance responsible for inferring and accumulating schema details based on the record data.
+        datetime_format_inferrer (DatetimeFormatInferrer): An instance responsible for inferring and accumulating datetime format information from the record data.
+        records_count (int): The current count of processed records. This value is incremented if the message is a record.
+        current_page_records (List[Mapping[str, Any]]): A list where the data of processed record messages is accumulated.
+
+    Returns:
+        int: The updated record count after processing the message.
+    """
+    if message.type == MessageType.RECORD:
+        current_page_records.append(message.record.data)  # type: ignore
+        records_count += 1
+        schema_inferrer.accumulate(message.record)  # type: ignore
+        datetime_format_inferrer.accumulate(message.record)  # type: ignore
+
+    return records_count