flink-python/pyflink/datastream/connectors/elasticsearch.py - flink - Git at Google

 ################################################################################
 #  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.
 ################################################################################
 import abc
 from enum import Enum
 from typing import List, Union

 from pyflink.datastream.connectors import Sink, DeliveryGuarantee
 from pyflink.java_gateway import get_gateway
 from pyflink.util.java_utils import to_jarray


 __all__ = ['FlushBackoffType',
            'ElasticsearchEmitter',
            'Elasticsearch6SinkBuilder',
            'Elasticsearch7SinkBuilder',
            'ElasticsearchSink']


 class FlushBackoffType(Enum):
     """
     Used to control whether the sink should retry failed requests at all or with which kind back off
     strategy.

     :data: `CONSTANT`:

     After every failure, it waits a configured time until the retries are exhausted.

     :data: `EXPONENTIAL`:

     After every failure, it waits initially the configured time and increases the waiting time
     exponentially until the retries are exhausted.

     :data: `NONE`:

     The failure is not retried.
     """

     CONSTANT = 0,
     EXPONENTIAL = 1,
     NONE = 2,

     def _to_j_flush_backoff_type(self):
         JFlushBackoffType = get_gateway().jvm \
             .org.apache.flink.connector.elasticsearch.sink.FlushBackoffType
         return getattr(JFlushBackoffType, self.name)


 class ElasticsearchEmitter(object):
     """
     Emitter which is used by sinks to prepare elements for sending them to Elasticsearch.
     """

     def __init__(self, j_emitter):
         self._j_emitter = j_emitter

     @staticmethod
     def static_index(index: str, key_field: str = None, doc_type: str = None) \
             -> 'ElasticsearchEmitter':
         """
         Creates an emitter with static index which is invoked on every record to convert it to
         Elasticsearch actions.
         """
         JMapElasticsearchEmitter = get_gateway().jvm \
             .org.apache.flink.connector.elasticsearch.sink.MapElasticsearchEmitter
         j_emitter = JMapElasticsearchEmitter(index, doc_type, key_field, False)
         return ElasticsearchEmitter(j_emitter)

     @staticmethod
     def dynamic_index(index_field: str, key_field: str = None, doc_type: str = None) \
             -> 'ElasticsearchEmitter':
         """
         Creates an emitter with dynamic index which is invoked on every record to convert it to
         Elasticsearch actions.
         """
         JMapElasticsearchEmitter = get_gateway().jvm \
             .org.apache.flink.connector.elasticsearch.sink.MapElasticsearchEmitter
         j_emitter = JMapElasticsearchEmitter(index_field, doc_type, key_field, True)
         return ElasticsearchEmitter(j_emitter)


 class ElasticsearchSinkBuilderBase(abc.ABC):
     """
     Base builder to construct a ElasticsearchSink.
     """

     @abc.abstractmethod
     def __init__(self):
         self._j_elasticsearch_sink_builder = None

     @abc.abstractmethod
     def get_http_host_class(self):
         """
         Gets the org.apache.http.HttpHost class which path is different in different Elasticsearch
         version.
         """
         pass

     def set_emitter(self, emitter: ElasticsearchEmitter) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the emitter which is invoked on every record to convert it to Elasticsearch actions.

         :param emitter: The emitter to process records into Elasticsearch actions.
         """
         self._j_elasticsearch_sink_builder.setEmitter(emitter._j_emitter)
         return self

     def set_hosts(self, hosts: Union[str, List[str]]) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the hosts where the Elasticsearch cluster nodes are reachable.
         """
         if not isinstance(hosts, list):
             hosts = [hosts]
         JHttpHost = self.get_http_host_class()
         j_http_hosts_list = [JHttpHost.create(x) for x in hosts]
         j_http_hosts_array = to_jarray(JHttpHost, j_http_hosts_list)
         self._j_elasticsearch_sink_builder.setHosts(j_http_hosts_array)
         return self

     def set_delivery_guarantee(self, delivery_guarantee: DeliveryGuarantee) \
             -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the wanted DeliveryGuarantee. The default delivery guarantee is DeliveryGuarantee#NONE
         """
         j_delivery_guarantee = delivery_guarantee._to_j_delivery_guarantee()
         self._j_elasticsearch_sink_builder.setDeliveryGuarantee(j_delivery_guarantee)
         return self

     def set_bulk_flush_max_actions(self, num_max_actions: int) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the maximum number of actions to buffer for each bulk request. You can pass -1 to
         disable it. The default flush size 1000.
         """
         self._j_elasticsearch_sink_builder.setBulkFlushMaxActions(num_max_actions)
         return self

     def set_bulk_flush_max_size_mb(self, max_size_mb: int) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the maximum size of buffered actions, in mb, per bulk request. You can pass -1 to
         disable it.
         """
         self._j_elasticsearch_sink_builder.setBulkFlushMaxSizeMb(max_size_mb)
         return self

     def set_bulk_flush_interval(self, interval_millis: int) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the bulk flush interval, in milliseconds. You can pass -1 to disable it.
         """
         self._j_elasticsearch_sink_builder.setBulkFlushInterval(interval_millis)
         return self

     def set_bulk_flush_backoff_strategy(self,
                                         flush_backoff_type: FlushBackoffType,
                                         max_retries: int,
                                         delay_millis: int) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the type of back off to use when flushing bulk requests. The default bulk flush back
         off type is FlushBackoffType#NONE.

         Sets the amount of delay between each backoff attempt when flushing bulk requests, in
         milliseconds.

         Sets the maximum number of retries for a backoff attempt when flushing bulk requests.
         """
         self._j_elasticsearch_sink_builder.setBulkFlushBackoffStrategy(
             flush_backoff_type._to_j_flush_backoff_type(), max_retries, delay_millis)
         return self

     def set_connection_username(self, username: str) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the username used to authenticate the connection with the Elasticsearch cluster.
         """
         self._j_elasticsearch_sink_builder.setConnectionUsername(username)
         return self

     def set_connection_password(self, password: str) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the password used to authenticate the connection with the Elasticsearch cluster.
         """
         self._j_elasticsearch_sink_builder.setConnectionPassword(password)
         return self

     def set_connection_path_prefix(self, prefix: str) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets a prefix which used for every REST communication to the Elasticsearch cluster.
         """
         self._j_elasticsearch_sink_builder.setConnectionPathPrefix(prefix)
         return self

     def set_connection_request_timeout(self, timeout: int) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the timeout for requesting the connection of the Elasticsearch cluster from the
         connection manager.
         """
         self._j_elasticsearch_sink_builder.setConnectionRequestTimeout(timeout)
         return self

     def set_connection_timeout(self, timeout: int) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the timeout for establishing a connection of the Elasticsearch cluster.
         """
         self._j_elasticsearch_sink_builder.setConnectionTimeout(timeout)
         return self

     def set_socket_timeout(self, timeout: int) -> 'ElasticsearchSinkBuilderBase':
         """
         Sets the timeout for waiting for data or, put differently, a maximum period inactivity
         between two consecutive data packets.
         """
         self._j_elasticsearch_sink_builder.setSocketTimeout(timeout)
         return self

     def build(self) -> 'ElasticsearchSink':
         """
         Constructs the ElasticsearchSink with the properties configured this builder.
         """
         return ElasticsearchSink(self._j_elasticsearch_sink_builder.build())


 class Elasticsearch6SinkBuilder(ElasticsearchSinkBuilderBase):
     """
     Builder to construct an Elasticsearch 6 compatible ElasticsearchSink.

     The following example shows the minimal setup to create a ElasticsearchSink that submits
     actions on checkpoint or the default number of actions was buffered (1000).

     Example:
     ::

         >>> sink = Elasticsearch6SinkBuilder() \\
         ...     .set_hosts('localhost:9200') \\
         ...     .set_emitter(ElasticsearchEmitter.static_index("user", "key_col")) \\
         ...     .build()
     """

     def __init__(self):
         self._j_elasticsearch_sink_builder = get_gateway().jvm \
             .org.apache.flink.connector.elasticsearch.sink.Elasticsearch6SinkBuilder()

     def get_http_host_class(self):
         return get_gateway().jvm.org.apache.flink.elasticsearch6.shaded.org.apache.http.HttpHost


 class Elasticsearch7SinkBuilder(ElasticsearchSinkBuilderBase):
     """
     Builder to construct an Elasticsearch 7 compatible ElasticsearchSink.

     The following example shows the minimal setup to create a ElasticsearchSink that submits
     actions on checkpoint or the default number of actions was buffered (1000).

     Example:
     ::

         >>> sink = Elasticsearch7SinkBuilder() \\
         ...     .set_hosts('localhost:9200') \\
         ...     .set_emitter(ElasticsearchEmitter.dynamic_index("index_col", "key_col")) \\
         ...     .build()
     """

     def __init__(self):
         self._j_elasticsearch_sink_builder = get_gateway().jvm \
             .org.apache.flink.connector.elasticsearch.sink.Elasticsearch7SinkBuilder()

     def get_http_host_class(self):
         return get_gateway().jvm.org.apache.flink.elasticsearch7.shaded.org.apache.http.HttpHost


 class ElasticsearchSink(Sink):
     """
     Flink Sink to insert or update data in an Elasticsearch index. The sink supports the following
     delivery guarantees.

     DeliveryGuarantee.NONE does not provide any guarantees: actions are flushed to Elasticsearch
     only depending on the configurations of the bulk processor. In case of a failure, it might
     happen that actions are lost if the bulk processor still has buffered actions.

     DeliveryGuarantee.AT_LEAST_ONCE on a checkpoint the sink will wait until all buffered actions
     are flushed to and acknowledged by Elasticsearch. No actions will be lost but actions might be
     sent to Elasticsearch multiple times when Flink restarts. These additional requests may cause
     inconsistent data in ElasticSearch right after the restart, but eventually everything will be
     consistent again.
     """
     def __init__(self, j_elasticsearch_sink):
         super(ElasticsearchSink, self).__init__(sink=j_elasticsearch_sink)
	################################################################################
	# 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.
	################################################################################
	import abc
	from enum import Enum
	from typing import List, Union

	from pyflink.datastream.connectors import Sink, DeliveryGuarantee
	from pyflink.java_gateway import get_gateway
	from pyflink.util.java_utils import to_jarray


	__all__ = ['FlushBackoffType',
	'ElasticsearchEmitter',
	'Elasticsearch6SinkBuilder',
	'Elasticsearch7SinkBuilder',
	'ElasticsearchSink']


	class FlushBackoffType(Enum):
	"""
	Used to control whether the sink should retry failed requests at all or with which kind back off
	strategy.

	:data: `CONSTANT`:

	After every failure, it waits a configured time until the retries are exhausted.

	:data: `EXPONENTIAL`:

	After every failure, it waits initially the configured time and increases the waiting time
	exponentially until the retries are exhausted.

	:data: `NONE`:

	The failure is not retried.
	"""

	CONSTANT = 0,
	EXPONENTIAL = 1,
	NONE = 2,

	def _to_j_flush_backoff_type(self):
	JFlushBackoffType = get_gateway().jvm \
	.org.apache.flink.connector.elasticsearch.sink.FlushBackoffType
	return getattr(JFlushBackoffType, self.name)


	class ElasticsearchEmitter(object):
	"""
	Emitter which is used by sinks to prepare elements for sending them to Elasticsearch.
	"""

	def __init__(self, j_emitter):
	self._j_emitter = j_emitter

	@staticmethod
	def static_index(index: str, key_field: str = None, doc_type: str = None) \
	-> 'ElasticsearchEmitter':
	"""
	Creates an emitter with static index which is invoked on every record to convert it to
	Elasticsearch actions.
	"""
	JMapElasticsearchEmitter = get_gateway().jvm \
	.org.apache.flink.connector.elasticsearch.sink.MapElasticsearchEmitter
	j_emitter = JMapElasticsearchEmitter(index, doc_type, key_field, False)
	return ElasticsearchEmitter(j_emitter)

	@staticmethod
	def dynamic_index(index_field: str, key_field: str = None, doc_type: str = None) \
	-> 'ElasticsearchEmitter':
	"""
	Creates an emitter with dynamic index which is invoked on every record to convert it to
	Elasticsearch actions.
	"""
	JMapElasticsearchEmitter = get_gateway().jvm \
	.org.apache.flink.connector.elasticsearch.sink.MapElasticsearchEmitter
	j_emitter = JMapElasticsearchEmitter(index_field, doc_type, key_field, True)
	return ElasticsearchEmitter(j_emitter)


	class ElasticsearchSinkBuilderBase(abc.ABC):
	"""
	Base builder to construct a ElasticsearchSink.
	"""

	@abc.abstractmethod
	def __init__(self):
	self._j_elasticsearch_sink_builder = None

	@abc.abstractmethod
	def get_http_host_class(self):
	"""
	Gets the org.apache.http.HttpHost class which path is different in different Elasticsearch
	version.
	"""
	pass

	def set_emitter(self, emitter: ElasticsearchEmitter) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the emitter which is invoked on every record to convert it to Elasticsearch actions.

	:param emitter: The emitter to process records into Elasticsearch actions.
	"""
	self._j_elasticsearch_sink_builder.setEmitter(emitter._j_emitter)
	return self

	def set_hosts(self, hosts: Union[str, List[str]]) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the hosts where the Elasticsearch cluster nodes are reachable.
	"""
	if not isinstance(hosts, list):
	hosts = [hosts]
	JHttpHost = self.get_http_host_class()
	j_http_hosts_list = [JHttpHost.create(x) for x in hosts]
	j_http_hosts_array = to_jarray(JHttpHost, j_http_hosts_list)
	self._j_elasticsearch_sink_builder.setHosts(j_http_hosts_array)
	return self

	def set_delivery_guarantee(self, delivery_guarantee: DeliveryGuarantee) \
	-> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the wanted DeliveryGuarantee. The default delivery guarantee is DeliveryGuarantee#NONE
	"""
	j_delivery_guarantee = delivery_guarantee._to_j_delivery_guarantee()
	self._j_elasticsearch_sink_builder.setDeliveryGuarantee(j_delivery_guarantee)
	return self

	def set_bulk_flush_max_actions(self, num_max_actions: int) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the maximum number of actions to buffer for each bulk request. You can pass -1 to
	disable it. The default flush size 1000.
	"""
	self._j_elasticsearch_sink_builder.setBulkFlushMaxActions(num_max_actions)
	return self

	def set_bulk_flush_max_size_mb(self, max_size_mb: int) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the maximum size of buffered actions, in mb, per bulk request. You can pass -1 to
	disable it.
	"""
	self._j_elasticsearch_sink_builder.setBulkFlushMaxSizeMb(max_size_mb)
	return self

	def set_bulk_flush_interval(self, interval_millis: int) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the bulk flush interval, in milliseconds. You can pass -1 to disable it.
	"""
	self._j_elasticsearch_sink_builder.setBulkFlushInterval(interval_millis)
	return self

	def set_bulk_flush_backoff_strategy(self,
	flush_backoff_type: FlushBackoffType,
	max_retries: int,
	delay_millis: int) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the type of back off to use when flushing bulk requests. The default bulk flush back
	off type is FlushBackoffType#NONE.

	Sets the amount of delay between each backoff attempt when flushing bulk requests, in
	milliseconds.

	Sets the maximum number of retries for a backoff attempt when flushing bulk requests.
	"""
	self._j_elasticsearch_sink_builder.setBulkFlushBackoffStrategy(
	flush_backoff_type._to_j_flush_backoff_type(), max_retries, delay_millis)
	return self

	def set_connection_username(self, username: str) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the username used to authenticate the connection with the Elasticsearch cluster.
	"""
	self._j_elasticsearch_sink_builder.setConnectionUsername(username)
	return self

	def set_connection_password(self, password: str) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the password used to authenticate the connection with the Elasticsearch cluster.
	"""
	self._j_elasticsearch_sink_builder.setConnectionPassword(password)
	return self

	def set_connection_path_prefix(self, prefix: str) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets a prefix which used for every REST communication to the Elasticsearch cluster.
	"""
	self._j_elasticsearch_sink_builder.setConnectionPathPrefix(prefix)
	return self

	def set_connection_request_timeout(self, timeout: int) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the timeout for requesting the connection of the Elasticsearch cluster from the
	connection manager.
	"""
	self._j_elasticsearch_sink_builder.setConnectionRequestTimeout(timeout)
	return self

	def set_connection_timeout(self, timeout: int) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the timeout for establishing a connection of the Elasticsearch cluster.
	"""
	self._j_elasticsearch_sink_builder.setConnectionTimeout(timeout)
	return self

	def set_socket_timeout(self, timeout: int) -> 'ElasticsearchSinkBuilderBase':
	"""
	Sets the timeout for waiting for data or, put differently, a maximum period inactivity
	between two consecutive data packets.
	"""
	self._j_elasticsearch_sink_builder.setSocketTimeout(timeout)
	return self

	def build(self) -> 'ElasticsearchSink':
	"""
	Constructs the ElasticsearchSink with the properties configured this builder.
	"""
	return ElasticsearchSink(self._j_elasticsearch_sink_builder.build())


	class Elasticsearch6SinkBuilder(ElasticsearchSinkBuilderBase):
	"""
	Builder to construct an Elasticsearch 6 compatible ElasticsearchSink.

	The following example shows the minimal setup to create a ElasticsearchSink that submits
	actions on checkpoint or the default number of actions was buffered (1000).

	Example:
	::

	>>> sink = Elasticsearch6SinkBuilder() \\
	... .set_hosts('localhost:9200') \\
	... .set_emitter(ElasticsearchEmitter.static_index("user", "key_col")) \\
	... .build()
	"""

	def __init__(self):
	self._j_elasticsearch_sink_builder = get_gateway().jvm \
	.org.apache.flink.connector.elasticsearch.sink.Elasticsearch6SinkBuilder()

	def get_http_host_class(self):
	return get_gateway().jvm.org.apache.flink.elasticsearch6.shaded.org.apache.http.HttpHost


	class Elasticsearch7SinkBuilder(ElasticsearchSinkBuilderBase):
	"""
	Builder to construct an Elasticsearch 7 compatible ElasticsearchSink.

	The following example shows the minimal setup to create a ElasticsearchSink that submits
	actions on checkpoint or the default number of actions was buffered (1000).

	Example:
	::

	>>> sink = Elasticsearch7SinkBuilder() \\
	... .set_hosts('localhost:9200') \\
	... .set_emitter(ElasticsearchEmitter.dynamic_index("index_col", "key_col")) \\
	... .build()
	"""

	def __init__(self):
	self._j_elasticsearch_sink_builder = get_gateway().jvm \
	.org.apache.flink.connector.elasticsearch.sink.Elasticsearch7SinkBuilder()

	def get_http_host_class(self):
	return get_gateway().jvm.org.apache.flink.elasticsearch7.shaded.org.apache.http.HttpHost


	class ElasticsearchSink(Sink):
	"""
	Flink Sink to insert or update data in an Elasticsearch index. The sink supports the following
	delivery guarantees.

	DeliveryGuarantee.NONE does not provide any guarantees: actions are flushed to Elasticsearch
	only depending on the configurations of the bulk processor. In case of a failure, it might
	happen that actions are lost if the bulk processor still has buffered actions.

	DeliveryGuarantee.AT_LEAST_ONCE on a checkpoint the sink will wait until all buffered actions
	are flushed to and acknowledged by Elasticsearch. No actions will be lost but actions might be
	sent to Elasticsearch multiple times when Flink restarts. These additional requests may cause
	inconsistent data in ElasticSearch right after the restart, but eventually everything will be
	consistent again.
	"""
	def __init__(self, j_elasticsearch_sink):
	super(ElasticsearchSink, self).__init__(sink=j_elasticsearch_sink)