airbyte_cdk.models.airbyte_protocol

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5from dataclasses import InitVar, dataclass
  6from typing import Annotated, Any, Dict, List, Mapping, Optional, Union
  7
  8from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
  9from serpyco_rs.metadata import Alias
 10
 11# ruff: noqa: F405  # ignore fuzzy import issues with 'import *'
 12
 13
 14@dataclass
 15class AirbyteStateBlob:
 16    """
 17    A dataclass that dynamically sets attributes based on provided keyword arguments and positional arguments.
 18    Used to "mimic" pydantic Basemodel with ConfigDict(extra='allow') option.
 19
 20    The `AirbyteStateBlob` class allows for flexible instantiation by accepting any number of keyword arguments
 21    and positional arguments. These are used to dynamically update the instance's attributes. This class is useful
 22    in scenarios where the attributes of an object are not known until runtime and need to be set dynamically.
 23
 24    Attributes:
 25        kwargs (InitVar[Mapping[str, Any]]): A dictionary of keyword arguments used to set attributes dynamically.
 26
 27    Methods:
 28        __init__(*args: Any, **kwargs: Any) -> None:
 29            Initializes the `AirbyteStateBlob` by setting attributes from the provided arguments.
 30
 31        __eq__(other: object) -> bool:
 32            Checks equality between two `AirbyteStateBlob` instances based on their internal dictionaries.
 33            Returns `False` if the other object is not an instance of `AirbyteStateBlob`.
 34    """
 35
 36    kwargs: InitVar[Mapping[str, Any]]
 37
 38    def __init__(self, *args: Any, **kwargs: Any) -> None:
 39        # Set any attribute passed in through kwargs
 40        for arg in args:
 41            self.__dict__.update(arg)
 42        for key, value in kwargs.items():
 43            setattr(self, key, value)
 44
 45    def __eq__(self, other: object) -> bool:
 46        return (
 47            False
 48            if not isinstance(other, AirbyteStateBlob)
 49            else bool(self.__dict__ == other.__dict__)
 50        )
 51
 52
 53# The following dataclasses have been redeclared to include the new version of AirbyteStateBlob
 54@dataclass
 55class AirbyteStreamState:
 56    stream_descriptor: StreamDescriptor  # type: ignore [name-defined]
 57    stream_state: Optional[AirbyteStateBlob] = None
 58
 59
 60@dataclass
 61class AirbyteGlobalState:
 62    stream_states: List[AirbyteStreamState]
 63    shared_state: Optional[AirbyteStateBlob] = None
 64
 65
 66@dataclass
 67class AirbyteStateMessage:
 68    type: Optional[AirbyteStateType] = None  # type: ignore [name-defined]
 69    stream: Optional[AirbyteStreamState] = None
 70    global_: Annotated[AirbyteGlobalState | None, Alias("global")] = (
 71        None  # "global" is a reserved keyword in python ⇒ Alias is used for (de-)serialization
 72    )
 73    data: Optional[Dict[str, Any]] = None
 74    sourceStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]
 75    destinationStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]
 76
 77
 78# The following dataclasses have been redeclared to include scopes, optional_scopes,
 79# and scopes_join_strategy fields that are used by declarative OAuth connectors.
 80# The protocol model (OauthConnectorInputSpecification) does not include these fields,
 81# so serpyco_rs silently drops them during deserialization. By overriding the model here
 82# and cascading through OAuthConfigSpecification → AdvancedAuth → ConnectorSpecification,
 83# the fields are preserved in the connector's spec output.
 84# This follows the same override pattern used above for AirbyteStateBlob.
 85@dataclass
 86class OauthConnectorInputSpecification:
 87    consent_url: str
 88    access_token_url: str
 89    scope: Optional[str] = None
 90    scopes: Optional[List[Dict[str, Any]]] = None
 91    optional_scopes: Optional[List[Dict[str, Any]]] = None
 92    # Stored as str (not ScopesJoinStrategy enum) because spec.py converts the enum
 93    # to its .value before serialization. The protocol layer only sees plain strings.
 94    scopes_join_strategy: Optional[str] = None
 95    access_token_headers: Optional[Dict[str, Any]] = None
 96    access_token_params: Optional[Dict[str, Any]] = None
 97    extract_output: Optional[List[str]] = None
 98    state: Optional[State] = None  # type: ignore [name-defined]
 99    client_id_key: Optional[str] = None
100    client_secret_key: Optional[str] = None
101    scope_key: Optional[str] = None
102    state_key: Optional[str] = None
103    auth_code_key: Optional[str] = None
104    redirect_uri_key: Optional[str] = None
105    token_expiry_key: Optional[str] = None
106
107
108@dataclass
109class OAuthConfigSpecification:
110    oauth_user_input_from_connector_config_specification: Optional[Dict[str, Any]] = None
111    oauth_connector_input_specification: Optional[OauthConnectorInputSpecification] = None
112    complete_oauth_output_specification: Optional[Dict[str, Any]] = None
113    complete_oauth_server_input_specification: Optional[Dict[str, Any]] = None
114    complete_oauth_server_output_specification: Optional[Dict[str, Any]] = None
115
116
117@dataclass
118class AdvancedAuth:
119    auth_flow_type: Optional[AuthFlowType] = None  # type: ignore [name-defined]
120    predicate_key: Optional[List[str]] = None
121    predicate_value: Optional[str] = None
122    oauth_config_specification: Optional[OAuthConfigSpecification] = None
123
124
125@dataclass
126class ConnectorSpecification:
127    connectionSpecification: Dict[str, Any]
128    documentationUrl: Optional[str] = None
129    changelogUrl: Optional[str] = None
130    supportsIncremental: Optional[bool] = None
131    supportsNormalization: Optional[bool] = False
132    supportsDBT: Optional[bool] = False
133    supported_destination_sync_modes: Optional[List[DestinationSyncMode]] = None  # type: ignore [name-defined]
134    authSpecification: Optional[AuthSpecification] = None  # type: ignore [name-defined]
135    advanced_auth: Optional[AdvancedAuth] = None
136    protocol_version: Optional[str] = None
137
138
139@dataclass
140class AirbyteMessage:
141    type: Type  # type: ignore [name-defined]
142    log: Optional[AirbyteLogMessage] = None  # type: ignore [name-defined]
143    spec: Optional[ConnectorSpecification] = None
144    connectionStatus: Optional[AirbyteConnectionStatus] = None  # type: ignore [name-defined]
145    catalog: Optional[AirbyteCatalog] = None  # type: ignore [name-defined]
146    record: Optional[AirbyteRecordMessage] = None  # type: ignore [name-defined]
147    state: Optional[AirbyteStateMessage] = None
148    trace: Optional[AirbyteTraceMessage] = None  # type: ignore [name-defined]
149    control: Optional[AirbyteControlMessage] = None  # type: ignore [name-defined]