2ac5248387
Co-authored-by: brianjlai <brian.lai@airbyte.io> Co-authored-by: Brian Lai <51336873+brianjlai@users.noreply.github.com>
18 lines
953 B
Python
18 lines
953 B
Python
# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
|
|
|
|
from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
|
|
from airbyte_protocol.models import AirbyteMessage, Type
|
|
|
|
|
|
def get_stream_descriptor(message: AirbyteMessage) -> HashableStreamDescriptor:
|
|
if message.type == Type.RECORD:
|
|
return HashableStreamDescriptor(name=message.record.stream, namespace=message.record.namespace)
|
|
elif message.type == Type.STATE:
|
|
if not message.state.stream or not message.state.stream.stream_descriptor:
|
|
raise ValueError("State message was not in per-stream state format, which is required for record counts.")
|
|
return HashableStreamDescriptor(
|
|
name=message.state.stream.stream_descriptor.name, namespace=message.state.stream.stream_descriptor.namespace
|
|
)
|
|
else:
|
|
raise NotImplementedError(f"get_stream_descriptor is not implemented for message type '{message.type}'.")
|