diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_stream.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_stream.py new file mode 100644 index 00000000..e70bf6de --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/dispatch/document_stream.py @@ -0,0 +1,65 @@ + +import asyncio +import uuid +import logging +from . librarian import LibrarianRequestor + +# Module logger +logger = logging.getLogger(__name__) + +class DocumentStreamExport: + + def __init__(self, backend): + self.backend = backend + + async def process(self, data, error, ok, request): + + user = request.query.get("user") + document_id = request.query.get("document-id") + chunk_size = int(request.query.get("chunk-size", 1024 * 1024)) + + if not user or not document_id: + return await error("Missing required parameters: user, document-id") + + response = await ok() + + lr = LibrarianRequestor( + backend=self.backend, + consumer="api-gateway-doc-stream-" + str(uuid.uuid4()), + subscriber="api-gateway-doc-stream-" + str(uuid.uuid4()), + ) + + try: + + await lr.start() + + async def responder(resp, fin): + if "content" in resp: + content = resp["content"] + # Content is base64 encoded, write as-is for client to decode + # Or decode here and write raw bytes + import base64 + chunk_data = base64.b64decode(content) + await response.write(chunk_data) + + await lr.process( + { + "operation": "stream-document", + "user": user, + "document-id": document_id, + "chunk-size": chunk_size, + }, + responder + ) + + except Exception as e: + + logger.error(f"Document stream exception: {e}", exc_info=True) + + finally: + + await lr.stop() + + await response.write_eof() + + return response diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 35edad76..d068ecef 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -45,6 +45,7 @@ from . rows_import import RowsImport from . core_export import CoreExport from . core_import import CoreImport +from . document_stream import DocumentStreamExport from . mux import Mux @@ -135,6 +136,14 @@ class DispatcherManager: def dispatch_core_import(self): return DispatcherWrapper(self.process_core_import) + def dispatch_document_stream(self): + return DispatcherWrapper(self.process_document_stream) + + async def process_document_stream(self, data, error, ok, request): + + ds = DocumentStreamExport(self.backend) + return await ds.process(data, error, ok, request) + async def process_core_import(self, data, error, ok, request): ci = CoreImport(self.backend) diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py b/trustgraph-flow/trustgraph/gateway/endpoint/manager.py index f9616d9a..ff92dca2 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/manager.py @@ -64,6 +64,12 @@ class EndpointManager: method = "GET", dispatcher = dispatcher_manager.dispatch_core_export(), ), + StreamEndpoint( + endpoint_path = "/api/v1/document-stream", + auth = auth, + method = "GET", + dispatcher = dispatcher_manager.dispatch_document_stream(), + ), ] def add_routes(self, app):