def poll(self, timeout: float = -1) -> (str, dict):
""" Consume one message from Fink server
Parameters
----------
timeout: float, optional
maximum time to block waiting for a message
if not set default is None i.e. wait indefinitely
Returns
----------
(topic, alert, key): tuple(str, dict, str)
returns (None, None, None) on timeout
"""
msg = self._consumer.poll(timeout)
if msg is None:
return None, None, None
# msg.error() returns None or KafkaError
if msg.error():
error_message = """
Error: {} topic: {}[{}] at offset: {} with key: {}
""".format(msg.error(), msg.topic(), msg.partition(), msg.offset(),
str(msg.key()))
raise AlertError(error_message)
topic = msg.topic()
# decode the key if it is bytes
key = msg.key()
if type(key) == bytes:
key = key.decode('utf8')
if key is None:
# backward compatibility
key = '1.0_0.4.3'
# Get the schema
if self.schema_path is not None:
_parsed_schema = _get_alert_schema(schema_path=self.schema_path)
elif key is not None:
_parsed_schema = _get_alert_schema(key=key)
else:
msg = """
The message cannot be decoded as there is no key (None). Alternatively
specify manually the schema path when instantiating ``AlertConsumer`` (or from fink_consumer).
"""
raise NotImplementedError(msg)
self._parsed_schema = _parsed_schema
avro_alert = io.BytesIO(msg.value())
alert = _decode_avro_alert(avro_alert, _parsed_schema)
return topic, alert, key
def consume(self, num_alerts: int = 1, timeout: float = -1) -> list:
""" Consume and return list of messages
Parameters
----------
num_messages: int
maximum number of messages to return
timeout: float
maximum time to block waiting for messages
if not set default is None i.e. wait indefinitely
Returns
----------
list: [tuple(str, dict, str)]
list of topic, alert, key
returns an empty list on timeout
"""
alerts = []
msg_list = self._consumer.consume(num_alerts, timeout)
for msg in msg_list:
topic = msg.topic()
# decode the key if it is bytes
key = msg.key()
if type(key) == bytes:
key = key.decode('utf8')
if key is None:
# backward compatibility
key = '1.0_0.4.3'
# Get the schema
if self.schema_path is not None:
_parsed_schema = _get_alert_schema(
schema_path=self.schema_path)
elif key is not None:
_parsed_schema = _get_alert_schema(key=key)
else:
msg = """
The message cannot be decoded as there is no key (None). Either specify a
key when writing the alert, or specify manually the schema path when
instantiating ``AlertConsumer`` (or from fink_consumer).
"""
raise NotImplementedError(msg)
avro_alert = io.BytesIO(msg.value())
alert = _decode_avro_alert(avro_alert, _parsed_schema)
alerts.append((topic, alert, key))
return alerts
def poll_and_write(self,
outdir: str,
timeout: float = -1,
overwrite: bool = False) -> (str, dict):
""" Consume one message from Fink server, save alert on disk and
return (topic, alert, key)
Parameters
----------
outdir: str
Folder to store the alert. It must exists.
timeout: float, optional
maximum time to block waiting for messages
if not set default is None i.e. wait indefinitely
overwrite: bool, optional
If True, allow an existing alert to be overwritten.
Default is False.
Returns
----------
(topic, alert): tuple(str, dict)
returns (None, None) on timeout
"""
topic, alert, key = self.poll(timeout)
if topic is not None:
# Get the schema
if self.schema_path is not None:
_parsed_schema = _get_alert_schema(
schema_path=self.schema_path)
elif key is not None:
_parsed_schema = _get_alert_schema(key=key)
else:
msg = """
The message cannot be written as there is no key (None). Either specify a
key when writing the alert, or specify manually the schema path when
instantiating ``AlertConsumer`` (or from fink_consumer).
"""
raise NotImplementedError(msg)
# print('Alert written at {}'.format(outdir))
write_alert(alert, _parsed_schema, outdir, overwrite=overwrite)
return topic, alert, key
def __init__(self, topics: list, config: dict, schema=None):
"""Creates an instance of `AlertConsumer`
Parameters
----------
topics : list of str
list of topics to subscribe
config: dict
Dictionary of configurations. Allowed keys are:
username: str
username for API access
password: str
password for API access
group_id: str
group.id for Kafka consumer
bootstrap.servers: str, optional
Kafka servers to connect to
"""
self._topics = topics
self._kafka_config = _get_kafka_config(config)
self._parsed_schema = _get_alert_schema(schema_path=schema)
self._consumer = confluent_kafka.Consumer(self._kafka_config)
self._consumer.subscribe(self._topics)