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

.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.

airbyte/sources/base.py

@@ -174,10 +174,10 @@ def _discover(self) -> AirbyteCatalog:
         """Call discover on the connector.
 
         This involves the following steps:
-        * Write the config to a temporary file
-        * execute the connector with discover --config <config_file>
-        * Listen to the messages and return the first AirbyteCatalog that comes along.
-        * Make sure the subprocess is killed when the function returns.
+        - Write the config to a temporary file
+        - execute the connector with discover --config <config_file>
+        - Listen to the messages and return the first AirbyteCatalog that comes along.
+        - Make sure the subprocess is killed when the function returns.
         """
         with as_temp_files([self._config]) as [config_file]:
             for msg in self._execute(["discover", "--config", config_file]):
@@ -298,7 +298,7 @@ def _yaml_spec(self) -> str:
         for each connector.
         """
         spec_obj: ConnectorSpecification = self._get_spec()
-        spec_dict = spec_obj.dict(exclude_unset=True)
+        spec_dict: dict[str, Any] = spec_obj.model_dump(exclude_unset=True)
         # convert to a yaml string
         return yaml.dump(spec_dict)
 
@@ -537,7 +537,7 @@ def _read_with_catalog(
         with as_temp_files(
             [
                 self._config,
-                catalog.json(),
+                catalog.model_dump_json(),
                 state.to_state_input_file_text() if state else "[]",
             ]
         ) as [
@@ -579,7 +579,7 @@ def _execute(self, args: list[str]) -> Iterator[AirbyteMessage]:
             self._last_log_messages = []
             for line in self.executor.execute(args):
                 try:
-                    message = AirbyteMessage.parse_raw(line)
+                    message: AirbyteMessage = AirbyteMessage.model_validate_json(json_data=line)
                     if message.type is Type.RECORD:
                         self._processed_records += 1
                     if message.type == Type.LOG:

airbyte/sources/declarative.py

@@ -4,9 +4,12 @@
 from __future__ import annotations
 
 import json
+import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, cast
 
+import pydantic
+
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 
@@ -19,6 +22,18 @@
     from collections.abc import Iterator
 
 
+def _suppress_cdk_pydantic_deprecation_warnings() -> None:
+    """Suppress deprecation warnings from Pydantic in the CDK.
+
+    CDK has deprecated uses of `json()` and `parse_obj()`, and we don't want users
+    to see these warnings.
+    """
+    warnings.filterwarnings(
+        "ignore",
+        category=pydantic.warnings.PydanticDeprecatedSince20,
+    )
+
+
 class DeclarativeExecutor(Executor):
     """An executor for declarative sources."""
 
@@ -32,6 +47,7 @@ def __init__(
         - If `manifest` is a string, it will be parsed as an HTTP path.
         - If `manifest` is a dict, it will be used as is.
         """
+        _suppress_cdk_pydantic_deprecation_warnings()
         self._manifest_dict: dict
         if isinstance(manifest, Path):
             self._manifest_dict = cast(dict, json.loads(manifest.read_text()))
@@ -91,6 +107,8 @@ def __init__(
             manifest: The manifest for the declarative source. This can be a path to a yaml file, a
             yaml string, or a dict.
         """
+        _suppress_cdk_pydantic_deprecation_warnings()
+
         # TODO: Conform manifest to a dict or str (TBD)
         self.manifest = manifest
 

airbyte/sources/registry.py

@@ -75,6 +75,7 @@
     "source-salesloft",
     "source-slack",
     "source-surveymonkey",
+    "source-tiktok-marketing",
     "source-the-guardian-api",
     "source-trello",
     "source-typeform",

examples/run_faker.py

@@ -13,12 +13,9 @@
 import airbyte as ab
 
 
-SCALE = 50_000  # Number of records to generate between users and purchases.
+SCALE = 200_000  # Number of records to generate between users and purchases.
 FORCE_FULL_REFRESH = True  # Whether to force a full refresh on the source.
 
-# This is a dummy secret, just to test functionality.
-DUMMY_SECRET = ab.get_secret("DUMMY_SECRET")
-
 
 print("Initializing cache...")
 cache = ab.get_default_cache()

examples/run_pokeapi.py

@@ -11,12 +11,13 @@
 from __future__ import annotations
 
 import airbyte as ab
+from airbyte.experimental import get_source
 
 
-source = ab.get_source(
+source = get_source(
     "source-pokeapi",
     config={"pokemon_name": "bulbasaur"},
-    install_if_missing=True,
+    source_manifest=True,
 )
 source.check()