Feat: Upgrade to Pydantic 2.0 and CDK 2.0 (#291)

Author: aaronsteers

.github/workflows/python_pytest.yml

@@ -142,6 +142,10 @@ jobs:
       fail-fast: false
 
     runs-on: "${{ matrix.os }}-latest"
+    env:
+      # Enforce UTF-8 encoding so Windows runners don't fail inside the connector code.
+      # TODO: See if we can fully enforce this within PyAirbyte itself.
+      PYTHONIOENCODING: utf-8
     steps:
     # Common steps:
     - name: Checkout code

.github/workflows/test-pr-command.yml

@@ -63,8 +63,11 @@ jobs:
           Windows,
         ]
       fail-fast: false
-
     runs-on: "${{ matrix.os }}-latest"
+    env:
+      # Enforce UTF-8 encoding so Windows runners don't fail inside the connector code.
+      # TODO: See if we can fully enforce this within PyAirbyte itself.
+      PYTHONIOENCODING: utf-8
     steps:
 
     # Custom steps to fetch the PR and checkout the code:

airbyte/_executor.py

@@ -100,6 +100,7 @@ def _stream_from_subprocess(args: list[str]) -> Generator[Iterable[str], None, N
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
         universal_newlines=True,
+        encoding="utf-8",
     )
 
     def _stream_from_file(file: IO[str]) -> Generator[str, Any, None]:

airbyte/_future_cdk/record_processor.py

@@ -133,7 +133,7 @@ def _airbyte_messages_from_buffer(
         buffer: io.TextIOBase,
     ) -> Iterator[AirbyteMessage]:
         """Yield messages from a buffer."""
-        yield from (AirbyteMessage.parse_raw(line) for line in buffer)
+        yield from (AirbyteMessage.model_validate_json(line) for line in buffer)
 
     @final
     def process_input_stream(

airbyte/_future_cdk/state_providers.py

@@ -78,7 +78,10 @@ def to_state_input_file_text(self) -> str:
         return (
             "["
             + "\n, ".join(
-                [state_artifact.json() for state_artifact in (self._state_message_artifacts or [])]
+                [
+                    state_artifact.model_dump_json()
+                    for state_artifact in (self._state_message_artifacts or [])
+                ]
             )
             + "]"
         )

airbyte/_future_cdk/state_writers.py

@@ -36,4 +36,4 @@ def write_state(
         state_message: AirbyteStateMessage,
     ) -> None:
         """Save or 'write' a state artifact."""
-        print(state_message.json())
+        print(state_message.model_dump_json())

airbyte/_util/document_rendering.py

@@ -28,7 +28,7 @@ def _to_title_case(name: str, /) -> str:
 class CustomRenderingInstructions(BaseModel):
     """Instructions for rendering a stream's records as documents."""
 
-    title_property: Optional[str]
+    title_property: Optional[str] = None
     content_properties: list[str]
     frontmatter_properties: list[str]
     metadata_properties: list[str]
@@ -37,9 +37,9 @@ class CustomRenderingInstructions(BaseModel):
 class DocumentRenderer(BaseModel):
     """Instructions for rendering a stream's records as documents."""
 
-    title_property: Optional[str]
-    content_properties: Optional[list[str]]
-    metadata_properties: Optional[list[str]]
+    title_property: Optional[str] = None
+    content_properties: Optional[list[str]] = None
+    metadata_properties: Optional[list[str]] = None
     render_metadata: bool = False
 
     # TODO: Add primary key and cursor key support:

airbyte/caches/_state_backend.py

@@ -104,7 +104,7 @@ def write_state(
                     source_name=self.source_name,
                     stream_name=stream_name,
                     table_name=table_prefix + stream_name,
-                    state_json=state_message.json(),
+                    state_json=state_message.model_dump_json(),
                 )
             )
             session.commit()
@@ -170,7 +170,7 @@ def get_state_provider(
 
         return StaticInputState(
             from_state_messages=[
-                AirbyteStateMessage.parse_raw(state.state_json) for state in states
+                AirbyteStateMessage.model_validate_json(state.state_json) for state in states
             ]
         )
 

airbyte/caches/snowflake.py

@@ -31,6 +31,6 @@
 class SnowflakeCache(SnowflakeConfig, CacheBase):
     """Configuration for the Snowflake cache."""
 
-    dedupe_mode = RecordDedupeMode.APPEND
+    dedupe_mode: RecordDedupeMode = RecordDedupeMode.APPEND
 
     _sql_processor_class = PrivateAttr(default=SnowflakeSqlProcessor)

airbyte/secrets/base.py

@@ -6,11 +6,18 @@
 import json
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import cast
+from typing import TYPE_CHECKING, Any, cast
+
+from pydantic_core import CoreSchema, core_schema
 
 from airbyte import exceptions as exc
 
 
+if TYPE_CHECKING:
+    from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler, ValidationInfo
+    from pydantic.json_schema import JsonSchemaValue
+
+
 class SecretSourceEnum(str, Enum):
     ENV = "env"
     DOTENV = "dotenv"
@@ -65,6 +72,49 @@ def parse_json(self) -> dict:
                 },
             ) from None
 
+    # Pydantic compatibility
+
+    @classmethod
+    def validate(
+        cls,
+        v: Any,  # noqa: ANN401  # Must allow `Any` to match Pydantic signature
+        info: ValidationInfo,
+    ) -> SecretString:
+        """Validate the input value is valid as a secret string."""
+        _ = info  # Unused
+        if not isinstance(v, str):
+            raise exc.PyAirbyteInputError(
+                message="A valid `str` or `SecretString` object is required.",
+            )
+        return cls(v)
+
+    @classmethod
+    def __get_pydantic_core_schema__(  # noqa: PLW3201  # Pydantic dunder
+        cls,
+        source_type: Any,  # noqa: ANN401  # Must allow `Any` to match Pydantic signature
+        handler: GetCoreSchemaHandler,
+    ) -> CoreSchema:
+        return core_schema.with_info_after_validator_function(
+            function=cls.validate, schema=handler(str), field_name=handler.field_name
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(  # noqa: PLW3201  # Pydantic dunder method
+        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        """
+        Return a modified JSON schema for the secret string.
+
+        - `writeOnly=True` is the official way to prevent secrets from being exposed inadvertently.
+        - `Format=password` is a popular and readable convention to indicate the field is sensitive.
+        """
+        _ = _core_schema, handler  # Unused
+        return {
+            "type": "string",
+            "format": "password",
+            "writeOnly": True,
+        }
+
 
 class SecretManager(ABC):
     """Abstract base class for secret managers.