def from_proto(data_source):
"""
Convert data source config in FeatureTable spec to a DataSource class object.
"""
if data_source.file_options.file_format and data_source.file_options.file_url:
data_source_obj = FileSource(
field_mapping=data_source.field_mapping,
file_format=FileFormat.from_proto(data_source.file_options.file_format),
file_url=data_source.file_options.file_url,
event_timestamp_column=data_source.event_timestamp_column,
created_timestamp_column=data_source.created_timestamp_column,
date_partition_column=data_source.date_partition_column,
)
elif data_source.bigquery_options.table_ref:
data_source_obj = BigQuerySource(
field_mapping=data_source.field_mapping,
table_ref=data_source.bigquery_options.table_ref,
event_timestamp_column=data_source.event_timestamp_column,
created_timestamp_column=data_source.created_timestamp_column,
date_partition_column=data_source.date_partition_column,
)
elif (
data_source.kafka_options.bootstrap_servers
and data_source.kafka_options.topic
and data_source.kafka_options.message_format
):
data_source_obj = KafkaSource(
field_mapping=data_source.field_mapping,
bootstrap_servers=data_source.kafka_options.bootstrap_servers,
message_format=StreamFormat.from_proto(
data_source.kafka_options.message_format
),
topic=data_source.kafka_options.topic,
event_timestamp_column=data_source.event_timestamp_column,
created_timestamp_column=data_source.created_timestamp_column,
date_partition_column=data_source.date_partition_column,
)
elif (
data_source.kinesis_options.record_format
and data_source.kinesis_options.region
and data_source.kinesis_options.stream_name
):
data_source_obj = KinesisSource(
field_mapping=data_source.field_mapping,
record_format=StreamFormat.from_proto(
data_source.kinesis_options.record_format
),
region=data_source.kinesis_options.region,
stream_name=data_source.kinesis_options.stream_name,
event_timestamp_column=data_source.event_timestamp_column,
created_timestamp_column=data_source.created_timestamp_column,
date_partition_column=data_source.date_partition_column,
)
else:
raise ValueError("Could not identify the source type being added")
return data_source_obj
def from_proto(data_source: DataSourceProto):
return KinesisSource(
field_mapping=dict(data_source.field_mapping),
record_format=StreamFormat.from_proto(
data_source.kinesis_options.record_format),
region=data_source.kinesis_options.region,
stream_name=data_source.kinesis_options.stream_name,
event_timestamp_column=data_source.event_timestamp_column,
created_timestamp_column=data_source.created_timestamp_column,
date_partition_column=data_source.date_partition_column,
)
def from_proto(data_source: DataSourceProto):
return KafkaSource(
field_mapping=dict(data_source.field_mapping),
bootstrap_servers=data_source.kafka_options.bootstrap_servers,
message_format=StreamFormat.from_proto(
data_source.kafka_options.message_format),
topic=data_source.kafka_options.topic,
event_timestamp_column=data_source.event_timestamp_column,
created_timestamp_column=data_source.created_timestamp_column,
date_partition_column=data_source.date_partition_column,
)
def from_proto(data_source: DataSourceProto):
return KinesisSource(
name=data_source.name,
event_timestamp_column=data_source.timestamp_field,
field_mapping=dict(data_source.field_mapping),
record_format=StreamFormat.from_proto(
data_source.kinesis_options.record_format),
region=data_source.kinesis_options.region,
stream_name=data_source.kinesis_options.stream_name,
created_timestamp_column=data_source.created_timestamp_column,
timestamp_field=data_source.timestamp_field,
date_partition_column=data_source.date_partition_column,
description=data_source.description,
tags=dict(data_source.tags),
owner=data_source.owner,
batch_source=DataSource.from_proto(data_source.batch_source),
)
def from_proto(data_source: DataSourceProto):
return KafkaSource(
name=data_source.name,
event_timestamp_column=data_source.timestamp_field,
field_mapping=dict(data_source.field_mapping),
bootstrap_servers=data_source.kafka_options.bootstrap_servers,
message_format=StreamFormat.from_proto(
data_source.kafka_options.message_format),
topic=data_source.kafka_options.topic,
created_timestamp_column=data_source.created_timestamp_column,
timestamp_field=data_source.timestamp_field,
date_partition_column=data_source.date_partition_column,
description=data_source.description,
tags=dict(data_source.tags),
owner=data_source.owner,
batch_source=DataSource.from_proto(data_source.batch_source),
)
def from_proto(cls, kinesis_options_proto: DataSourceProto.KinesisOptions):
"""
Creates a KinesisOptions from a protobuf representation of a kinesis option
Args:
kinesis_options_proto: A protobuf representation of a DataSource
Returns:
Returns a KinesisOptions object based on the kinesis_options protobuf
"""
kinesis_options = cls(
record_format=StreamFormat.from_proto(kinesis_options_proto.record_format),
region=kinesis_options_proto.region,
stream_name=kinesis_options_proto.stream_name,
)
return kinesis_options
def from_proto(cls, kafka_options_proto: DataSourceProto.KafkaOptions):
"""
Creates a KafkaOptions from a protobuf representation of a kafka option
Args:
kafka_options_proto: A protobuf representation of a DataSource
Returns:
Returns a BigQueryOptions object based on the kafka_options protobuf
"""
kafka_options = cls(
bootstrap_servers=kafka_options_proto.bootstrap_servers,
message_format=StreamFormat.from_proto(kafka_options_proto.message_format),
topic=kafka_options_proto.topic,
)
return kafka_options