From 1d150cdc21e2e7bb6aa1c3cd7a418b6faae1a6b3 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 14 Apr 2025 17:51:23 +0100 Subject: [PATCH] Working on processor API --- trustgraph-base/trustgraph/base/__init__.py | 6 +- .../trustgraph/base/async_processor.py | 200 ++++++++++ .../trustgraph/base/base_processor.py | 375 ------------------ trustgraph-base/trustgraph/base/consumer.py | 151 ++----- .../trustgraph/base/consumer.py.OLD | 187 +++++++++ trustgraph-base/trustgraph/base/producer.py | 62 +-- .../trustgraph/base/producer.py.OLD | 57 +++ trustgraph-base/trustgraph/base/pubsub.py | 96 +++++ .../trustgraph/config/service/service.py | 1 - 9 files changed, 588 insertions(+), 547 deletions(-) create mode 100644 trustgraph-base/trustgraph/base/async_processor.py delete mode 100644 trustgraph-base/trustgraph/base/base_processor.py create mode 100644 trustgraph-base/trustgraph/base/consumer.py.OLD create mode 100644 trustgraph-base/trustgraph/base/producer.py.OLD create mode 100644 trustgraph-base/trustgraph/base/pubsub.py diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 3a58d51e..074bc286 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -1,8 +1,6 @@ -from . base_processor import BaseProcessor +from . pubsub import PulsarClient +from . async_processor import AsyncProcessor from . consumer import Consumer from . producer import Producer -from . consumer_producer import ConsumerProducer -from . publisher import Publisher -from . subscriber import Subscriber diff --git a/trustgraph-base/trustgraph/base/async_processor.py b/trustgraph-base/trustgraph/base/async_processor.py new file mode 100644 index 00000000..b355d3a2 --- /dev/null +++ b/trustgraph-base/trustgraph/base/async_processor.py @@ -0,0 +1,200 @@ + +import asyncio +import os +import argparse +import pulsar +from pulsar.schema import JsonSchema +import _pulsar +import time +import uuid +from prometheus_client import start_http_server, Info, Enum + +from .. schema import ConfigPush, config_push_queue +from .. log_level import LogLevel +from .. exceptions import TooManyRequests +from . pubsub import PulsarClient +from . producer import Producer +from . consumer import Consumer + +default_config_queue = config_push_queue +config_subscriber_id = str(uuid.uuid4()) + +class AsyncProcessor: + + + def __init__(self, **params): + + self.client = PulsarClient(**params) + + self.taskgroup = params.get("taskgroup") + if self.taskgroup is None: + raise RuntimeError("Essential taskgroup missing") + + if not hasattr(__class__, "params_metric"): + __class__.params_metric = Info( + 'params', 'Parameters configuration' + ) + + # FIXME: Maybe outputs information it should not + __class__.params_metric.info({ + k: str(params[k]) + for k in params + }) + + self.config_push_queue = params.get( + "config_push_queue", default_config_queue + ) + + self.config_handlers = [] + + @property + def pulsar_host(self): return self.client.pulsar_host + + async def start(self): + + self.config_sub_task = self.subscribe( + self.config_push_queue, config_subscriber_id, + schema=ConfigPush, handler=self.on_config_change + ) + + async def on_config_change(self, config, version): + for ch in self.config_handlers: + ch(config, version) + + async def run_config_queue(self): + + if self.module == "config.service": + print("I am config-svc, not looking at config queue", flush=True) + return + + print("Config thread running", flush=True) + + async def run(self): + while True: + await asyncio.sleep(2) + + def subscribe(self, queue, subscriber, schema, handler): + + return Consumer( + self.taskgroup, self.client, queue, subscriber, schema, handler + ) + + def publish(self, queue, schema): + + return Producer( + self.client, queue, schema + ) + + def set_processor_state(self, flow, state): + + if not hasattr(__class__, "state_metric"): + __class__.state_metric = Enum( + 'processor_state', 'Processor state', + ["flow"], + states=['starting', 'running', 'stopped'] + ) + + __class__.state_metric.labels("flow").state(state) + + def set_pubsub_info(self, flow, info) : + + if not hasattr(__class__, "pubsub_metric"): + __class__.pubsub_metric = Info( + 'pubsub', 'Pub/sub configuration', + ["flow"] + ) + + __class__.pubsub_metric.labels(flow=flow).info(info) + + @classmethod + async def launch_async(cls, args, ident): + + async with asyncio.TaskGroup() as tg: + + p = cls(**args | { "taskgroup": tg }) + + # FIXME: Two sort of 'ident' things going on here? + p.module = ident + p.config_ident = args.get("ident", "FIXME") + + await p.start() + +# task1 = tg.create_task(p.run_config_queue()) + task2 = tg.create_task(p.run()) + + @classmethod + def launch(cls, ident, doc): + + parser = argparse.ArgumentParser( + prog=ident, + description=doc + ) + + parser.add_argument( + '--id', + default=ident, + help=f'Configuration identity (default: {ident})', + ) + + cls.add_args(parser) + + args = parser.parse_args() + args = vars(args) + + print(args) + + if args["metrics"]: + start_http_server(args["metrics_port"]) + + while True: + + try: + + asyncio.run(cls.launch_async( + args, ident + )) + + except KeyboardInterrupt: + print("Keyboard interrupt.") + return + + except _pulsar.Interrupted: + print("Pulsar Interrupted.") + return + + except Exception as e: + + print(type(e)) + +# print(e.message) +# print(e.exceptions) + + print("Exception:", e, flush=True) + print("Will retry...", flush=True) + + time.sleep(4) + + @staticmethod + def add_args(parser): + + PulsarClient.add_args(parser) + + parser.add_argument( + '--config-push-queue', + default=default_config_queue, + help=f'Config push queue {default_config_queue}', + ) + + parser.add_argument( + '--metrics', + action=argparse.BooleanOptionalAction, + default=True, + help=f'Metrics enabled (default: true)', + ) + + parser.add_argument( + '-P', '--metrics-port', + type=int, + default=8000, + help=f'Pulsar host (default: 8000)', + ) diff --git a/trustgraph-base/trustgraph/base/base_processor.py b/trustgraph-base/trustgraph/base/base_processor.py deleted file mode 100644 index d5627fba..00000000 --- a/trustgraph-base/trustgraph/base/base_processor.py +++ /dev/null @@ -1,375 +0,0 @@ - -import asyncio -import os -import argparse -import pulsar -from pulsar.schema import JsonSchema -import _pulsar -import time -import uuid -from prometheus_client import start_http_server, Info, Enum - -from .. schema import ConfigPush, config_push_queue -from .. log_level import LogLevel -from .. exceptions import TooManyRequests - -default_config_queue = config_push_queue -config_subscriber_id = str(uuid.uuid4()) - -class Subscription: - def __init__(self, consumer, handler, taskgroup): - self.running = True - self.task = None - self.consumer = consumer - self.handle = handler - self.taskgroup=taskgroup - - async def start(self): - self.running = True - self.task = self.taskgroup.create_task(self.run()) - - async def run(self): - - while self.running: - - print("Waiting...") - msg = await asyncio.to_thread(self.consumer.receive) - print("Got", msg) - -# expiry = time.time() + self.rate_limit_timeout - expiry = time.time() + 10 - - # This loop is for retry on rate-limit / resource limits - while True: - - if time.time() > expiry: - - print("Gave up waiting for rate-limit retry", flush=True) - - # Message failed to be processed, this causes it to - # be retried - self.consumer.negative_acknowledge(msg) - - # FIXME -# __class__.processing_metric.labels(status="error").inc() - - # Break out of retry loop, processes next message - break - - try: - - print("Handle...") - # FIXME -# with __class__.request_metric.time(): - await self.handle(msg, self.consumer) - print("Handled.") - - # Acknowledge successful processing of the message - self.consumer.acknowledge(msg) - -# __class__.processing_metric.labels(status="success").inc() - - # Break out of retry loop - break - - except TooManyRequests: - - print("TooManyRequests: will retry...", flush=True) - - # FIXME -# __class__.rate_limit_metric.inc() - - # Sleep - time.sleep(self.rate_limit_retry) - - # Contine from retry loop, just causes a reprocessing - continue - - except Exception as e: - - print("Exception:", e, flush=True) - - # Message failed to be processed, this causes it to - # be retried - self.consumer.negative_acknowledge(msg) - -# __class__.processing_metric.labels(status="error").inc() - - # Break out of retry loop, processes next message - break - -class Publisher: - - def __init__(self, producer): - self.producer = producer -# self.running = True - - async def send(self, msg, properties={}): - self.producer.send(msg, properties) - - # FIXME -# __class__.output_metric.inc() - - -class BaseProcessor: - - default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650') - default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) - - def __init__(self, **params): - - self.taskgroup = params.get("taskgroup") - if self.taskgroup is None: - raise RuntimeError("Essential taskgroup missing") - - self.client = None - - if not hasattr(__class__, "params_metric"): - __class__.params_metric = Info( - 'params', 'Parameters configuration' - ) - - # FIXME: Maybe outputs information it should not - __class__.params_metric.info({ - k: str(params[k]) - for k in params - }) - - pulsar_host = params.get("pulsar_host", self.default_pulsar_host) - pulsar_listener = params.get("pulsar_listener", None) - pulsar_api_key = params.get("pulsar_api_key", None) - log_level = params.get("log_level", LogLevel.INFO) - - self.config_push_queue = params.get( - "config_push_queue", - default_config_queue - ) - - self.pulsar_host = pulsar_host - self.pulsar_api_key = pulsar_api_key - - if pulsar_api_key: - auth = pulsar.AuthenticationToken(pulsar_api_key) - self.client = pulsar.Client( - pulsar_host, - authentication=auth, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - else: - self.client = pulsar.Client( - pulsar_host, - listener_name=pulsar_listener, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - - self.pulsar_listener = pulsar_listener - - self.config_subscriber = self.client.subscribe( - self.config_push_queue, config_subscriber_id, - consumer_type=pulsar.ConsumerType.Shared, - initial_position=pulsar.InitialPosition.Earliest, - schema=JsonSchema(ConfigPush), - ) - - self.config_handlers = [] - - def __del__(self): - - if hasattr(self, "client"): - if self.client: - self.client.close() - - @staticmethod - def add_args(parser): - - parser.add_argument( - '-p', '--pulsar-host', - default=__class__.default_pulsar_host, - help=f'Pulsar host (default: {__class__.default_pulsar_host})', - ) - - parser.add_argument( - '--pulsar-api-key', - default=__class__.default_pulsar_api_key, - help=f'Pulsar API key', - ) - - parser.add_argument( - '--config-push-queue', - default=default_config_queue, - help=f'Config push queue {default_config_queue}', - ) - - parser.add_argument( - '--pulsar-listener', - help=f'Pulsar listener (default: none)', - ) - - parser.add_argument( - '-l', '--log-level', - type=LogLevel, - default=LogLevel.INFO, - choices=list(LogLevel), - help=f'Output queue (default: info)' - ) - - parser.add_argument( - '--metrics', - action=argparse.BooleanOptionalAction, - default=True, - help=f'Metrics enabled (default: true)', - ) - - parser.add_argument( - '-P', '--metrics-port', - type=int, - default=8000, - help=f'Pulsar host (default: 8000)', - ) - - async def start(self): - pass - - async def run_config_queue(self): - - if self.module == "config.service": - print("I am config-svc, not looking at config queue", flush=True) - return - - print("Config thread running", flush=True) - - while True: - - try: - msg = await asyncio.to_thread( - self.config_subscriber.receive, timeout_millis=2000 - ) - except pulsar.Timeout: - continue - - v = msg.value() - print("Got config version", v.version, flush=True) - - for h in self.config_handlers: - await h(v.version, v.config) - - async def run(self): - while True: - await asyncio.sleep(2) - -# raise RuntimeError("Something should have implemented the run method") - - def subscribe(self, input_queue, subscriber, schema, handler): - - consumer = self.client.subscribe( - input_queue, subscriber, - consumer_type=pulsar.ConsumerType.Shared, - schema=JsonSchema(schema), - ) - - s = Subscription(consumer, handler, self.taskgroup) - - return s - - def publish(self, output_queue, schema): - - producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(schema), - chunking_enabled=True, - ) - - p = Publisher(producer) - - return p - - def set_processor_state(self, flow, state): - - if not hasattr(__class__, "state_metric"): - __class__.state_metric = Enum( - 'processor_state', 'Processor state', - ["flow"], - states=['starting', 'running', 'stopped'] - ) - - __class__.state_metric.labels("flow").state(state) - - def set_pubsub_info(self, flow, info) : - - if not hasattr(__class__, "pubsub_metric"): - __class__.pubsub_metric = Info( - 'pubsub', 'Pub/sub configuration', - ["flow"] - ) - - __class__.pubsub_metric.labels(flow=flow).info(info) - - @classmethod - async def launch_async(cls, args, ident): - - async with asyncio.TaskGroup() as tg: - - p = cls(taskgroup=tg, **args) - - # FIXME: Two sort of 'ident' things going on here? - p.module = ident - p.config_ident = args.get("ident", "FIXME") - - await p.start() - - task1 = tg.create_task(p.run_config_queue()) - task2 = tg.create_task(p.run()) - - @classmethod - def launch(cls, ident, doc): - - parser = argparse.ArgumentParser( - prog=ident, - description=doc - ) - - parser.add_argument( - '--id', - default=ident, - help=f'Configuration identity (default: {ident})', - ) - - cls.add_args(parser) - - args = parser.parse_args() - args = vars(args) - - print(args) - - if args["metrics"]: - start_http_server(args["metrics_port"]) - - while True: - - try: - - asyncio.run(cls.launch_async( - args, ident - )) - - except KeyboardInterrupt: - print("Keyboard interrupt.") - return - - except _pulsar.Interrupted: - print("Pulsar Interrupted.") - return - - except Exception as e: - - print(type(e)) - - print(e.message) - print(e.exceptions) - print(e) - - print("Exception:", e, flush=True) - print("Will retry...", flush=True) - - time.sleep(4) - diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index 6ba2650e..7ea9ef49 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -1,110 +1,45 @@ import asyncio -from pulsar.schema import JsonSchema -import pulsar -from prometheus_client import Histogram, Info, Counter, Enum import time -from . base_processor import BaseProcessor from .. exceptions import TooManyRequests -default_rate_limit_retry = 10 -default_rate_limit_timeout = 7200 +class Consumer: -class Consumer(BaseProcessor): + def __init__( + self, taskgroup, client, queue, subscriber, schema, + handler + ): - def __init__(self, **params): + self.taskgroup = taskgroup + self.client = client + self.queue = queue + self.subscriber = subscriber + self.schema = schema + self.handler = handler - if not hasattr(__class__, "state_metric"): - __class__.state_metric = Enum( - 'processor_state', 'Processor state', - states=['starting', 'running', 'stopped'] - ) - __class__.state_metric.state('starting') + self.running = True + self.task = None - __class__.state_metric.state('starting') + async def start(self): - super(Consumer, self).__init__(**params) + self.running = True - self.subscriber = params.get("subscriber") - self.input_schema = params.get("input_schema") - - self.rate_limit_retry = params.get( - "rate_limit_retry", default_rate_limit_retry - ) - self.rate_limit_timeout = params.get( - "rate_limit_timeout", default_rate_limit_timeout + self.consumer = self.client.subscribe( + self.queue, self.subscriber, self.schema ) - if self.input_schema == None: - raise RuntimeError("input_schema must be specified") + self.task = self.taskgroup.create_task(self.run()) - if not hasattr(__class__, "request_metric"): - __class__.request_metric = Histogram( - 'request_latency', 'Request latency (seconds)', - ["flow"] - ) - - if not hasattr(__class__, "pubsub_metric"): - __class__.pubsub_metric = Info( - 'pubsub', 'Pub/sub configuration', - ["flow"] - ) - - if not hasattr(__class__, "processing_metric"): - __class__.processing_metric = Counter( - 'processing_count', 'Processing count', - ["flow", "status"] - ) - - if not hasattr(__class__, "rate_limit_metric"): - __class__.rate_limit_metric = Counter( - 'rate_limit_count', 'Rate limit event count', - ["flow"] - ) - - # __class__.pubsub_metric.info({ - # "input_queue": self.input_queue, - # "subscriber": self.subscriber, - # "input_schema": self.input_schema.__name__, - # "rate_limit_retry": str(self.rate_limit_retry), - # "rate_limit_timeout": str(self.rate_limit_timeout), - # }) - - # self.consumer = self.client.subscribe( - # self.input_queue, self.subscriber, - # consumer_type=pulsar.ConsumerType.Shared, - # schema=JsonSchema(self.input_schema), - # ) - - self.config_queues = {} - self.config_handlers = [self.on_config_pubsub] - - print("Initialised consumer.", flush=True) - - async def on_config_pubsub(self, version, config): - - print("Configuring using config", version) - - flows = config.get("flows", {}).get(self.ident, {}) - - print(flows) - - print("Configuration completed OK") async def run(self): - __class__.state_metric.state('running') - - while True: - - await asyncio.sleep(1) - - continue + while self.running: msg = await asyncio.to_thread(self.consumer.receive) - expiry = time.time() + self.rate_limit_timeout +# expiry = time.time() + self.rate_limit_timeout + expiry = time.time() + 10 # This loop is for retry on rate-limit / resource limits while True: @@ -117,20 +52,24 @@ class Consumer(BaseProcessor): # be retried self.consumer.negative_acknowledge(msg) - __class__.processing_metric.labels(status="error").inc() + # FIXME +# __class__.processing_metric.labels(status="error").inc() # Break out of retry loop, processes next message break try: - with __class__.request_metric.time(): - await self.handle(msg) + print("Handle...") + # FIXME +# with __class__.request_metric.time(): + await self.handler(msg, self.consumer) + print("Handled.") # Acknowledge successful processing of the message self.consumer.acknowledge(msg) - __class__.processing_metric.labels(status="success").inc() +# __class__.processing_metric.labels(status="success").inc() # Break out of retry loop break @@ -139,14 +78,15 @@ class Consumer(BaseProcessor): print("TooManyRequests: will retry...", flush=True) - __class__.rate_limit_metric.inc() + # FIXME +# __class__.rate_limit_metric.inc() # Sleep time.sleep(self.rate_limit_retry) # Contine from retry loop, just causes a reprocessing continue - + except Exception as e: print("Exception:", e, flush=True) @@ -155,33 +95,8 @@ class Consumer(BaseProcessor): # be retried self.consumer.negative_acknowledge(msg) - __class__.processing_metric.labels(status="error").inc() +# __class__.processing_metric.labels(status="error").inc() # Break out of retry loop, processes next message break - @staticmethod - def add_args(parser, default_input_queue, default_subscriber): - - BaseProcessor.add_args(parser) - - parser.add_argument( - '-s', '--subscriber', - default=default_subscriber, - help=f'Queue subscriber name (default: {default_subscriber})' - ) - - parser.add_argument( - '--rate-limit-retry', - type=int, - default=default_rate_limit_retry, - help=f'Rate limit retry (default: {default_rate_limit_retry})' - ) - - parser.add_argument( - '--rate-limit-timeout', - type=int, - default=default_rate_limit_timeout, - help=f'Rate limit timeout (default: {default_rate_limit_timeout})' - ) - diff --git a/trustgraph-base/trustgraph/base/consumer.py.OLD b/trustgraph-base/trustgraph/base/consumer.py.OLD new file mode 100644 index 00000000..6ba2650e --- /dev/null +++ b/trustgraph-base/trustgraph/base/consumer.py.OLD @@ -0,0 +1,187 @@ + +import asyncio +from pulsar.schema import JsonSchema +import pulsar +from prometheus_client import Histogram, Info, Counter, Enum +import time + +from . base_processor import BaseProcessor +from .. exceptions import TooManyRequests + +default_rate_limit_retry = 10 +default_rate_limit_timeout = 7200 + +class Consumer(BaseProcessor): + + def __init__(self, **params): + + if not hasattr(__class__, "state_metric"): + __class__.state_metric = Enum( + 'processor_state', 'Processor state', + states=['starting', 'running', 'stopped'] + ) + __class__.state_metric.state('starting') + + __class__.state_metric.state('starting') + + super(Consumer, self).__init__(**params) + + self.subscriber = params.get("subscriber") + self.input_schema = params.get("input_schema") + + self.rate_limit_retry = params.get( + "rate_limit_retry", default_rate_limit_retry + ) + self.rate_limit_timeout = params.get( + "rate_limit_timeout", default_rate_limit_timeout + ) + + if self.input_schema == None: + raise RuntimeError("input_schema must be specified") + + if not hasattr(__class__, "request_metric"): + __class__.request_metric = Histogram( + 'request_latency', 'Request latency (seconds)', + ["flow"] + ) + + if not hasattr(__class__, "pubsub_metric"): + __class__.pubsub_metric = Info( + 'pubsub', 'Pub/sub configuration', + ["flow"] + ) + + if not hasattr(__class__, "processing_metric"): + __class__.processing_metric = Counter( + 'processing_count', 'Processing count', + ["flow", "status"] + ) + + if not hasattr(__class__, "rate_limit_metric"): + __class__.rate_limit_metric = Counter( + 'rate_limit_count', 'Rate limit event count', + ["flow"] + ) + + # __class__.pubsub_metric.info({ + # "input_queue": self.input_queue, + # "subscriber": self.subscriber, + # "input_schema": self.input_schema.__name__, + # "rate_limit_retry": str(self.rate_limit_retry), + # "rate_limit_timeout": str(self.rate_limit_timeout), + # }) + + # self.consumer = self.client.subscribe( + # self.input_queue, self.subscriber, + # consumer_type=pulsar.ConsumerType.Shared, + # schema=JsonSchema(self.input_schema), + # ) + + self.config_queues = {} + self.config_handlers = [self.on_config_pubsub] + + print("Initialised consumer.", flush=True) + + async def on_config_pubsub(self, version, config): + + print("Configuring using config", version) + + flows = config.get("flows", {}).get(self.ident, {}) + + print(flows) + + print("Configuration completed OK") + + async def run(self): + + __class__.state_metric.state('running') + + while True: + + await asyncio.sleep(1) + + continue + + msg = await asyncio.to_thread(self.consumer.receive) + + expiry = time.time() + self.rate_limit_timeout + + # This loop is for retry on rate-limit / resource limits + while True: + + if time.time() > expiry: + + print("Gave up waiting for rate-limit retry", flush=True) + + # Message failed to be processed, this causes it to + # be retried + self.consumer.negative_acknowledge(msg) + + __class__.processing_metric.labels(status="error").inc() + + # Break out of retry loop, processes next message + break + + try: + + with __class__.request_metric.time(): + await self.handle(msg) + + # Acknowledge successful processing of the message + self.consumer.acknowledge(msg) + + __class__.processing_metric.labels(status="success").inc() + + # Break out of retry loop + break + + except TooManyRequests: + + print("TooManyRequests: will retry...", flush=True) + + __class__.rate_limit_metric.inc() + + # Sleep + time.sleep(self.rate_limit_retry) + + # Contine from retry loop, just causes a reprocessing + continue + + except Exception as e: + + print("Exception:", e, flush=True) + + # Message failed to be processed, this causes it to + # be retried + self.consumer.negative_acknowledge(msg) + + __class__.processing_metric.labels(status="error").inc() + + # Break out of retry loop, processes next message + break + + @staticmethod + def add_args(parser, default_input_queue, default_subscriber): + + BaseProcessor.add_args(parser) + + parser.add_argument( + '-s', '--subscriber', + default=default_subscriber, + help=f'Queue subscriber name (default: {default_subscriber})' + ) + + parser.add_argument( + '--rate-limit-retry', + type=int, + default=default_rate_limit_retry, + help=f'Rate limit retry (default: {default_rate_limit_retry})' + ) + + parser.add_argument( + '--rate-limit-timeout', + type=int, + default=default_rate_limit_timeout, + help=f'Rate limit timeout (default: {default_rate_limit_timeout})' + ) + diff --git a/trustgraph-base/trustgraph/base/producer.py b/trustgraph-base/trustgraph/base/producer.py index 619670db..9cb1197d 100644 --- a/trustgraph-base/trustgraph/base/producer.py +++ b/trustgraph-base/trustgraph/base/producer.py @@ -1,57 +1,21 @@ -from pulsar.schema import JsonSchema -from prometheus_client import Info, Counter +class Producer: -from . base_processor import BaseProcessor + def __init__(self, client, queue, schema): + self.client = client + self.queue = queue + self.schema = schema -class Producer(BaseProcessor): - - def __init__(self, **params): - - output_queue = params.get("output_queue") - output_schema = params.get("output_schema") - - if not hasattr(__class__, "output_metric"): - __class__.output_metric = Counter( - 'output_count', 'Output items created' - ) - - if not hasattr(__class__, "pubsub_metric"): - __class__.pubsub_metric = Info( - 'pubsub', 'Pub/sub configuration' - ) - - __class__.pubsub_metric.info({ - "output_queue": output_queue, - "output_schema": output_schema.__name__, - }) - - super(Producer, self).__init__(**params) - - if output_schema == None: - raise RuntimeError("output_schema must be specified") - - self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(output_schema), - chunking_enabled=True, - ) + self.producer = None async def send(self, msg, properties={}): + + if self.producer == None: + self.producer = self.client.publish(self.queue, self.schema) + + self.producer.send(msg, properties) - __class__.output_metric.inc() - @staticmethod - def add_args( - parser, default_input_queue, default_subscriber, - default_output_queue, - ): - - BaseProcessor.add_args(parser) - - parser.add_argument( - '-o', '--output-queue', - default=default_output_queue, - help=f'Output queue (default: {default_output_queue})' - ) + # FIXME +# __class__.output_metric.inc() diff --git a/trustgraph-base/trustgraph/base/producer.py.OLD b/trustgraph-base/trustgraph/base/producer.py.OLD new file mode 100644 index 00000000..619670db --- /dev/null +++ b/trustgraph-base/trustgraph/base/producer.py.OLD @@ -0,0 +1,57 @@ + +from pulsar.schema import JsonSchema +from prometheus_client import Info, Counter + +from . base_processor import BaseProcessor + +class Producer(BaseProcessor): + + def __init__(self, **params): + + output_queue = params.get("output_queue") + output_schema = params.get("output_schema") + + if not hasattr(__class__, "output_metric"): + __class__.output_metric = Counter( + 'output_count', 'Output items created' + ) + + if not hasattr(__class__, "pubsub_metric"): + __class__.pubsub_metric = Info( + 'pubsub', 'Pub/sub configuration' + ) + + __class__.pubsub_metric.info({ + "output_queue": output_queue, + "output_schema": output_schema.__name__, + }) + + super(Producer, self).__init__(**params) + + if output_schema == None: + raise RuntimeError("output_schema must be specified") + + self.producer = self.client.create_producer( + topic=output_queue, + schema=JsonSchema(output_schema), + chunking_enabled=True, + ) + + async def send(self, msg, properties={}): + self.producer.send(msg, properties) + __class__.output_metric.inc() + + @staticmethod + def add_args( + parser, default_input_queue, default_subscriber, + default_output_queue, + ): + + BaseProcessor.add_args(parser) + + parser.add_argument( + '-o', '--output-queue', + default=default_output_queue, + help=f'Output queue (default: {default_output_queue})' + ) + diff --git a/trustgraph-base/trustgraph/base/pubsub.py b/trustgraph-base/trustgraph/base/pubsub.py new file mode 100644 index 00000000..c9c4ed92 --- /dev/null +++ b/trustgraph-base/trustgraph/base/pubsub.py @@ -0,0 +1,96 @@ + +import os +import pulsar +import uuid +from pulsar.schema import JsonSchema + +from .. log_level import LogLevel + +class PulsarClient: + + default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650') + default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) + + def __init__(self, **params): + + self.client = None + + pulsar_host = params.get("pulsar_host", self.default_pulsar_host) + pulsar_listener = params.get("pulsar_listener", None) + pulsar_api_key = params.get( + "pulsar_api_key", + self.default_pulsar_api_key + ) + log_level = params.get("log_level", LogLevel.INFO) + + self.pulsar_host = pulsar_host + self.pulsar_api_key = pulsar_api_key + + if pulsar_api_key: + auth = pulsar.AuthenticationToken(pulsar_api_key) + self.client = pulsar.Client( + pulsar_host, + authentication=auth, + logger=pulsar.ConsoleLogger(log_level.to_pulsar()) + ) + else: + self.client = pulsar.Client( + pulsar_host, + listener_name=pulsar_listener, + logger=pulsar.ConsoleLogger(log_level.to_pulsar()) + ) + + self.pulsar_listener = pulsar_listener + + def close(self): + self.client.close() + + def __del__(self): + + if hasattr(self, "client"): + if self.client: + self.client.close() + + def subscribe(self, queue, subscriber, schema): + + return self.client.subscribe( + queue, subscriber, + consumer_type=pulsar.ConsumerType.Shared, + schema=JsonSchema(schema), + ) + + def publish(self, queue, schema): + + return self.client.create_producer( + topic=queue, + schema=JsonSchema(schema), + chunking_enabled=True, + ) + + @staticmethod + def add_args(parser): + + parser.add_argument( + '-p', '--pulsar-host', + default=__class__.default_pulsar_host, + help=f'Pulsar host (default: {__class__.default_pulsar_host})', + ) + + parser.add_argument( + '--pulsar-api-key', + default=__class__.default_pulsar_api_key, + help=f'Pulsar API key', + ) + + parser.add_argument( + '--pulsar-listener', + help=f'Pulsar listener (default: none)', + ) + + parser.add_argument( + '-l', '--log-level', + type=LogLevel, + default=LogLevel.INFO, + choices=list(LogLevel), + help=f'Output queue (default: info)' + ) diff --git a/trustgraph-flow/trustgraph/config/service/service.py b/trustgraph-flow/trustgraph/config/service/service.py index 2c2f143b..b05ae446 100644 --- a/trustgraph-flow/trustgraph/config/service/service.py +++ b/trustgraph-flow/trustgraph/config/service/service.py @@ -43,7 +43,6 @@ class Processor(BaseProcessor): push_queue = params.get("push_queue", default_push_queue) subscriber = params.get("subscriber", default_subscriber) - input_schema = ConfigRequest output_schema = ConfigResponse push_schema = ConfigResponse