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()