diff --git a/trustgraph-base/trustgraph/base/request_response_spec.py b/trustgraph-base/trustgraph/base/request_response_spec.py index 82574e9d..e4c80c74 100644 --- a/trustgraph-base/trustgraph/base/request_response_spec.py +++ b/trustgraph-base/trustgraph/base/request_response_spec.py @@ -14,7 +14,7 @@ logger = logging.getLogger(__name__) class RequestResponse(Subscriber): def __init__( - self, client, subscription, consumer_name, + self, backend, subscription, consumer_name, request_topic, request_schema, request_metrics, response_topic, response_schema, @@ -22,7 +22,7 @@ class RequestResponse(Subscriber): ): super(RequestResponse, self).__init__( - client = client, + backend = backend, subscription = subscription, consumer_name = consumer_name, topic = response_topic, @@ -31,7 +31,7 @@ class RequestResponse(Subscriber): ) self.producer = Producer( - client = client, + backend = backend, topic = request_topic, schema = request_schema, metrics = request_metrics, @@ -126,7 +126,7 @@ class RequestResponseSpec(Spec): ) rr = self.impl( - client = processor.pulsar_client, + backend = processor.pubsub, # Make subscription names unique, so that all subscribers get # to see all response messages diff --git a/trustgraph-base/trustgraph/base/subscriber_spec.py b/trustgraph-base/trustgraph/base/subscriber_spec.py index 7dca09db..b408366c 100644 --- a/trustgraph-base/trustgraph/base/subscriber_spec.py +++ b/trustgraph-base/trustgraph/base/subscriber_spec.py @@ -16,7 +16,7 @@ class SubscriberSpec(Spec): ) subscriber = Subscriber( - client = processor.pulsar_client, + backend = processor.pubsub, topic = definition[self.name], subscription = flow.id, consumer_name = flow.id, diff --git a/trustgraph-cli/trustgraph/cli/dump_queues.py b/trustgraph-cli/trustgraph/cli/dump_queues.py index 93151858..0a298450 100644 --- a/trustgraph-cli/trustgraph/cli/dump_queues.py +++ b/trustgraph-cli/trustgraph/cli/dump_queues.py @@ -17,6 +17,7 @@ from datetime import datetime import argparse from trustgraph.base.subscriber import Subscriber +from trustgraph.base.pubsub import get_pubsub def format_message(queue_name, msg): """Format a message with timestamp and queue name.""" @@ -167,11 +168,11 @@ async def async_main(queues, output_file, pulsar_host, listener_name, subscriber print(f"Mode: {'append' if append_mode else 'overwrite'}") print(f"Press Ctrl+C to stop\n") - # Connect to Pulsar + # Create backend connection try: - client = pulsar.Client(pulsar_host, listener_name=listener_name) + backend = get_pubsub(pulsar_host=pulsar_host, pulsar_listener=listener_name, pubsub_backend='pulsar') except Exception as e: - print(f"Error connecting to Pulsar at {pulsar_host}: {e}", file=sys.stderr) + print(f"Error connecting to backend at {pulsar_host}: {e}", file=sys.stderr) sys.exit(1) # Create Subscribers and central queue @@ -181,7 +182,7 @@ async def async_main(queues, output_file, pulsar_host, listener_name, subscriber for queue_name in queues: try: sub = Subscriber( - client=client, + backend=backend, topic=queue_name, subscription=subscriber_name, consumer_name=f"{subscriber_name}-{queue_name}", @@ -195,7 +196,7 @@ async def async_main(queues, output_file, pulsar_host, listener_name, subscriber if not subscribers: print("\nNo subscribers created. Exiting.", file=sys.stderr) - client.close() + backend.close() sys.exit(1) print(f"\nListening for messages...\n") @@ -256,7 +257,7 @@ async def async_main(queues, output_file, pulsar_host, listener_name, subscriber # Clean shutdown of Subscribers for _, sub in subscribers: await sub.stop() - client.close() + backend.close() print(f"\nMessages logged to: {output_file}") diff --git a/trustgraph-flow/trustgraph/config/service/service.py b/trustgraph-flow/trustgraph/config/service/service.py index 414ad847..455e567e 100644 --- a/trustgraph-flow/trustgraph/config/service/service.py +++ b/trustgraph-flow/trustgraph/config/service/service.py @@ -112,7 +112,7 @@ class Processor(AsyncProcessor): self.config_request_consumer = Consumer( taskgroup = self.taskgroup, - client = self.pulsar_client, + backend = self.pubsub, flow = None, topic = config_request_queue, subscriber = id, @@ -122,14 +122,14 @@ class Processor(AsyncProcessor): ) self.config_response_producer = Producer( - client = self.pulsar_client, + backend = self.pubsub, topic = config_response_queue, schema = ConfigResponse, metrics = config_response_metrics, ) self.config_push_producer = Producer( - client = self.pulsar_client, + backend = self.pubsub, topic = config_push_queue, schema = ConfigPush, metrics = config_push_metrics, @@ -137,7 +137,7 @@ class Processor(AsyncProcessor): self.flow_request_consumer = Consumer( taskgroup = self.taskgroup, - client = self.pulsar_client, + backend = self.pubsub, flow = None, topic = flow_request_queue, subscriber = id, @@ -147,7 +147,7 @@ class Processor(AsyncProcessor): ) self.flow_response_producer = Producer( - client = self.pulsar_client, + backend = self.pubsub, topic = flow_response_queue, schema = FlowResponse, metrics = flow_response_metrics, diff --git a/trustgraph-flow/trustgraph/cores/service.py b/trustgraph-flow/trustgraph/cores/service.py index 4b1573a9..18154fc5 100755 --- a/trustgraph-flow/trustgraph/cores/service.py +++ b/trustgraph-flow/trustgraph/cores/service.py @@ -84,7 +84,7 @@ class Processor(AsyncProcessor): self.knowledge_request_consumer = Consumer( taskgroup = self.taskgroup, - client = self.pulsar_client, + backend = self.pubsub, flow = None, topic = knowledge_request_queue, subscriber = id, @@ -94,7 +94,7 @@ class Processor(AsyncProcessor): ) self.knowledge_response_producer = Producer( - client = self.pulsar_client, + backend = self.pubsub, topic = knowledge_response_queue, schema = KnowledgeResponse, metrics = knowledge_response_metrics, diff --git a/trustgraph-flow/trustgraph/gateway/config/receiver.py b/trustgraph-flow/trustgraph/gateway/config/receiver.py index 0427e236..bdd123a9 100755 --- a/trustgraph-flow/trustgraph/gateway/config/receiver.py +++ b/trustgraph-flow/trustgraph/gateway/config/receiver.py @@ -34,9 +34,9 @@ logger.setLevel(logging.INFO) class ConfigReceiver: - def __init__(self, pulsar_client): + def __init__(self, backend): - self.pulsar_client = pulsar_client + self.backend = backend self.flow_handlers = [] @@ -104,8 +104,8 @@ class ConfigReceiver: self.config_cons = Consumer( taskgroup = tg, flow = None, - client = self.pulsar_client, - subscriber = f"gateway-{id}", + backend = self.backend, + subscriber = f"gateway-{id}", topic = config_push_queue, schema = ConfigPush, handler = self.on_config, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py b/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py index 61b0bcbc..62626046 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/core_export.py @@ -10,9 +10,9 @@ logger = logging.getLogger(__name__) class CoreExport: - def __init__(self, pulsar_client): - self.pulsar_client = pulsar_client - + def __init__(self, backend): + self.backend = backend + async def process(self, data, error, ok, request): id = request.query["id"] @@ -21,7 +21,7 @@ class CoreExport: response = await ok() kr = KnowledgeRequestor( - pulsar_client = self.pulsar_client, + backend = self.backend, consumer = "api-gateway-core-export-" + str(uuid.uuid4()), subscriber = "api-gateway-core-export-" + str(uuid.uuid4()), ) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_export.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_export.py index f7d53005..8866972d 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_export.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_export.py @@ -15,12 +15,12 @@ logger = logging.getLogger(__name__) class DocumentEmbeddingsExport: def __init__( - self, ws, running, pulsar_client, queue, consumer, subscriber + self, ws, running, backend, queue, consumer, subscriber ): self.ws = ws self.running = running - self.pulsar_client = pulsar_client + self.backend = backend self.queue = queue self.consumer = consumer self.subscriber = subscriber @@ -48,9 +48,9 @@ class DocumentEmbeddingsExport: async def run(self): """Enhanced run with better error handling""" self.subs = Subscriber( - client = self.pulsar_client, + backend = self.backend, topic = self.queue, - consumer_name = self.consumer, + consumer_name = self.consumer, subscription = self.subscriber, schema = DocumentEmbeddings, backpressure_strategy = "block" # Configurable diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/entity_contexts_export.py b/trustgraph-flow/trustgraph/gateway/dispatch/entity_contexts_export.py index 2be9c703..c03bdda6 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/entity_contexts_export.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/entity_contexts_export.py @@ -15,12 +15,12 @@ logger = logging.getLogger(__name__) class EntityContextsExport: def __init__( - self, ws, running, pulsar_client, queue, consumer, subscriber + self, ws, running, backend, queue, consumer, subscriber ): self.ws = ws self.running = running - self.pulsar_client = pulsar_client + self.backend = backend self.queue = queue self.consumer = consumer self.subscriber = subscriber @@ -48,9 +48,9 @@ class EntityContextsExport: async def run(self): """Enhanced run with better error handling""" self.subs = Subscriber( - client = self.pulsar_client, + backend = self.backend, topic = self.queue, - consumer_name = self.consumer, + consumer_name = self.consumer, subscription = self.subscriber, schema = EntityContexts, backpressure_strategy = "block" # Configurable diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_export.py b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_export.py index d4abec73..d6d7a1c5 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_export.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_export.py @@ -15,12 +15,12 @@ logger = logging.getLogger(__name__) class GraphEmbeddingsExport: def __init__( - self, ws, running, pulsar_client, queue, consumer, subscriber + self, ws, running, backend, queue, consumer, subscriber ): self.ws = ws self.running = running - self.pulsar_client = pulsar_client + self.backend = backend self.queue = queue self.consumer = consumer self.subscriber = subscriber @@ -48,9 +48,9 @@ class GraphEmbeddingsExport: async def run(self): """Enhanced run with better error handling""" self.subs = Subscriber( - client = self.pulsar_client, + backend = self.backend, topic = self.queue, - consumer_name = self.consumer, + consumer_name = self.consumer, subscription = self.subscriber, schema = GraphEmbeddings, backpressure_strategy = "block" # Configurable diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index d35e5525..7978c3a0 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -98,9 +98,9 @@ class DispatcherWrapper: class DispatcherManager: - def __init__(self, pulsar_client, config_receiver, prefix="api-gateway", + def __init__(self, backend, config_receiver, prefix="api-gateway", queue_overrides=None): - self.pulsar_client = pulsar_client + self.backend = backend self.config_receiver = config_receiver self.config_receiver.add_handler(self) self.prefix = prefix @@ -161,7 +161,7 @@ class DispatcherManager: response_queue = self.queue_overrides[kind].get("response") dispatcher = global_dispatchers[kind]( - pulsar_client = self.pulsar_client, + backend = self.backend, timeout = 120, consumer = f"{self.prefix}-{kind}-request", subscriber = f"{self.prefix}-{kind}-request", @@ -216,7 +216,7 @@ class DispatcherManager: id = str(uuid.uuid4()) dispatcher = import_dispatchers[kind]( - pulsar_client = self.pulsar_client, + backend = self.backend, ws = ws, running = running, queue = qconfig, @@ -254,7 +254,7 @@ class DispatcherManager: id = str(uuid.uuid4()) dispatcher = export_dispatchers[kind]( - pulsar_client = self.pulsar_client, + backend = self.backend, ws = ws, running = running, queue = qconfig, @@ -296,7 +296,7 @@ class DispatcherManager: if kind in request_response_dispatchers: dispatcher = request_response_dispatchers[kind]( - pulsar_client = self.pulsar_client, + backend = self.backend, request_queue = qconfig["request"], response_queue = qconfig["response"], timeout = 120, @@ -305,7 +305,7 @@ class DispatcherManager: ) elif kind in sender_dispatchers: dispatcher = sender_dispatchers[kind]( - pulsar_client = self.pulsar_client, + backend = self.backend, queue = qconfig, ) else: diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py index ff91e461..69fc588d 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py @@ -15,12 +15,12 @@ logger = logging.getLogger(__name__) class TriplesExport: def __init__( - self, ws, running, pulsar_client, queue, consumer, subscriber + self, ws, running, backend, queue, consumer, subscriber ): self.ws = ws self.running = running - self.pulsar_client = pulsar_client + self.backend = backend self.queue = queue self.consumer = consumer self.subscriber = subscriber @@ -48,9 +48,9 @@ class TriplesExport: async def run(self): """Enhanced run with better error handling""" self.subs = Subscriber( - client = self.pulsar_client, + backend = self.backend, topic = self.queue, - consumer_name = self.consumer, + consumer_name = self.consumer, subscription = self.subscriber, schema = Triples, backpressure_strategy = "block" # Configurable diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index 6303e4fc..e49b083c 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -80,7 +80,7 @@ class Api: else: self.auth = Authenticator(allow_all=True) - self.config_receiver = ConfigReceiver(self.pulsar_client) + self.config_receiver = ConfigReceiver(self.pubsub_backend) # Build queue overrides dictionary from CLI arguments queue_overrides = {} @@ -126,7 +126,7 @@ class Api: queue_overrides["librarian"]["response"] = librarian_resp self.dispatcher_manager = DispatcherManager( - pulsar_client = self.pulsar_client, + backend = self.pubsub_backend, config_receiver = self.config_receiver, prefix = "gateway", queue_overrides = queue_overrides, diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index e3ec6977..55f7ec27 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -143,7 +143,7 @@ class Processor(AsyncProcessor): self.librarian_request_consumer = Consumer( taskgroup = self.taskgroup, - client = self.pulsar_client, + backend = self.pubsub, flow = None, topic = librarian_request_queue, subscriber = id, @@ -153,7 +153,7 @@ class Processor(AsyncProcessor): ) self.librarian_response_producer = Producer( - client = self.pulsar_client, + backend = self.pubsub, topic = librarian_response_queue, schema = LibrarianResponse, metrics = librarian_response_metrics, @@ -161,7 +161,7 @@ class Processor(AsyncProcessor): self.collection_request_consumer = Consumer( taskgroup = self.taskgroup, - client = self.pulsar_client, + backend = self.pubsub, flow = None, topic = collection_request_queue, subscriber = id, @@ -171,7 +171,7 @@ class Processor(AsyncProcessor): ) self.collection_response_producer = Producer( - client = self.pulsar_client, + backend = self.pubsub, topic = collection_response_queue, schema = CollectionManagementResponse, metrics = collection_response_metrics, @@ -183,7 +183,7 @@ class Processor(AsyncProcessor): ) self.config_request_producer = Producer( - client = self.pulsar_client, + backend = self.pubsub, topic = config_request_queue, schema = ConfigRequest, metrics = config_request_metrics, @@ -195,7 +195,7 @@ class Processor(AsyncProcessor): self.config_response_consumer = Consumer( taskgroup = self.taskgroup, - client = self.pulsar_client, + backend = self.pubsub, flow = None, topic = config_response_queue, subscriber = f"{id}-config", diff --git a/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py b/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py index c5b8af06..05b8100d 100644 --- a/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py @@ -78,7 +78,7 @@ class Processor(FlowProcessor): # Create storage management consumer self.storage_request_consumer = Consumer( taskgroup=self.taskgroup, - client=self.pulsar_client, + backend=self.pubsub, flow=None, topic=object_storage_management_topic, subscriber=f"{id}-storage", @@ -89,7 +89,7 @@ class Processor(FlowProcessor): # Create storage management response producer self.storage_response_producer = Producer( - client=self.pulsar_client, + backend=self.pubsub, topic=storage_management_response_topic, schema=StorageManagementResponse, metrics=storage_response_metrics,