26 lines
705 B
Python
26 lines
705 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
|