20 lines
609 B
Python
20 lines
609 B
Python
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
|
|
from typing import Any, List
|
|
|
|
from airbyte_protocol.models import AirbyteStateMessage
|
|
|
|
|
|
class StateBuilder:
|
|
def __init__(self) -> None:
|
|
self._state: List[AirbyteStateMessage] = []
|
|
|
|
def with_stream_state(self, stream_name: str, state: Any) -> "StateBuilder":
|
|
self._state.append(
|
|
AirbyteStateMessage.parse_obj({"type": "STREAM", "stream": {"stream_state": state, "stream_descriptor": {"name": stream_name}}})
|
|
)
|
|
return self
|
|
|
|
def build(self) -> List[AirbyteStateMessage]:
|
|
return self._state
|