diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index e1e3f367..fc582d07 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -37,6 +37,7 @@ from . triples_import import TriplesImport from . graph_embeddings_import import GraphEmbeddingsImport from . document_embeddings_import import DocumentEmbeddingsImport from . entity_contexts_import import EntityContextsImport +from . objects_import import ObjectsImport from . core_export import CoreExport from . core_import import CoreImport @@ -82,6 +83,7 @@ import_dispatchers = { "graph-embeddings": GraphEmbeddingsImport, "document-embeddings": DocumentEmbeddingsImport, "entity-contexts": EntityContextsImport, + "objects": ObjectsImport, } class DispatcherWrapper: diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/objects_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/objects_import.py new file mode 100644 index 00000000..f7891f96 --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/dispatch/objects_import.py @@ -0,0 +1,70 @@ +import asyncio +import uuid +import logging +from aiohttp import WSMsgType + +from ... schema import Metadata +from ... schema import ExtractedObject +from ... base import Publisher + +from . serialize import to_subgraph + +# Module logger +logger = logging.getLogger(__name__) + +class ObjectsImport: + + def __init__( + self, ws, running, pulsar_client, queue + ): + + self.ws = ws + self.running = running + + self.publisher = Publisher( + pulsar_client, topic = queue, schema = ExtractedObject + ) + + async def start(self): + await self.publisher.start() + + async def destroy(self): + # Step 1: Stop accepting new messages + self.running.stop() + + # Step 2: Wait for publisher to drain its queue + logger.info("Draining publisher queue...") + await self.publisher.stop() + + # Step 3: Close websocket only after queue is drained + if self.ws: + await self.ws.close() + + async def receive(self, msg): + + data = msg.json() + + elt = ExtractedObject( + metadata=Metadata( + id=data["metadata"]["id"], + metadata=to_subgraph(data["metadata"].get("metadata", [])), + user=data["metadata"]["user"], + collection=data["metadata"]["collection"], + ), + schema_name=data["schema_name"], + values=data["values"], + confidence=data.get("confidence", 1.0), + source_span=data.get("source_span", ""), + ) + + await self.publisher.send(None, elt) + + async def run(self): + + while self.running.get(): + await asyncio.sleep(0.5) + + if self.ws: + await self.ws.close() + + self.ws = None \ No newline at end of file