sink.py

import logging

from ..abc.sink import Sink
from .connection import ElasticSearchBulk

from .data_feeder import data_feeder_create_or_index

import asab

#

L = logging.getLogger(__name__)

#


class ElasticSearchSink(Sink):
	"""
	ElasticSearchSink allows you to insert events into ElasticSearch through POST requests

	The following attributes can be passed to the context and thus override the default behavior
	of the sink:

	es_index (STRING): ElasticSearch index name

	data_feeder accepts the event as its only parameter and yields data as Python generator
	The example implementation is:

	def data_feeder_create_or_index(event):
			_id = event.pop("_id", None)

			if _id is None:
					yield b'{"create":{}}\n'
			else:
					yield orjson.dumps(
							{"index": {"_id": _id}}, option=orjson.OPT_APPEND_NEWLINE
					)

			yield orjson.dumps(event, option=orjson.OPT_APPEND_NEWLINE)

	|

	"""

	ConfigDefaults = {
		"index_prefix": "bspump_",  # Obsolete, use 'index'
		"index": "bspump_",
	}

	def __init__(self, app, pipeline, connection, id=None, config=None, bulk_class=ElasticSearchBulk, data_feeder=data_feeder_create_or_index):
		"""
		Description:

		**Parameters**

		app : Application
				Name of the Application

		pipeline : Pipeline
				Name of the Pipeline

		connection : Connection
				Name of the Connection

		id : ID, default= None
				ID

		config : JSON, default= None
				Configuration file with additional information.

		bulk_class=ElasticBulk :

		data_feeder=data_feeder_create_or_index :

		|

		"""
		super().__init__(app, pipeline, id=id, config=config)

		self.Connection = pipeline.locate_connection(app, connection)
		self.BulkClass = bulk_class

		self.Index = self.Config.get('index')

		# intex_prefix is obsolete. It is supported currently ensure backward compatibility
		if self.Index == "bspump_" and self.Config.get('index_prefix') != "bspump_" and len(self.Config.get('index_prefix')) > 0:
			asab.LogObsolete.warning("The 'index_prefix' has been renamed to 'index', please adjust the configuration.")
			self.Index = self.Config.get('index_prefix')

		if data_feeder is None:
			raise RuntimeError("data_feeder must not be None.")

		self.__data_feeder = data_feeder

		app.PubSub.subscribe("ElasticSearchConnection.pause!", self._connection_throttle)
		app.PubSub.subscribe("ElasticSearchConnection.unpause!", self._connection_throttle)


	def process(self, context, event):


		"""
		Description:

		**Parameters**

		context :

		event : any data type
				Information with timestamp.


		"""
		try:
			_id = event.pop("_id", None)
		except TypeError:
			if isinstance(event, dict) is False:
				L.error("You are trying to pass event of type: {} to ElasticSearchSink, but only dict is supported".format(type(event)))
			raise
		self.Connection.consume(
			context.get("es_index", self.Index),
			self.__data_feeder(event, _id),
			bulk_class=self.BulkClass,
		)

	def _connection_throttle(self, event_name, connection):
		if connection != self.Connection:
			return

		if event_name == "ElasticSearchConnection.pause!":
			self.Pipeline.throttle(self, True)
		elif event_name == "ElasticSearchConnection.unpause!":
			self.Pipeline.throttle(self, False)
		else:
			raise RuntimeError("Unexpected event name '{}'".format(event_name))