| ################################################################################ |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| ################################################################################ |
| from pyflink.common import SerializationSchema, DeserializationSchema |
| from pyflink.datastream.functions import SinkFunction, SourceFunction |
| from pyflink.java_gateway import get_gateway |
| |
| |
| __all__ = [ |
| 'RMQConnectionConfig', |
| 'RMQSource', |
| 'RMQSink' |
| ] |
| |
| |
| class RMQConnectionConfig(object): |
| """ |
| Connection Configuration for RMQ. |
| """ |
| |
| def __init__(self, j_rmq_connection_config): |
| self._j_rmq_connection_config = j_rmq_connection_config |
| |
| def get_host(self) -> str: |
| return self._j_rmq_connection_config.getHost() |
| |
| def get_port(self) -> int: |
| return self._j_rmq_connection_config.getPort() |
| |
| def get_virtual_host(self) -> str: |
| return self._j_rmq_connection_config.getVirtualHost() |
| |
| def get_user_name(self) -> str: |
| return self._j_rmq_connection_config.getUsername() |
| |
| def get_password(self) -> str: |
| return self._j_rmq_connection_config.getPassword() |
| |
| def get_uri(self) -> str: |
| return self._j_rmq_connection_config.getUri() |
| |
| def get_network_recovery_interval(self) -> int: |
| return self._j_rmq_connection_config.getNetworkRecoveryInterval() |
| |
| def is_automatic_recovery(self) -> bool: |
| return self._j_rmq_connection_config.isAutomaticRecovery() |
| |
| def is_topology_recovery(self) -> bool: |
| return self._j_rmq_connection_config.isTopologyRecovery() |
| |
| def get_connection_timeout(self) -> int: |
| return self._j_rmq_connection_config.getConnectionTimeout() |
| |
| def get_requested_channel_max(self) -> int: |
| return self._j_rmq_connection_config.getRequestedChannelMax() |
| |
| def get_requested_frame_max(self) -> int: |
| return self._j_rmq_connection_config.getRequestedFrameMax() |
| |
| def get_requested_heartbeat(self) -> int: |
| return self._j_rmq_connection_config.getRequestedHeartbeat() |
| |
| class Builder(object): |
| """ |
| Builder for RMQConnectionConfig. |
| """ |
| |
| def __init__(self): |
| self._j_options_builder = get_gateway().jvm.org.apache.flink.streaming.connectors\ |
| .rabbitmq.common.RMQConnectionConfig.Builder() |
| |
| def set_port(self, port: int) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setPort(port) |
| return self |
| |
| def set_host(self, host: str) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setHost(host) |
| return self |
| |
| def set_virtual_host(self, vhost: str) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setVirtualHost(vhost) |
| return self |
| |
| def set_user_name(self, user_name: str) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setUserName(user_name) |
| return self |
| |
| def set_password(self, password: str) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setPassword(password) |
| return self |
| |
| def set_uri(self, uri: str) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setUri(uri) |
| return self |
| |
| def set_topology_recovery_enabled( |
| self, topology_recovery_enabled: bool) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setTopologyRecoveryEnabled(topology_recovery_enabled) |
| return self |
| |
| def set_requested_heartbeat( |
| self, requested_heartbeat: int) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setRequestedHeartbeat(requested_heartbeat) |
| return self |
| |
| def set_requested_frame_max( |
| self, requested_frame_max: int) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setRequestedFrameMax(requested_frame_max) |
| return self |
| |
| def set_requested_channel_max( |
| self, requested_channel_max: int) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setRequestedChannelMax(requested_channel_max) |
| return self |
| |
| def set_network_recovery_interval( |
| self, network_recovery_interval: int) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setNetworkRecoveryInterval(network_recovery_interval) |
| return self |
| |
| def set_connection_timeout(self, connection_timeout: int) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setConnectionTimeout(connection_timeout) |
| return self |
| |
| def set_automatic_recovery(self, automatic_recovery: bool) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setAutomaticRecovery(automatic_recovery) |
| return self |
| |
| def set_prefetch_count(self, prefetch_count: int) -> 'RMQConnectionConfig.Builder': |
| self._j_options_builder.setPrefetchCount(prefetch_count) |
| return self |
| |
| def build(self) -> 'RMQConnectionConfig': |
| return RMQConnectionConfig(self._j_options_builder.build()) |
| |
| |
| class RMQSource(SourceFunction): |
| def __init__(self, |
| connection_config: 'RMQConnectionConfig', |
| queue_name: str, |
| use_correlation_id: bool, |
| deserialization_schema: DeserializationSchema |
| ): |
| """ |
| Creates a new RabbitMQ source. |
| |
| For exactly-once, you must set the correlation ids of messages at the producer. |
| The correlation id must be unique. Otherwise the behavior of the source is undefined. |
| |
| If in doubt, set use_correlation_id to False. |
| |
| When correlation ids are not used, this source has at-least-once processing semantics |
| when checkpointing is enabled. |
| |
| :param connection_config: The RabbiMQ connection configuration. |
| :param queue_name: The queue to receive messages from. |
| :param use_correlation_id: Whether the messages received are supplied with a unique id |
| to deduplicate messages (in case of failed acknowledgments). |
| Only used when checkpointing is enabled. |
| :param deserialization_schema: A deserializer used to convert between RabbitMQ's |
| messages and Flink's objects. |
| """ |
| JRMQSource = get_gateway().jvm.org.apache.flink.streaming.connectors.rabbitmq.RMQSource |
| j_rmq_source = JRMQSource( |
| connection_config._j_rmq_connection_config, |
| queue_name, |
| use_correlation_id, |
| deserialization_schema._j_deserialization_schema |
| ) |
| super(RMQSource, self).__init__(source_func=j_rmq_source) |
| |
| |
| class RMQSink(SinkFunction): |
| def __init__(self, connection_config: 'RMQConnectionConfig', |
| queue_name: str, serialization_schema: SerializationSchema): |
| """ |
| Creates a new RabbitMQ sink. |
| |
| :param connection_config: The RabbiMQ connection configuration. |
| :param queue_name: The queue to publish messages to. |
| :param serialization_schema: A serializer used to convert Flink objects to bytes. |
| """ |
| JRMQSink = get_gateway().jvm.org.apache.flink.streaming.connectors.rabbitmq.RMQSink |
| j_rmq_sink = JRMQSink( |
| connection_config._j_rmq_connection_config, |
| queue_name, |
| serialization_schema._j_serialization_schema, |
| ) |
| super(RMQSink, self).__init__(sink_func=j_rmq_sink) |
