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