From 2f0ee1c09b1452a0d2c9aee4d452daae73b64ef0 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 3 Dec 2024 18:01:42 +0000 Subject: [PATCH] Remove aiopulsar dependency --- trustgraph-flow/setup.py | 1 - .../trustgraph/api/gateway/endpoint.py | 2 + .../api/gateway/graph_embeddings_stream.py | 3 +- .../trustgraph/api/gateway/publisher.py | 3 - .../trustgraph/api/gateway/subscriber.py | 56 +++++++++++++------ .../trustgraph/api/gateway/triples_stream.py | 5 +- 6 files changed, 47 insertions(+), 23 deletions(-) diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 65bb7326..e6c732a3 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -59,7 +59,6 @@ setuptools.setup( "ibis", "jsonschema", "aiohttp", - "aiopulsar-py", "pinecone[grpc]", ], scripts=[ diff --git a/trustgraph-flow/trustgraph/api/gateway/endpoint.py b/trustgraph-flow/trustgraph/api/gateway/endpoint.py index 26bbc364..2b246361 100644 --- a/trustgraph-flow/trustgraph/api/gateway/endpoint.py +++ b/trustgraph-flow/trustgraph/api/gateway/endpoint.py @@ -93,6 +93,8 @@ class ServiceEndpoint: except Exception as e: raise RuntimeError("Timeout") + print(resp) + if resp.error: print("Error") return web.json_response( diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py index 7b3a9524..3d4efd45 100644 --- a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py @@ -1,5 +1,6 @@ import asyncio +import queue from pulsar.schema import JsonSchema import uuid @@ -43,7 +44,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): resp = await asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_graph_embeddings(resp)) - except TimeoutError: + except queue.Empty: continue except Exception as e: diff --git a/trustgraph-flow/trustgraph/api/gateway/publisher.py b/trustgraph-flow/trustgraph/api/gateway/publisher.py index 7e9c9677..89c612ce 100644 --- a/trustgraph-flow/trustgraph/api/gateway/publisher.py +++ b/trustgraph-flow/trustgraph/api/gateway/publisher.py @@ -38,8 +38,6 @@ class Publisher: id, item = self.q.get() - print("THING ON Q") - if id: producer.send(item, { "id": id }) else: @@ -53,4 +51,3 @@ class Publisher: def send(self, id, msg): self.q.put((id, msg)) - print("PUT ON Q") diff --git a/trustgraph-flow/trustgraph/api/gateway/subscriber.py b/trustgraph-flow/trustgraph/api/gateway/subscriber.py index ae947aff..cccfc5b4 100644 --- a/trustgraph-flow/trustgraph/api/gateway/subscriber.py +++ b/trustgraph-flow/trustgraph/api/gateway/subscriber.py @@ -7,7 +7,7 @@ import time class Subscriber: def __init__(self, pulsar_host, topic, subscription, consumer_name, - schema=None, max_size=10): + schema=None, max_size=100): self.pulsar_host = pulsar_host self.topic = topic self.subscription = subscription @@ -16,6 +16,7 @@ class Subscriber: self.q = {} self.full = {} self.max_size = max_size + self.lock = threading.Lock() def start(self): self.task = threading.Thread(target=self.run) @@ -37,13 +38,11 @@ class Subscriber: consumer_name=self.consumer_name, schema=self.schema, ) - + while True: msg = consumer.receive() - print("Received", msg) - # Acknowledge successful reception of the message consumer.acknowledge(msg) @@ -53,11 +52,20 @@ class Subscriber: id = None value = msg.value() - if id in self.q: - self.q[id].put(value) - for q in self.full.values(): - q.put(value) + with self.lock: + + if id in self.q: + try: + self.q[id].put(value, timeout=0.5) + except: + pass + + for q in self.full.values(): + try: + q.put(value, timeout=0.5) + except: + pass except Exception as e: print("Exception:", e, flush=True) @@ -66,20 +74,36 @@ class Subscriber: time.sleep(2) def subscribe(self, id): - q = queue.Queue(maxsize=self.max_size) - self.q[id] = q + + with self.lock: + + q = queue.Queue(maxsize=self.max_size) + self.q[id] = q + return q def unsubscribe(self, id): - if id in self.q: - del self.q[id] + + with self.lock: + + if id in self.q: +# self.q[id].shutdown(immediate=True) + del self.q[id] def subscribe_all(self, id): - q = queue.Queue(maxsize=self.max_size) - self.full[id] = q + + with self.lock: + + q = queue.Queue(maxsize=self.max_size) + self.full[id] = q + return q def unsubscribe_all(self, id): - if id in self.full: - del self.full[id] + + with self.lock: + + if id in self.full: +# self.full[id].shutdown(immediate=True) + del self.full[id] diff --git a/trustgraph-flow/trustgraph/api/gateway/triples_stream.py b/trustgraph-flow/trustgraph/api/gateway/triples_stream.py index 0875938e..4638e08d 100644 --- a/trustgraph-flow/trustgraph/api/gateway/triples_stream.py +++ b/trustgraph-flow/trustgraph/api/gateway/triples_stream.py @@ -1,5 +1,6 @@ import asyncio +import queue from pulsar.schema import JsonSchema import uuid @@ -38,10 +39,10 @@ class TriplesStreamEndpoint(SocketEndpoint): while running.get(): try: - resp = asyncio.to_thread(q.get, timeout=0.5) + resp = await asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_triples(resp)) - except TimeoutError: + except queue.Empty: continue except Exception as e: