airbyte-cdk 6.5.3rc2__py3-none-any.whl → 6.6.0__py3-none-any.whl

airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py

@@ -4,7 +4,11 @@ from typing import Any, Mapping
 
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
-from airbyte_cdk.sources.declarative.models import DatetimeBasedCursor, SubstreamPartitionRouter
+from airbyte_cdk.sources.declarative.models import (
+    DatetimeBasedCursor,
+    SubstreamPartitionRouter,
+    CustomIncrementalSync,
+)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import ParentStreamConfig
 
 
@@ -32,7 +36,7 @@ class LegacyToPerPartitionStateMigration(StateMigration):
     def __init__(
         self,
         partition_router: SubstreamPartitionRouter,
-        cursor: DatetimeBasedCursor,
+        cursor: CustomIncrementalSync | DatetimeBasedCursor,
         config: Mapping[str, Any],
         parameters: Mapping[str, Any],
     ):
@@ -43,7 +47,9 @@ class LegacyToPerPartitionStateMigration(StateMigration):
         self._partition_key_field = InterpolatedString.create(
             self._get_partition_field(self._partition_router), parameters=self._parameters
         ).eval(self._config)
-        self._cursor_field = InterpolatedString.create(self._cursor.cursor_field, parameters=self._parameters).eval(self._config)
+        self._cursor_field = InterpolatedString.create(
+            self._cursor.cursor_field, parameters=self._parameters
+        ).eval(self._config)
 
     def _get_partition_field(self, partition_router: SubstreamPartitionRouter) -> str:
         parent_stream_config = partition_router.parent_stream_configs[0]
@@ -62,7 +68,7 @@ class LegacyToPerPartitionStateMigration(StateMigration):
             return False
 
         # There is exactly one parent stream
-        number_of_parent_streams = len(self._partition_router.parent_stream_configs)
+        number_of_parent_streams = len(self._partition_router.parent_stream_configs)  # type: ignore # custom partition will introduce this attribute if needed
         if number_of_parent_streams != 1:
             # There should be exactly one parent stream
             return False
@@ -85,5 +91,8 @@ class LegacyToPerPartitionStateMigration(StateMigration):
         return True
 
     def migrate(self, stream_state: Mapping[str, Any]) -> Mapping[str, Any]:
-        states = [{"partition": {self._partition_key_field: key}, "cursor": value} for key, value in stream_state.items()]
+        states = [
+            {"partition": {self._partition_key_field: key}, "cursor": value}
+            for key, value in stream_state.items()
+        ]
         return {"states": states}