airbyte_cdk.sources.declarative.migrations.legacy_to_per_partition_state_migration

  1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
  2
  3from typing import Any, Mapping, Union
  4
  5from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
  6from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
  7from airbyte_cdk.sources.declarative.models import (
  8    CustomPartitionRouter,
  9    DatetimeBasedCursor,
 10    SubstreamPartitionRouter,
 11    UnionPartitionRouter,
 12)
 13from airbyte_cdk.sources.declarative.models.declarative_component_schema import ParentStreamConfig
 14
 15
 16def _is_already_migrated(stream_state: Mapping[str, Any]) -> bool:
 17    return "states" in stream_state
 18
 19
 20class LegacyToPerPartitionStateMigration(StateMigration):
 21    """
 22    Transforms the input state for per-partitioned streams from the legacy format to the low-code format.
 23    The cursor field and partition ID fields are automatically extracted from the stream's DatetimebasedCursor and SubstreamPartitionRouter.
 24
 25    Example input state:
 26    {
 27    "13506132": {
 28      "last_changed": "2022-12-27T08:34:39+00:00"
 29    }
 30    Example output state:
 31    {
 32      "partition": {"id": "13506132"},
 33      "cursor": {"last_changed": "2022-12-27T08:34:39+00:00"}
 34    }
 35    """
 36
 37    def __init__(
 38        self,
 39        partition_router: Union[
 40            SubstreamPartitionRouter, UnionPartitionRouter, CustomPartitionRouter
 41        ],
 42        cursor: DatetimeBasedCursor,
 43        config: Mapping[str, Any],
 44        parameters: Mapping[str, Any],
 45    ):
 46        self._partition_router = partition_router
 47        self._cursor = cursor
 48        self._config = config
 49        self._parameters = parameters
 50        self._partition_key_field = InterpolatedString.create(
 51            self._get_partition_field(partition_router), parameters=self._parameters
 52        ).eval(self._config)
 53        self._cursor_field = InterpolatedString.create(
 54            self._cursor.cursor_field, parameters=self._parameters
 55        ).eval(self._config)
 56
 57    def _get_partition_field(
 58        self,
 59        partition_router: Union[
 60            SubstreamPartitionRouter, UnionPartitionRouter, CustomPartitionRouter
 61        ],
 62    ) -> str:
 63        if isinstance(partition_router, UnionPartitionRouter):
 64            return partition_router.partition_field
 65
 66        parent_stream_config = partition_router.parent_stream_configs[0]  # type: ignore # custom partition routers are expected to expose parent_stream_configs
 67
 68        # Retrieve the partition field with a condition, as properties are returned as a dictionary for custom components.
 69        partition_field = (
 70            parent_stream_config.partition_field
 71            if isinstance(parent_stream_config, ParentStreamConfig)
 72            else parent_stream_config.get("partition_field")  # type: ignore # See above comment on why parent_stream_config might be a dict
 73        )
 74
 75        return partition_field
 76
 77    def should_migrate(self, stream_state: Mapping[str, Any]) -> bool:
 78        if _is_already_migrated(stream_state):
 79            return False
 80
 81        # UnionPartitionRouter has no parent_stream_configs; its partitions are already
 82        # normalized to a single partition field so the parent stream check does not apply.
 83        if not isinstance(self._partition_router, UnionPartitionRouter):
 84            # There is exactly one parent stream
 85            number_of_parent_streams = len(self._partition_router.parent_stream_configs)  # type: ignore # custom partition will introduce this attribute if needed
 86            if number_of_parent_streams != 1:
 87                # There should be exactly one parent stream
 88                return False
 89        """
 90        The expected state format is
 91        "<parent_key_id>" : {
 92          "<cursor_field>" : "<cursor_value>"
 93        }
 94        """
 95        if not stream_state:
 96            return False
 97        for key, value in stream_state.items():
 98            # it is expected the internal value to be a dictionary according to docstring
 99            if not isinstance(value, dict):
100                return False
101            keys = list(value.keys())
102            if len(keys) != 1:
103                # The input partitioned state should only have one key
104                return False
105            if keys[0] != self._cursor_field:
106                # Unexpected key. Found {keys[0]}. Expected {self._cursor.cursor_field}
107                return False
108
109        return True
110
111    def migrate(self, stream_state: Mapping[str, Any]) -> Mapping[str, Any]:
112        states = [
113            {"partition": {self._partition_key_field: key}, "cursor": value}
114            for key, value in stream_state.items()
115        ]
116        return {"states": states}