From e02cf31ac654045241d1ba48e84439a921bc57ab Mon Sep 17 00:00:00 2001
From: Aaron Steers <aj@airbyte.io>
Date: Sun, 10 Nov 2024 11:47:07 -0800
Subject: [PATCH] revert changes for serpyco data classes

---
 airbyte_cdk/models/airbyte_protocol.py | 68 +++++++++++++++-----------
 1 file changed, 39 insertions(+), 29 deletions(-)

diff --git a/airbyte_cdk/models/airbyte_protocol.py b/airbyte_cdk/models/airbyte_protocol.py
index 0b9ede19..e74e9195 100644
--- a/airbyte_cdk/models/airbyte_protocol.py
+++ b/airbyte_cdk/models/airbyte_protocol.py
@@ -1,29 +1,36 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-from __future__ import annotations
 
-from dataclasses import InitVar, dataclass
-from typing import TYPE_CHECKING, Annotated, Any
+# Serpyco uses type definitions at runtime during SerDes operations.
+# For this reason we have some exceptions to normal linting rules.
+# ruff: noqa: TCH001, TCH002, TCH003  # Don't auto-move imports to `TYPE_CHECKING` block.
+# ruff: noqa: F403    # Allow '*' import to shadow everything from protocols package.
 
-from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
+# Allow camelCase names (imported from java library)
+# ruff: noqa: N815
 
+from __future__ import annotations
 
-if TYPE_CHECKING:
-    from collections.abc import Mapping
+from collections.abc import Mapping
+from dataclasses import InitVar, dataclass
 
-    from serpyco_rs.metadata import Alias
+# Serpyco does not support the 3.10-style "|" operator.
+# ruff: noqa: UP007  # Allow deprecated `Union` and `Optional`
+from typing import Annotated, Any, Optional, Union
 
-    from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
+from serpyco_rs.metadata import Alias
 
+from airbyte_protocol_dataclasses import models
+from airbyte_protocol_dataclasses.models import *
 
-# ruff: noqa: F405  # ignore fuzzy import issues with 'import *'
+from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
 
 
 @dataclass
-class AirbyteStateBlob:
+class AirbyteStateBlob:  # noqa: PLW1641  # Should implement __hash__
     """A dataclass that dynamically sets attributes based on provided keyword arguments and positional arguments.
-    Used to "mimic" pydantic Basemodel with ConfigDict(extra='allow') option.
+    Used to "mimic" pydantic BaseModel with ConfigDict(extra='allow') option.
 
     The `AirbyteStateBlob` class allows for flexible instantiation by accepting any number of keyword arguments
     and positional arguments. These are used to dynamically update the instance's attributes. This class is useful
@@ -61,36 +68,39 @@ def __eq__(self, other: object) -> bool:
 # The following dataclasses have been redeclared to include the new version of AirbyteStateBlob
 @dataclass
 class AirbyteStreamState:
-    stream_descriptor: StreamDescriptor  # type: ignore [name-defined]
-    stream_state: AirbyteStateBlob | None = None
+    stream_descriptor: models.StreamDescriptor
+    stream_state: Optional[AirbyteStateBlob] = None
 
 
 @dataclass
 class AirbyteGlobalState:
     stream_states: list[AirbyteStreamState]
-    shared_state: AirbyteStateBlob | None = None
+    shared_state: Optional[AirbyteStateBlob] = None
 
 
 @dataclass
 class AirbyteStateMessage:
-    type: AirbyteStateType | None = None  # type: ignore [name-defined]
-    stream: AirbyteStreamState | None = None
-    global_: Annotated[AirbyteGlobalState | None, Alias("global")] = (
+    type: Optional[models.AirbyteStateType] = None
+    stream: Optional[models.AirbyteStreamState] = None
+    global_: Annotated[Optional[AirbyteGlobalState], Alias("global")] = (
         None  # "global" is a reserved keyword in python ⇒ Alias is used for (de-)serialization
     )
-    data: dict[str, Any] | None = None
-    sourceStats: AirbyteStateStats | None = None  # type: ignore [name-defined]
-    destinationStats: AirbyteStateStats | None = None  # type: ignore [name-defined]
+    data: Optional[dict[str, Any]] = None
+    sourceStats: Optional[models.AirbyteStateStats] = None
+    destinationStats: Optional[models.AirbyteStateStats] = None
 
 
 @dataclass
 class AirbyteMessage:
-    type: Type  # type: ignore [name-defined]
-    log: AirbyteLogMessage | None = None  # type: ignore [name-defined]
-    spec: ConnectorSpecification | None = None  # type: ignore [name-defined]
-    connectionStatus: AirbyteConnectionStatus | None = None  # type: ignore [name-defined]
-    catalog: AirbyteCatalog | None = None  # type: ignore [name-defined]
-    record: AirbyteFileTransferRecordMessage | AirbyteRecordMessage | None = None  # type: ignore [name-defined]
-    state: AirbyteStateMessage | None = None
-    trace: AirbyteTraceMessage | None = None  # type: ignore [name-defined]
-    control: AirbyteControlMessage | None = None  # type: ignore [name-defined]
+    type: models.Type
+    log: Optional[models.AirbyteLogMessage] = None
+    spec: Optional[models.ConnectorSpecification] = None
+    connectionStatus: Optional[models.AirbyteConnectionStatus] = None
+    catalog: Optional[models.AirbyteCatalog] = None
+
+    # These two differ from the original dataclasses:
+    record: Optional[Union[AirbyteFileTransferRecordMessage, models.AirbyteRecordMessage]] = None
+    state: Optional[AirbyteStateMessage] = None
+
+    trace: Optional[models.AirbyteTraceMessage] = None
+    control: Optional[models.AirbyteControlMessage] = None