class VideoPipeline:
def __init__(self, func: Callable[[bytes], bytes]):
self.nc = None
self.func = func
self.msg_idx = 0
self.nats_endpoint = _env_var('NATS_ENDPOINT', 'NATS broker endpoint')
self.src_topic = _env_var('NATS_SRC_TOPIC', 'source NATS topic')
self.dst_topic = _env_var('NATS_DST_TOPIC', 'destination NATS topic')
logging.info("[{}]: {} -> {}".format(self.nats_endpoint,
self.src_topic, self.dst_topic))
async def run(self, loop):
self.nc = NATS()
await self.nc.connect(loop=loop, servers=[str(self.nats_endpoint)])
logging.info('connected to NATS')
await self.nc.subscribe(str(self.src_topic), cb=self._message_handler)
async def _message_handler(self, msg):
ds_msg = xi_iot_pb2.DataStreamMessage()
ds_msg.ParseFromString(msg.data)
try:
ds_msg.payload = self.func(ds_msg.payload) # TODO error handling
except Exception as e:
logging.error("failed to process message #{}: {}".format(
self.msg_idx, e))
if self.msg_idx % 100 == 0:
logging.info("processed message #{}".format(self.msg_idx))
try:
await self.nc.publish(str(self.dst_topic),
ds_msg.SerializeToString())
except Exception as e:
logging.error("failed to publish message #{}: {}".format(
self.msg_idx, e))
self.msg_idx += 1
def __del__(self):
if self.nc:
self.nc.drain()
class Matcher:
fut_stop: asyncio.Future
def __init__(self):
self.nc = NATS()
client_stub = pydgraph.DgraphClientStub("dgraph:9080")
self.dg = pydgraph.DgraphClient(client_stub)
loop = asyncio.get_running_loop()
loop.add_signal_handler(signal.SIGTERM, self.stop)
loop.add_signal_handler(signal.SIGHUP, self.stop)
loop.add_signal_handler(signal.SIGINT, self.stop)
self.fut_stop = loop.create_future()
async def setup(self):
await self.nc.connect("nats://nats")
await self.nc.subscribe(TOPIC, "matcher", self.on_update)
print("Listening on ", TOPIC)
self.setup_schema()
def setup_schema(self):
schema = """
email_address: string @index(exact) .
id: string @index(exact) .
"""
op = pydgraph.Operation(schema=schema)
self.dg.alter(op)
async def reply(self, msg, resp):
if msg.reply:
await self.nc.publish(msg.reply, resp)
async def on_update(self, msg) -> None:
tr = None
try:
tr = ExtractionTrigger.from_bytes(msg.data)
except Exception:
traceback.print_exc()
await self.reply(msg, orjson.dumps({"ok": False}))
return
txn = self.dg.txn()
nq = None
query = None
try:
nq = tr.nquads()
query = tr.query()
mutation = txn.create_mutation(set_nquads=nq)
request = txn.create_request(query=query,
mutations=[mutation],
commit_now=True)
txn.commit()
await self.reply(msg, orjson.dumps({"ok": True}))
except Exception:
traceback.print_exc()
print(nq, query)
txn.discard()
await self.reply(msg, orjson.dumps({"ok": False}))
def stop(self) -> None:
self.fut_stop.set_result(True)
async def run_until_done(self) -> None:
await self.fut_stop
await asyncio.gather(self.nc.drain())
class NATSHelper(object):
def __init__(self,
nats_broker_url=None,
nats_src_topic=None,
nats_dst_topic=None):
self.subscribe_id = None
self.connected = False
self.nats_client = NATS()
self._get_config_from_env_var_('nats_broker_url', 'NATS_ENDPOINT',
nats_broker_url, 'NATS broker')
self._get_config_from_env_var_('nats_src_topic', 'NATS_SRC_TOPIC',
nats_src_topic, 'NATS source topic')
self._get_config_from_env_var_('nats_dst_topic', 'NATS_DST_TOPIC',
nats_dst_topic,
'NATS destination topic')
logger.info("broker: {b}, src topic: {s}, dst_topic: {d}".format(
b=self.nats_broker_url,
s=self.nats_src_topic,
d=self.nats_dst_topic))
def __del__(self):
self.close()
def _get_config_from_env_var_(self, attr, var, default, what):
if default is None:
value = os.environ.get(var)
else:
value = default
if not value: # also check empty strings
raise Exception(
'{what} not provided in environment var {var}'.format(
what=what, var=var))
setattr(self, attr, value)
@staticmethod
def payload_from_message(msg):
"""
Convert input payload into an image
"""
_msg = xi_iot_pb2.DataStreamMessage()
_msg.ParseFromString(msg.data)
return _msg.payload
@staticmethod
def message_from_payload(payload):
"""
Convert image into output payload
"""
msg = xi_iot_pb2.DataStreamMessage()
msg.payload = payload
return msg.SerializeToString()
async def publish(self, payload):
try:
payload = self.message_from_payload(payload)
# RFC: We could leverage `reply` topic as the destination topic which would not require NATS_DST_TOPIC to be provided
# await nc.publish(reply, data)
logger.info("Sending message to topic '{}'".format(
self.nats_dst_topic))
await self.nats_client.publish(self.nats_dst_topic, payload)
except Exception as e:
# Catch an display errors which are otherwise not shown
logger.error("{}".format(e))
raise
async def connect(self, loop, message_handler_cb):
# Define a helper function to unbox the message and catch errors for display
async def receive_cb(msg):
try:
logger.info("Received a message!")
payload = self.payload_from_message(msg)
await message_handler_cb(payload)
except Exception as e:
# Catch an display errors which are otherwise not shown
logger.error("{}".format(e))
raise
try:
# This will return immediately if the server is not listening on the given URL
await self.nats_client.connect(self.nats_broker_url, loop=loop)
self.connected = True
logger.info(
"Connected to broker, subscribing to topic '{}'".format(
self.nats_src_topic))
self.subscribe_id = await self.nats_client.subscribe(
self.nats_src_topic, cb=receive_cb)
except Exception as e:
# Catch an display errors which are otherwise not shown
logger.error("{}".format(e))
raise
def close(self):
# Remove interest in subscription.
loop = asyncio.get_event_loop()
if self.subscribe_id is not None:
loop.run_until_complete(
self.nats_client.unsubscribe(self.subscribe_id))
self.subscribe_id = None
# Terminate connection to NATS.
if self.nats_client is not None and self.connected:
self.nats_client.drain()
loop.run_until_complete(self.nats_client.close())
self.nats_client = None
self.connected = False
