def get_payload_from_typing_hints(self, values):
TestSchema = typing.NamedTuple('TestSchema', [
('integer_example', int),
('string_example', unicode),
('list_of_strings', typing.List[unicode]),
('optional_kv', typing.Optional[typing.Tuple[unicode, float]]),
('optional_integer', typing.Optional[int]),
])
builder = NamedTupleBasedPayloadBuilder(TestSchema(**values))
return builder.build()
def get_payload_from_typing_hints(self, values):
TestSchema = typing.NamedTuple('TestSchema', [
('integer_example', int),
('boolean', bool),
('string_example', str),
('list_of_strings', typing.List[str]),
('mapping', typing.Mapping[str, float]),
('optional_integer', typing.Optional[int]),
])
builder = NamedTupleBasedPayloadBuilder(TestSchema(**values))
return builder.build()
def __init__(
self,
subscription_path,
min_bundle_timeout=None,
deduplicate=None,
expansion_service=None,
):
"""
Initializes a read operation from Pub/Sub Lite, returning the serialized
bytes of SequencedMessage protos.
Args:
subscription_path: A Pub/Sub Lite Subscription path.
min_bundle_timeout: The minimum wall time to pass before allowing
bundle closure. Setting this to too small of a value will result in
increased compute costs and lower throughput per byte. Immediate
timeouts (0) may be useful for testing.
deduplicate: Whether to deduplicate messages based on the value of
the 'x-goog-pubsublite-dataflow-uuid' attribute.
"""
if min_bundle_timeout is None:
min_bundle_timeout = 60 * 1000
if deduplicate is None:
deduplicate = False
if expansion_service is None:
expansion_service = _default_io_expansion_service()
super().__init__(
'beam:transform:org.apache.beam:pubsublite_read:v1',
NamedTupleBasedPayloadBuilder(
_ReadSchema(subscription_path=subscription_path,
min_bundle_timeout=min_bundle_timeout,
deduplicate=deduplicate)), expansion_service)
def expand(self, pbegin):
return (pbegin
| 'User data mapper' >> beam.Map(
self.user_data_mapper).with_output_types(List[bytes])
| ExternalTransform(self.URN,
NamedTupleBasedPayloadBuilder(self.params),
self.expansion_service))
def __init__(
self,
driver_class_name,
jdbc_url,
username,
password,
statement,
connection_properties=None,
connection_init_sqls=None,
expansion_service=None,
):
super(WriteToJdbc, self).__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
WriteToJdbcSchema(
driver_class_name=driver_class_name,
jdbc_url=jdbc_url,
username=username,
password=password,
statement=statement,
connection_properties=connection_properties,
connection_init_sqls=connection_init_sqls,
),
),
expansion_service or default_io_expansion_service(),
)
def __init__(
self,
topic_path,
add_uuids=None,
expansion_service=None,
):
"""
Initializes a write operation to Pub/Sub Lite, writing the serialized bytes
of PubSubMessage protos.
Args:
topic_path: A Pub/Sub Lite Topic path.
add_uuids: Whether to add uuids to the 'x-goog-pubsublite-dataflow-uuid'
uuid attribute.
"""
if add_uuids is None:
add_uuids = False
if expansion_service is None:
expansion_service = _default_io_expansion_service()
super().__init__(
'beam:transform:org.apache.beam:pubsublite_write:v1',
NamedTupleBasedPayloadBuilder(
_WriteSchema(
topic_path=topic_path,
add_uuids=add_uuids,
)),
expansion_service)
def __init__(
self,
subscription_path,
deduplicate=None,
expansion_service=None,
):
"""
Initializes a read operation from Pub/Sub Lite, returning the serialized
bytes of SequencedMessage protos.
Args:
subscription_path: A Pub/Sub Lite Subscription path.
deduplicate: Whether to deduplicate messages based on the value of
the 'x-goog-pubsublite-dataflow-uuid' attribute.
"""
if deduplicate is None:
deduplicate = False
if expansion_service is None:
expansion_service = _default_io_expansion_service()
super().__init__(
'beam:transform:org.apache.beam:pubsublite_read:v1',
NamedTupleBasedPayloadBuilder(
_ReadSchema(
subscription_path=subscription_path, deduplicate=deduplicate)),
expansion_service)
def __init__(self,
consumer_config,
topics,
key_deserializer=byte_array_deserializer,
value_deserializer=byte_array_deserializer,
expansion_service=None):
"""
Initializes a read operation from Kafka.
:param consumer_config: A dictionary containing the consumer configuration.
:param topics: A list of topic strings.
:param key_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's key, e.g.
'org.apache.kafka.common.
serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.
serialization.ByteArrayDeserializer'.
:param value_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's value, e.g.
'org.apache.kafka.common.
serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.
serialization.ByteArrayDeserializer'.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
super(ReadFromKafka, self).__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
ReadFromKafkaSchema(
consumer_config=list(consumer_config.items()),
topics=topics,
key_deserializer=key_deserializer,
value_deserializer=value_deserializer,
)), expansion_service or default_io_expansion_service())
def __init__(self,
producer_config,
topic,
key_serializer=byte_array_serializer,
value_serializer=byte_array_serializer,
expansion_service=None):
"""
Initializes a write operation to Kafka.
:param producer_config: A dictionary containing the producer configuration.
:param topic: A Kafka topic name.
:param key_deserializer: A fully-qualified Java class name of a Kafka
Serializer for the topic's key, e.g.
'org.apache.kafka.common.serialization.LongSerializer'.
Default: 'org.apache.kafka.common.serialization.ByteArraySerializer'.
:param value_deserializer: A fully-qualified Java class name of a Kafka
Serializer for the topic's value, e.g.
'org.apache.kafka.common.serialization.LongSerializer'.
Default: 'org.apache.kafka.common.serialization.ByteArraySerializer'.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
WriteToKafkaSchema(
producer_config=producer_config,
topic=topic,
key_serializer=key_serializer,
value_serializer=value_serializer,
)), expansion_service or default_io_expansion_service())
def __init__(self, query, dialect=None):
super(SqlTransform, self).__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
SqlTransformSchema(query=query, dialect=dialect)),
BeamJarExpansionService(
':sdks:java:extensions:sql:expansion-service:shadowJar'))
def __init__(
self,
table_name,
driver_class_name,
jdbc_url,
username,
password,
statement=None,
connection_properties=None,
connection_init_sqls=None,
expansion_service=None,
classpath=None,
):
"""
Initializes a write operation to Jdbc.
:param driver_class_name: name of the jdbc driver class
:param jdbc_url: full jdbc url to the database.
:param username: database username
:param password: database password
:param statement: sql statement to be executed
:param connection_properties: properties of the jdbc connection
passed as string with format
[propertyName=property;]*
:param connection_init_sqls: required only for MySql and MariaDB.
passed as list of strings
:param expansion_service: The address (host:port) of the ExpansionService.
:param classpath: A list of JARs or Java packages to include in the
classpath for the expansion service. This option is
usually needed for `jdbc` to include extra JDBC driver
packages.
The packages can be in these three formats: (1) A local
file, (2) A URL, (3) A gradle-style identifier of a Maven
package (e.g. "org.postgresql:postgresql:42.3.1").
By default, this argument includes a Postgres SQL JDBC
driver.
"""
classpath = classpath or DEFAULT_JDBC_CLASSPATH
super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
JdbcConfigSchema(
location=table_name,
config=RowCoder(
typing_to_runner_api(Config).row_type.schema).encode(
Config(
driver_class_name=driver_class_name,
jdbc_url=jdbc_url,
username=username,
password=password,
connection_properties=connection_properties,
connection_init_sqls=connection_init_sqls,
write_statement=statement,
read_query=None,
fetch_size=None,
output_parallelization=None,
))), ),
expansion_service or default_io_expansion_service(classpath),
)
def __init__(
self,
consumer_config,
topics,
key_deserializer=byte_array_deserializer,
value_deserializer=byte_array_deserializer,
start_read_time=None,
max_num_records=None,
max_read_time=None,
commit_offset_in_finalize=False,
timestamp_policy=processing_time_policy,
expansion_service=None,
):
"""
Initializes a read operation from Kafka.
:param consumer_config: A dictionary containing the consumer configuration.
:param topics: A list of topic strings.
:param key_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's key, e.g.
'org.apache.kafka.common.serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
:param value_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's value, e.g.
'org.apache.kafka.common.serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
:param start_read_time: Use timestamp to set up start offset in milliseconds
epoch.
:param max_num_records: Maximum amount of records to be read. Mainly used
for tests and demo applications.
:param max_read_time: Maximum amount of time in seconds the transform
executes. Mainly used for tests and demo applications.
:param commit_offset_in_finalize: Whether to commit offsets when finalizing.
:param timestamp_policy: The built-in timestamp policy which is used for
extracting timestamp from KafkaRecord.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
if timestamp_policy not in [
ReadFromKafka.processing_time_policy,
ReadFromKafka.create_time_policy, ReadFromKafka.log_append_time
]:
raise ValueError('timestamp_policy should be one of '
'[ProcessingTime, CreateTime, LogAppendTime]')
super(ReadFromKafka, self).__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
ReadFromKafkaSchema(
consumer_config=consumer_config,
topics=topics,
key_deserializer=key_deserializer,
value_deserializer=value_deserializer,
max_num_records=max_num_records,
max_read_time=max_read_time,
start_read_time=start_read_time,
commit_offset_in_finalize=commit_offset_in_finalize,
timestamp_policy=timestamp_policy)), expansion_service
or default_io_expansion_service())
def expand(self, pbegin):
return (
pbegin
| ExternalTransform(
self.URN,
NamedTupleBasedPayloadBuilder(self.params),
self.expansion_service,
)
| 'CSV to array mapper' >> beam.Map(lambda csv: csv.split(b','))
| 'CSV mapper' >> beam.Map(self.csv_mapper))
def __init__(
self,
table_name,
driver_class_name,
jdbc_url,
username,
password,
query=None,
output_parallelization=None,
fetch_size=None,
connection_properties=None,
connection_init_sqls=None,
expansion_service=None,
):
"""
Initializes a read operation from Jdbc.
:param driver_class_name: name of the jdbc driver class
:param jdbc_url: full jdbc url to the database.
:param username: database username
:param password: database password
:param query: sql query to be executed
:param output_parallelization: is output parallelization on
:param fetch_size: how many rows to fetch
:param connection_properties: properties of the jdbc connection
passed as string with format
[propertyName=property;]*
:param connection_init_sqls: required only for MySql and MariaDB.
passed as list of strings
:param expansion_service: The address (host:port) of the ExpansionService.
"""
super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
JdbcConfigSchema(
location=table_name,
config=RowCoder(
typing_to_runner_api(Config).row_type.schema).encode(
Config(
driver_class_name=driver_class_name,
jdbc_url=jdbc_url,
username=username,
password=password,
connection_properties=connection_properties,
connection_init_sqls=connection_init_sqls,
write_statement=None,
read_query=query,
fetch_size=fetch_size,
output_parallelization=output_parallelization,
))), ),
expansion_service or default_io_expansion_service(),
)
def expand(self, pbegin):
pcoll = pbegin.apply(
ExternalTransform(self.URN,
NamedTupleBasedPayloadBuilder(self.params),
self.expansion_service))
if self.params.with_attributes:
pcoll = pcoll | 'FromProto' >> Map(
pubsub.PubsubMessage._from_proto_str)
pcoll.element_type = pubsub.PubsubMessage
else:
pcoll.element_type = bytes
return pcoll
def expand(self, pvalue):
if self.with_attributes:
pcoll = pvalue | 'ToProto' >> Map(
pubsub.WriteToPubSub.to_proto_str)
else:
pcoll = pvalue | 'ToProto' >> Map(
lambda x: pubsub.PubsubMessage(x, {})._to_proto_str())
pcoll.element_type = bytes
return pcoll.apply(
ExternalTransform(self.URN,
NamedTupleBasedPayloadBuilder(self.params),
self.expansion_service))
def __init__(self, query, dialect=None, expansion_service=None):
"""
Creates a SqlTransform which will be expanded to Java's SqlTransform.
(See class docs).
:param query: The SQL query.
:param dialect: (optional) The dialect, e.g. use 'zetasql' for ZetaSQL.
:param expansion_service: (optional) The URL of the expansion service to use
"""
expansion_service = expansion_service or BeamJarExpansionService(
':sdks:java:extensions:sql:expansion-service:shadowJar')
super().__init__(self.URN,
NamedTupleBasedPayloadBuilder(
SqlTransformSchema(query=query, dialect=dialect)),
expansion_service=expansion_service)
def __init__(
self,
consumer_config,
topics,
key_deserializer=byte_array_deserializer,
value_deserializer=byte_array_deserializer,
start_read_time=None,
max_num_records=None,
max_read_time=None,
expansion_service=None,
):
"""
Initializes a read operation from Kafka.
:param consumer_config: A dictionary containing the consumer configuration.
:param topics: A list of topic strings.
:param key_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's key, e.g.
'org.apache.kafka.common.serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
:param value_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's value, e.g.
'org.apache.kafka.common.serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
:param start_read_time: Use timestamp to set up start offset in milliseconds
epoch.
:param max_num_records: Maximum amount of records to be read. Mainly used
for tests and demo applications.
:param max_read_time: Maximum amount of time in seconds the transform
executes. Mainly used for tests and demo applications.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
super(ReadFromKafka, self).__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
ReadFromKafkaSchema(
consumer_config=consumer_config,
topics=topics,
key_deserializer=key_deserializer,
value_deserializer=value_deserializer,
max_num_records=max_num_records,
max_read_time=max_read_time,
start_read_time=start_read_time,
)), expansion_service or default_io_expansion_service())
def __init__(
self,
project_id,
instance_id,
database_id,
table,
max_batch_size_bytes=None,
max_number_mutations=None,
max_number_rows=None,
grouping_factor=None,
host=None,
emulator_host=None,
commit_deadline=None,
max_cumulative_backoff=None,
expansion_service=None,
):
max_cumulative_backoff = int(
max_cumulative_backoff) if max_cumulative_backoff else None
commit_deadline = int(commit_deadline) if commit_deadline else None
super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
WriteToSpannerSchema(
project_id=project_id,
instance_id=instance_id,
database_id=database_id,
table=table,
max_batch_size_bytes=max_batch_size_bytes,
max_number_mutations=max_number_mutations,
max_number_rows=max_number_rows,
grouping_factor=grouping_factor,
host=host,
emulator_host=emulator_host,
commit_deadline=commit_deadline,
max_cumulative_backoff=max_cumulative_backoff,
), ),
expansion_service=expansion_service
or default_io_expansion_service(),
)
def expand(self, pbegin):
return (pbegin | ExternalTransform(
self.URN,
NamedTupleBasedPayloadBuilder(self.params),
self.expansion_service,
) | ParDo(_JsonStringToDictionaries()))
def __init__(
self,
consumer_config,
topics,
key_deserializer=byte_array_deserializer,
value_deserializer=byte_array_deserializer,
start_read_time=None,
max_num_records=None,
max_read_time=None,
commit_offset_in_finalize=False,
timestamp_policy=processing_time_policy,
with_metadata=False,
expansion_service=None,
):
"""
Initializes a read operation from Kafka.
:param consumer_config: A dictionary containing the consumer configuration.
:param topics: A list of topic strings.
:param key_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's key, e.g.
'org.apache.kafka.common.serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
:param value_deserializer: A fully-qualified Java class name of a Kafka
Deserializer for the topic's value, e.g.
'org.apache.kafka.common.serialization.LongDeserializer'.
Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
:param start_read_time: Use timestamp to set up start offset in milliseconds
epoch.
:param max_num_records: Maximum amount of records to be read. Mainly used
for tests and demo applications.
:param max_read_time: Maximum amount of time in seconds the transform
executes. Mainly used for tests and demo applications.
:param commit_offset_in_finalize: Whether to commit offsets when finalizing.
:param timestamp_policy: The built-in timestamp policy which is used for
extracting timestamp from KafkaRecord.
:param with_metadata: whether the returned PCollection should contain
Kafka related metadata or not. If False (default), elements of the
returned PCollection will be of type 'bytes' if True, elements of the
returned PCollection will be of the type 'Row'. Note that, currently
this only works when using default key and value deserializers where
Java Kafka Reader reads keys and values as 'byte[]'.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
if timestamp_policy not in [
ReadFromKafka.processing_time_policy,
ReadFromKafka.create_time_policy, ReadFromKafka.log_append_time
]:
raise ValueError('timestamp_policy should be one of '
'[ProcessingTime, CreateTime, LogAppendTime]')
super().__init__(
self.URN_WITH_METADATA
if with_metadata else self.URN_WITHOUT_METADATA,
NamedTupleBasedPayloadBuilder(
ReadFromKafkaSchema(
consumer_config=consumer_config,
topics=topics,
key_deserializer=key_deserializer,
value_deserializer=value_deserializer,
max_num_records=max_num_records,
max_read_time=max_read_time,
start_read_time=start_read_time,
commit_offset_in_finalize=commit_offset_in_finalize,
timestamp_policy=timestamp_policy)), expansion_service
or default_io_expansion_service())
def __init__(
self,
project_id,
instance_id,
database_id,
row_type=None,
sql=None,
table=None,
host=None,
emulator_host=None,
batching=None,
timestamp_bound_mode=None,
read_timestamp=None,
staleness=None,
time_unit=None,
expansion_service=None,
):
"""
Initializes a read operation from Spanner.
:param project_id: Specifies the Cloud Spanner project.
:param instance_id: Specifies the Cloud Spanner instance.
:param database_id: Specifies the Cloud Spanner database.
:param row_type: Row type that fits the given query or table. Passed as
NamedTuple, e.g. NamedTuple('name', [('row_name', unicode)])
:param sql: An sql query to execute. It's results must fit the
provided row_type. Don't use when table is set.
:param table: A spanner table. When provided all columns from row_type
will be selected to query. Don't use when query is set.
:param batching: By default Batch API is used to read data from Cloud
Spanner. It is useful to disable batching when the underlying query
is not root-partitionable.
:param host: Specifies the Cloud Spanner host.
:param emulator_host: Specifies Spanner emulator host.
:param timestamp_bound_mode: Defines how Cloud Spanner will choose a
timestamp for a read-only transaction or a single read/query.
Passed as TimestampBoundMode enum. Possible values:
STRONG: A timestamp bound that will perform reads and queries at a
timestamp where all previously committed transactions are visible.
READ_TIMESTAMP: Returns a timestamp bound that will perform reads
and queries at the given timestamp.
MIN_READ_TIMESTAMP: Returns a timestamp bound that will perform reads
and queries at a timestamp chosen to be at least given timestamp value.
EXACT_STALENESS: Returns a timestamp bound that will perform reads and
queries at an exact staleness. The timestamp is chosen soon after the
read is started.
MAX_STALENESS: Returns a timestamp bound that will perform reads and
queries at a timestamp chosen to be at most time_unit stale.
:param read_timestamp: Timestamp in string. Use only when
timestamp_bound_mode is set to READ_TIMESTAMP or MIN_READ_TIMESTAMP.
:param staleness: Staleness value as int. Use only when
timestamp_bound_mode is set to EXACT_STALENESS or MAX_STALENESS.
time_unit has to be set along with this param.
:param time_unit: Time unit for staleness_value passed as TimeUnit enum.
Possible values: NANOSECONDS, MICROSECONDS, MILLISECONDS, SECONDS,
HOURS, DAYS.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
assert row_type
assert sql or table and not (sql and table)
staleness_value = int(staleness) if staleness else None
if staleness_value or time_unit:
assert staleness_value and time_unit and \
timestamp_bound_mode is TimestampBoundMode.MAX_STALENESS or \
timestamp_bound_mode is TimestampBoundMode.EXACT_STALENESS
if read_timestamp:
assert timestamp_bound_mode is TimestampBoundMode.MIN_READ_TIMESTAMP\
or timestamp_bound_mode is TimestampBoundMode.READ_TIMESTAMP
super(ReadFromSpanner, self).__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
ReadFromSpannerSchema(
instance_id=instance_id,
database_id=database_id,
sql=sql,
table=table,
schema=named_tuple_to_schema(row_type).SerializeToString(),
project_id=project_id,
host=host,
emulator_host=emulator_host,
batching=batching,
timestamp_bound_mode=_get_enum_name(timestamp_bound_mode),
read_timestamp=read_timestamp,
staleness=staleness,
time_unit=_get_enum_name(time_unit),
), ),
expansion_service or default_io_expansion_service(),
)
