feat(airbyte-cdk): replace pydantic BaseModel with dataclasses + serpyco-rs in protocol

[artem1205]

What

Resolve https://github.com/airbytehq/airbyte-internal-issues/issues/8945 by use of dataclasses and serpyco-rs for serialization

Following #44026

How

1. Dataclass models of protocol is used from, PR: feat: add version with dataclasses airbyte-protocol#91
2. Refactor to use dataclasses and corresponding Serializers from serpyco-rs

Tip

Example of serialization dataclass -> str:

orjson.dumps(AirbyteMessageSerializer.dump(message)).decode()

Review guide

1. airbyte-cdk/python/airbyte_cdk/models/airbyte_protocol.py
2. airbyte-cdk/python/airbyte_cdk/entrypoint.py

Note

All Serializers are in airbyte-cdk/python/airbyte_cdk/models/airbyte_protocol.py

Warning

Excel Parser parse_ method changed to to_json, because to_dict() method returns datetime values in pandas.Timestamp values, which are not serializable by orjson. See airbyte-cdk/python/airbyte_cdk/sources/file_based/file_types/excel_parser.py
To align with previous pydantic serialization logic, date_format="iso", date_unit="us" (microseconds) will be used, since pydantic cannot handle nanoseconds and strips seconds to 6 digits, i.e. microseconds)
Polars dataframe.to_dicts() has the same behaviour

Code Example

from pydantic import BaseModel
import pandas as pd
from typing import Dict, Any
class MyModel(BaseModel):
    timestamp: Dict[str, Any]
# Create an instance of the model
model_instance = MyModel(timestamp={'ss': pd.Timestamp('2023-08-27 10:00:00')})
# Serialize to JSON
print(model_instance.model_dump_json())
{"timestamp":{"ss":"2023-08-27T10:00:00"}}
model_instance = MyModel(timestamp={'ss': pd.Timestamp('2023-08-27 10:00:00.12')})
print(model_instance.model_dump_json())
{"timestamp":{"ss":"2023-08-27T10:00:00.120000"}}
model_instance = MyModel(timestamp={'ss': pd.Timestamp('2023-08-27 10:00:00.123456')})
print(model_instance.model_dump_json())
{"timestamp":{"ss":"2023-08-27T10:00:00.123456"}}
model_instance = MyModel(timestamp={'ss': pd.Timestamp('2023-08-27 10:00:00.123456789')}) # nanoseconds are missing
print(model_instance.model_dump_json())
{"timestamp":{"ss":"2023-08-27T10:00:00.123456"}}

Warning

Avro Parser _to_output_value will decode bytes to string, cause bytes are not serializable by orjson. See airbyte-cdk/python/airbyte_cdk/sources/file_based/file_types/avro_parser.py, Slack thread

User Impact

-- changes in protocol models only

Tip

It is recommended to import protocol classes not directly by import airbyte_protocol statement, but from aribyte_cdk.models package

Note

All Serializers have omit_none=True parameter that is applied recursively. Thus, all None values are excluded from output.
This is expected behaviour and does not break anything in protocol.

Tests

Test with docker run; 10_000_000 records; stream: dummy_records

/usr/bin/time -h docker run --rm -v $(pwd)/secrets:/secrets -v $(pwd)/integration_tests:/integration_tests airbyte/source-hardcoded-records:dev read --config /secrets/config.json --catalog /integration_tests/configured_catalog.json > /tmp/test.txt

10,000,000 records are about 1,770 MB.

TIME	real	user	sys	Mb/s
Before (0.0.6)	2m37.51s	3.83s	14.85s	11.27
After (dev)	1m33.95s (1.67x)	2.71s	7.53s	18.8 (1.67x)

---

Test with local platform (deployed with abctl)

24,000,000 records are about 4,270 MB.

Version	Platform sync time E2E	Mb/s
Before (0.0.6)	10m 43s	6.64
After (dev)	5m 42s	12.48 (1.88x)

Can this PR be safely reverted and rolled back?

• YES 💚
• NO ❌

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Skipped Deployment

Name	Status	Preview	Comments	Updated (UTC)
airbyte-docs	⬜️ Ignored (Inspect)	Visit Preview		Aug 30, 2024 9:29pm

[girarda]

looks great just a few questions!

[girarda]

does the iterable really include Nones? This sounds difficult to work with

[artem1205]

mypy raises an error due to start_next_partition_generator, that can return None
I did not really dig into this class, only fixed mypy. I suppose we'll refactor this much during the next project: concurrency for Low-Code

airbyte/airbyte-cdk/python/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

Lines 160 to 179 in feded1f

	def start_next_partition_generator(self) -> Optional[AirbyteMessage]:
	"""
	Start the next partition generator.
	1. Pop the next stream to read from
	2. Submit the partition generator to the thread pool manager
	3. Add the stream to the list of streams currently generating partitions
	4. Return a stream status message
	"""
	if self._stream_instances_to_start_partition_generation:
	stream = self._stream_instances_to_start_partition_generation.pop(0)
	self._thread_pool_manager.submit(self._partition_enqueuer.generate_partitions, stream)
	self._streams_currently_generating_partitions.append(stream.name)
	self._logger.info(f"Marking stream {stream.name} as STARTED")
	self._logger.info(f"Syncing stream: {stream.name} ")
	return stream_status_as_airbyte_message(
	stream.as_airbyte_stream(),
	AirbyteStreamStatus.STARTED,
	)
	else:
	return None

[girarda]

if it's just for mypy, let's add a type ignore. Otherwise, the type will indicate that it is valid to yield None, which might not be the case?

[artem1205]

ignored

[girarda]

@artem1205 do you know if this will break any connectors?

If yes, I'm fine with upgrading to 5.0.0, but if not, I think it's preferable to stay on 4.X because of how difficult it still is to keep connectors on the latest major version.

[artem1205]

CAT tests were ✅ for source-:

• hardcoded-records
• zendesk-support
• shopify

These changes will not break anything, unless airbyte_protocol is imported directly (not as a part of airbyte_cdk.models)

I'll change the versioning to 4.x then.

[girarda]

can you check if any connector directly import airbyte_protocol? It might be worth fixing them first to avoid the need to flag it as a breaking change

[artem1205]

With grep and excluding test folders got 18 connectors:

grep -r --include="*.py" --exclude-dir="unit_tests" --exclude-dir="integration_tests" "from airbyte_protocol" ./airbyte-integrations/connectors/

1. /source-linkedin-ads/
2. /source-shopify/
3. /source-amazon-ads/
4. /source-mixpanel/
5. /source-amazon-seller-partner/
6. /source-bing-ads/
7. /source-google-ads/
8. /source-salesforce/
9. /source-stripe/
10. /source-mailchimp/
11. /destination-snowflake-cortex/
12. /source-airtable/
13. /source-zendesk-support/
14. /source-marketo/
15. /destination-pinecone/
16. /source-pinterest/
17. /source-github/

Do you want me to fix all of them before merging this PR?

[artem1205]

/approve-regression-tests

Check job output.

✅ Approving regression tests