diff --git a/trustgraph-base/trustgraph/api/library.py b/trustgraph-base/trustgraph/api/library.py index e50dc0aa..cfa55b57 100644 --- a/trustgraph-base/trustgraph/api/library.py +++ b/trustgraph-base/trustgraph/api/library.py @@ -6,6 +6,7 @@ including document storage, metadata management, and processing workflow coordin """ import datetime +import math import time import base64 import logging @@ -17,6 +18,12 @@ from . exceptions import * logger = logging.getLogger(__name__) +# Threshold for switching to chunked upload (10MB) +CHUNKED_UPLOAD_THRESHOLD = 10 * 1024 * 1024 + +# Default chunk size (5MB - S3 multipart minimum) +DEFAULT_CHUNK_SIZE = 5 * 1024 * 1024 + def to_value(x): """Convert wire format to Uri or Literal.""" @@ -67,13 +74,14 @@ class Library: def add_document( self, document, id, metadata, user, title, comments, - kind="text/plain", tags=[], + kind="text/plain", tags=[], on_progress=None, ): """ Add a document to the library. Stores a document with associated metadata in the library for - retrieval and processing. + retrieval and processing. For large documents (> 10MB), automatically + uses chunked upload for better reliability and progress tracking. Args: document: Document content as bytes @@ -84,6 +92,7 @@ class Library: comments: Document description or comments kind: MIME type of the document (default: "text/plain") tags: List of tags for categorization (default: []) + on_progress: Optional callback(bytes_sent, total_bytes) for progress updates Returns: dict: Response from the add operation @@ -107,6 +116,22 @@ class Library: kind="application/pdf", tags=["research", "physics"] ) + + # Add a large document with progress tracking + def progress(sent, total): + print(f"Uploaded {sent}/{total} bytes ({100*sent//total}%)") + + with open("large_document.pdf", "rb") as f: + library.add_document( + document=f.read(), + id="large-doc-001", + metadata=[], + user="trustgraph", + title="Large Document", + comments="A very large document", + kind="application/pdf", + on_progress=progress + ) ``` """ @@ -124,6 +149,21 @@ class Library: if not title: title = "" if not comments: comments = "" + # Check if we should use chunked upload + if len(document) >= CHUNKED_UPLOAD_THRESHOLD: + return self._add_document_chunked( + document=document, + id=id, + metadata=metadata, + user=user, + title=title, + comments=comments, + kind=kind, + tags=tags, + on_progress=on_progress, + ) + + # Small document: use single operation (existing behavior) triples = [] def emit(t): @@ -167,6 +207,101 @@ class Library: return self.request(input) + def _add_document_chunked( + self, document, id, metadata, user, title, comments, + kind, tags, on_progress=None, + ): + """ + Add a large document using chunked upload. + + Internal method that handles multipart upload for large documents. + """ + total_size = len(document) + chunk_size = DEFAULT_CHUNK_SIZE + + logger.info(f"Starting chunked upload for document {id} ({total_size} bytes)") + + # Begin upload session + begin_request = { + "operation": "begin-upload", + "document-metadata": { + "id": id, + "time": int(time.time()), + "kind": kind, + "title": title, + "comments": comments, + "user": user, + "tags": tags, + }, + "total-size": total_size, + "chunk-size": chunk_size, + } + + begin_response = self.request(begin_request) + + upload_id = begin_response.get("upload-id") + if not upload_id: + raise RuntimeError("Failed to begin upload: no upload_id returned") + + actual_chunk_size = begin_response.get("chunk-size", chunk_size) + total_chunks = begin_response.get("total-chunks", math.ceil(total_size / actual_chunk_size)) + + logger.info(f"Upload session {upload_id} created, {total_chunks} chunks") + + try: + # Upload chunks + bytes_sent = 0 + for chunk_index in range(total_chunks): + start = chunk_index * actual_chunk_size + end = min(start + actual_chunk_size, total_size) + chunk_data = document[start:end] + + chunk_request = { + "operation": "upload-chunk", + "upload-id": upload_id, + "chunk-index": chunk_index, + "content": base64.b64encode(chunk_data).decode("utf-8"), + "user": user, + } + + chunk_response = self.request(chunk_request) + + bytes_sent = end + + # Call progress callback if provided + if on_progress: + on_progress(bytes_sent, total_size) + + logger.debug(f"Chunk {chunk_index + 1}/{total_chunks} uploaded") + + # Complete upload + complete_request = { + "operation": "complete-upload", + "upload-id": upload_id, + "user": user, + } + + complete_response = self.request(complete_request) + + logger.info(f"Chunked upload completed for document {id}") + + return complete_response + + except Exception as e: + # Try to abort on failure + logger.error(f"Chunked upload failed: {e}") + try: + abort_request = { + "operation": "abort-upload", + "upload-id": upload_id, + "user": user, + } + self.request(abort_request) + logger.info(f"Aborted failed upload {upload_id}") + except Exception as abort_error: + logger.warning(f"Failed to abort upload: {abort_error}") + raise + def get_documents(self, user): """ List all documents for a user. @@ -535,3 +670,192 @@ class Library: logger.error("Failed to parse processing list response", exc_info=True) raise ProtocolException(f"Response not formatted correctly") + # Chunked upload management methods + + def get_pending_uploads(self, user): + """ + List all pending (in-progress) uploads for a user. + + Retrieves information about chunked uploads that have been started + but not yet completed. + + Args: + user: User identifier + + Returns: + list[dict]: List of pending upload information + + Example: + ```python + library = api.library() + pending = library.get_pending_uploads(user="trustgraph") + + for upload in pending: + print(f"Upload {upload['upload_id']}:") + print(f" Document: {upload['document_id']}") + print(f" Progress: {upload['chunks_received']}/{upload['total_chunks']}") + ``` + """ + input = { + "operation": "list-uploads", + "user": user, + } + + response = self.request(input) + + return response.get("upload-sessions", []) + + def get_upload_status(self, upload_id, user): + """ + Get the status of a specific upload. + + Retrieves detailed status information about a chunked upload, + including which chunks have been received and which are missing. + + Args: + upload_id: Upload session identifier + user: User identifier + + Returns: + dict: Upload status information including: + - upload_id: The upload session ID + - state: "in-progress", "completed", or "expired" + - chunks_received: Number of chunks received + - total_chunks: Total number of chunks expected + - received_chunks: List of received chunk indices + - missing_chunks: List of missing chunk indices + - bytes_received: Total bytes received + - total_bytes: Total expected bytes + + Example: + ```python + library = api.library() + status = library.get_upload_status( + upload_id="abc-123", + user="trustgraph" + ) + + if status['state'] == 'in-progress': + print(f"Missing chunks: {status['missing_chunks']}") + ``` + """ + input = { + "operation": "get-upload-status", + "upload-id": upload_id, + "user": user, + } + + return self.request(input) + + def abort_upload(self, upload_id, user): + """ + Abort an in-progress upload. + + Cancels a chunked upload and cleans up any uploaded chunks. + + Args: + upload_id: Upload session identifier + user: User identifier + + Returns: + dict: Empty response on success + + Example: + ```python + library = api.library() + library.abort_upload(upload_id="abc-123", user="trustgraph") + ``` + """ + input = { + "operation": "abort-upload", + "upload-id": upload_id, + "user": user, + } + + return self.request(input) + + def resume_upload(self, upload_id, document, user, on_progress=None): + """ + Resume an interrupted upload. + + Continues a chunked upload that was previously interrupted, + uploading only the missing chunks. + + Args: + upload_id: Upload session identifier to resume + document: Complete document content as bytes + user: User identifier + on_progress: Optional callback(bytes_sent, total_bytes) for progress updates + + Returns: + dict: Response from completing the upload + + Example: + ```python + library = api.library() + + # Check what's missing + status = library.get_upload_status( + upload_id="abc-123", + user="trustgraph" + ) + + if status['state'] == 'in-progress': + # Resume with the same document + with open("large_document.pdf", "rb") as f: + library.resume_upload( + upload_id="abc-123", + document=f.read(), + user="trustgraph" + ) + ``` + """ + # Get current status + status = self.get_upload_status(upload_id, user) + + if status.get("upload-state") == "expired": + raise RuntimeError("Upload session has expired, please start a new upload") + + if status.get("upload-state") == "completed": + return {"message": "Upload already completed"} + + missing_chunks = status.get("missing-chunks", []) + total_chunks = status.get("total-chunks", 0) + total_bytes = status.get("total-bytes", len(document)) + chunk_size = total_bytes // total_chunks if total_chunks > 0 else DEFAULT_CHUNK_SIZE + + logger.info(f"Resuming upload {upload_id}, {len(missing_chunks)} chunks remaining") + + # Upload missing chunks + for chunk_index in missing_chunks: + start = chunk_index * chunk_size + end = min(start + chunk_size, len(document)) + chunk_data = document[start:end] + + chunk_request = { + "operation": "upload-chunk", + "upload-id": upload_id, + "chunk-index": chunk_index, + "content": base64.b64encode(chunk_data).decode("utf-8"), + "user": user, + } + + self.request(chunk_request) + + if on_progress: + # Estimate progress including previously uploaded chunks + uploaded = total_chunks - len(missing_chunks) + missing_chunks.index(chunk_index) + 1 + bytes_sent = min(uploaded * chunk_size, total_bytes) + on_progress(bytes_sent, total_bytes) + + logger.debug(f"Resumed chunk {chunk_index}") + + # Complete upload + complete_request = { + "operation": "complete-upload", + "upload-id": upload_id, + "user": user, + } + + return self.request(complete_request) + diff --git a/trustgraph-base/trustgraph/schema/services/library.py b/trustgraph-base/trustgraph/schema/services/library.py index 391d49e1..30f8088b 100644 --- a/trustgraph-base/trustgraph/schema/services/library.py +++ b/trustgraph-base/trustgraph/schema/services/library.py @@ -49,6 +49,36 @@ from ..core.metadata import Metadata # <- (processing_metadata[]) # <- (error) +# begin-upload +# -> (document_metadata, total_size, chunk_size) +# <- (upload_id, chunk_size, total_chunks) +# <- (error) + +# upload-chunk +# -> (upload_id, chunk_index, content) +# <- (upload_id, chunk_index, chunks_received, total_chunks, bytes_received, total_bytes) +# <- (error) + +# complete-upload +# -> (upload_id) +# <- (document_id, object_id) +# <- (error) + +# abort-upload +# -> (upload_id) +# <- () +# <- (error) + +# get-upload-status +# -> (upload_id) +# <- (upload_id, state, chunks_received, missing_chunks, total_chunks, bytes_received, total_bytes) +# <- (error) + +# list-uploads +# -> (user) +# <- (uploads[]) +# <- (error) + @dataclass class DocumentMetadata: id: str = "" @@ -76,11 +106,33 @@ class Criteria: value: str = "" operator: str = "" +@dataclass +class UploadProgress: + """Progress information for chunked uploads.""" + upload_id: str = "" + chunks_received: int = 0 + total_chunks: int = 0 + bytes_received: int = 0 + total_bytes: int = 0 + +@dataclass +class UploadSession: + """Information about an in-progress upload.""" + upload_id: str = "" + document_id: str = "" + document_metadata_json: str = "" # JSON-encoded DocumentMetadata + total_size: int = 0 + chunk_size: int = 0 + total_chunks: int = 0 + chunks_received: int = 0 + created_at: str = "" + @dataclass class LibrarianRequest: # add-document, remove-document, update-document, get-document-metadata, # get-document-content, add-processing, remove-processing, list-documents, - # list-processing + # list-processing, begin-upload, upload-chunk, complete-upload, abort-upload, + # get-upload-status, list-uploads operation: str = "" # add-document, remove-document, update-document, get-document-metadata, @@ -90,16 +142,16 @@ class LibrarianRequest: # add-processing, remove-processing processing_id: str = "" - # add-document, update-document + # add-document, update-document, begin-upload document_metadata: DocumentMetadata | None = None # add-processing processing_metadata: ProcessingMetadata | None = None - # add-document + # add-document, upload-chunk content: bytes = b"" - # list-documents, list-processing + # list-documents, list-processing, list-uploads user: str = "" # list-documents?, list-processing? @@ -108,6 +160,16 @@ class LibrarianRequest: # criteria: list[Criteria] = field(default_factory=list) + # begin-upload + total_size: int = 0 + chunk_size: int = 0 + + # upload-chunk, complete-upload, abort-upload, get-upload-status + upload_id: str = "" + + # upload-chunk + chunk_index: int = 0 + @dataclass class LibrarianResponse: error: Error | None = None @@ -116,6 +178,29 @@ class LibrarianResponse: document_metadatas: list[DocumentMetadata] = field(default_factory=list) processing_metadatas: list[ProcessingMetadata] = field(default_factory=list) + # begin-upload response + upload_id: str = "" + chunk_size: int = 0 + total_chunks: int = 0 + + # upload-chunk response + chunk_index: int = 0 + chunks_received: int = 0 + bytes_received: int = 0 + total_bytes: int = 0 + + # complete-upload response + document_id: str = "" + object_id: str = "" + + # get-upload-status response + upload_state: str = "" # "in-progress", "completed", "expired" + received_chunks: list[int] = field(default_factory=list) + missing_chunks: list[int] = field(default_factory=list) + + # list-uploads response + upload_sessions: list[UploadSession] = field(default_factory=list) + # FIXME: Is this right? Using persistence on librarian so that # message chunking works diff --git a/trustgraph-flow/trustgraph/librarian/librarian.py b/trustgraph-flow/trustgraph/librarian/librarian.py index 8835cc73..fba89875 100644 --- a/trustgraph-flow/trustgraph/librarian/librarian.py +++ b/trustgraph-flow/trustgraph/librarian/librarian.py @@ -1,17 +1,24 @@ from .. schema import LibrarianRequest, LibrarianResponse, Error, Triple +from .. schema import UploadSession from .. knowledge import hash from .. exceptions import RequestError from .. tables.library import LibraryTableStore from . blob_store import BlobStore import base64 +import json import logging +import math +import time import uuid # Module logger logger = logging.getLogger(__name__) +# Default chunk size for multipart uploads (5MB - S3 minimum) +DEFAULT_CHUNK_SIZE = 5 * 1024 * 1024 + class Librarian: def __init__( @@ -275,3 +282,314 @@ class Librarian: processing_metadatas = procs, ) + # Chunked upload operations + + async def begin_upload(self, request): + """ + Initialize a chunked upload session. + + Creates an S3 multipart upload and stores session state in Cassandra. + """ + logger.info(f"Beginning chunked upload for document {request.document_metadata.id}") + + if request.document_metadata.kind not in ("text/plain", "application/pdf"): + raise RequestError( + "Invalid document kind: " + request.document_metadata.kind + ) + + if await self.table_store.document_exists( + request.document_metadata.user, + request.document_metadata.id + ): + raise RequestError("Document already exists") + + # Validate sizes + total_size = request.total_size + if total_size <= 0: + raise RequestError("total_size must be positive") + + # Use provided chunk size or default (minimum 5MB for S3) + chunk_size = request.chunk_size if request.chunk_size > 0 else DEFAULT_CHUNK_SIZE + if chunk_size < DEFAULT_CHUNK_SIZE: + chunk_size = DEFAULT_CHUNK_SIZE + + # Calculate total chunks + total_chunks = math.ceil(total_size / chunk_size) + + # Generate IDs + upload_id = str(uuid.uuid4()) + object_id = uuid.uuid4() + + # Create S3 multipart upload + s3_upload_id = self.blob_store.create_multipart_upload( + object_id, request.document_metadata.kind + ) + + # Serialize document metadata for storage + doc_meta_json = json.dumps({ + "id": request.document_metadata.id, + "time": request.document_metadata.time, + "kind": request.document_metadata.kind, + "title": request.document_metadata.title, + "comments": request.document_metadata.comments, + "user": request.document_metadata.user, + "tags": request.document_metadata.tags, + }) + + # Store session in Cassandra + await self.table_store.create_upload_session( + upload_id=upload_id, + user=request.document_metadata.user, + document_id=request.document_metadata.id, + document_metadata=doc_meta_json, + s3_upload_id=s3_upload_id, + object_id=object_id, + total_size=total_size, + chunk_size=chunk_size, + total_chunks=total_chunks, + ) + + logger.info(f"Created upload session {upload_id} with {total_chunks} chunks") + + return LibrarianResponse( + error=None, + upload_id=upload_id, + chunk_size=chunk_size, + total_chunks=total_chunks, + ) + + async def upload_chunk(self, request): + """ + Upload a single chunk of a document. + + Forwards the chunk to S3 and updates session state. + """ + logger.debug(f"Uploading chunk {request.chunk_index} for upload {request.upload_id}") + + # Get session + session = await self.table_store.get_upload_session(request.upload_id) + if session is None: + raise RequestError("Upload session not found or expired") + + # Validate ownership + if session["user"] != request.user: + raise RequestError("Not authorized to upload to this session") + + # Validate chunk index + if request.chunk_index < 0 or request.chunk_index >= session["total_chunks"]: + raise RequestError( + f"Invalid chunk index {request.chunk_index}, " + f"must be 0-{session['total_chunks']-1}" + ) + + # Decode content + content = base64.b64decode(request.content) + + # Upload to S3 (part numbers are 1-indexed in S3) + part_number = request.chunk_index + 1 + etag = self.blob_store.upload_part( + object_id=session["object_id"], + upload_id=session["s3_upload_id"], + part_number=part_number, + data=content, + ) + + # Update session with chunk info + await self.table_store.update_upload_session_chunk( + upload_id=request.upload_id, + chunk_index=request.chunk_index, + etag=etag, + ) + + # Calculate progress + chunks_received = session["chunks_received"] + # Add this chunk if not already present + if request.chunk_index not in chunks_received: + chunks_received[request.chunk_index] = etag + + num_chunks_received = len(chunks_received) + 1 # +1 for this chunk + bytes_received = num_chunks_received * session["chunk_size"] + # Adjust for last chunk potentially being smaller + if bytes_received > session["total_size"]: + bytes_received = session["total_size"] + + logger.debug(f"Chunk {request.chunk_index} uploaded, {num_chunks_received}/{session['total_chunks']} complete") + + return LibrarianResponse( + error=None, + upload_id=request.upload_id, + chunk_index=request.chunk_index, + chunks_received=num_chunks_received, + total_chunks=session["total_chunks"], + bytes_received=bytes_received, + total_bytes=session["total_size"], + ) + + async def complete_upload(self, request): + """ + Finalize a chunked upload and create the document. + + Completes the S3 multipart upload and creates the document metadata. + """ + logger.info(f"Completing upload {request.upload_id}") + + # Get session + session = await self.table_store.get_upload_session(request.upload_id) + if session is None: + raise RequestError("Upload session not found or expired") + + # Validate ownership + if session["user"] != request.user: + raise RequestError("Not authorized to complete this upload") + + # Verify all chunks received + chunks_received = session["chunks_received"] + if len(chunks_received) != session["total_chunks"]: + missing = [ + i for i in range(session["total_chunks"]) + if i not in chunks_received + ] + raise RequestError( + f"Missing chunks: {missing[:10]}{'...' if len(missing) > 10 else ''}" + ) + + # Build parts list for S3 (sorted by part number) + parts = [ + (chunk_index + 1, etag) # S3 part numbers are 1-indexed + for chunk_index, etag in sorted(chunks_received.items()) + ] + + # Complete S3 multipart upload + self.blob_store.complete_multipart_upload( + object_id=session["object_id"], + upload_id=session["s3_upload_id"], + parts=parts, + ) + + # Parse document metadata from session + doc_meta_dict = json.loads(session["document_metadata"]) + + # Create DocumentMetadata object + from .. schema import DocumentMetadata + doc_metadata = DocumentMetadata( + id=doc_meta_dict["id"], + time=doc_meta_dict.get("time", int(time.time())), + kind=doc_meta_dict["kind"], + title=doc_meta_dict.get("title", ""), + comments=doc_meta_dict.get("comments", ""), + user=doc_meta_dict["user"], + tags=doc_meta_dict.get("tags", []), + metadata=[], # Triples not supported in chunked upload yet + ) + + # Add document to table + await self.table_store.add_document(doc_metadata, session["object_id"]) + + # Delete upload session + await self.table_store.delete_upload_session(request.upload_id) + + logger.info(f"Upload {request.upload_id} completed, document {doc_metadata.id} created") + + return LibrarianResponse( + error=None, + document_id=doc_metadata.id, + object_id=str(session["object_id"]), + ) + + async def abort_upload(self, request): + """ + Cancel a chunked upload and clean up resources. + """ + logger.info(f"Aborting upload {request.upload_id}") + + # Get session + session = await self.table_store.get_upload_session(request.upload_id) + if session is None: + raise RequestError("Upload session not found or expired") + + # Validate ownership + if session["user"] != request.user: + raise RequestError("Not authorized to abort this upload") + + # Abort S3 multipart upload + self.blob_store.abort_multipart_upload( + object_id=session["object_id"], + upload_id=session["s3_upload_id"], + ) + + # Delete session from Cassandra + await self.table_store.delete_upload_session(request.upload_id) + + logger.info(f"Upload {request.upload_id} aborted") + + return LibrarianResponse(error=None) + + async def get_upload_status(self, request): + """ + Get the status of an in-progress upload. + """ + logger.debug(f"Getting status for upload {request.upload_id}") + + # Get session + session = await self.table_store.get_upload_session(request.upload_id) + if session is None: + return LibrarianResponse( + error=None, + upload_id=request.upload_id, + upload_state="expired", + ) + + # Validate ownership + if session["user"] != request.user: + raise RequestError("Not authorized to view this upload") + + chunks_received = session["chunks_received"] + received_list = sorted(chunks_received.keys()) + missing_list = [ + i for i in range(session["total_chunks"]) + if i not in chunks_received + ] + + bytes_received = len(chunks_received) * session["chunk_size"] + if bytes_received > session["total_size"]: + bytes_received = session["total_size"] + + return LibrarianResponse( + error=None, + upload_id=request.upload_id, + upload_state="in-progress", + received_chunks=received_list, + missing_chunks=missing_list, + chunks_received=len(chunks_received), + total_chunks=session["total_chunks"], + bytes_received=bytes_received, + total_bytes=session["total_size"], + ) + + async def list_uploads(self, request): + """ + List all in-progress uploads for a user. + """ + logger.debug(f"Listing uploads for user {request.user}") + + sessions = await self.table_store.list_upload_sessions(request.user) + + upload_sessions = [ + UploadSession( + upload_id=s["upload_id"], + document_id=s["document_id"], + document_metadata_json=s.get("document_metadata", ""), + total_size=s["total_size"], + chunk_size=s["chunk_size"], + total_chunks=s["total_chunks"], + chunks_received=s["chunks_received"], + created_at=str(s.get("created_at", "")), + ) + for s in sessions + ] + + return LibrarianResponse( + error=None, + upload_sessions=upload_sessions, + ) + diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index 7c1e428c..9ef449e8 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -361,6 +361,13 @@ class Processor(AsyncProcessor): "remove-processing": self.librarian.remove_processing, "list-documents": self.librarian.list_documents, "list-processing": self.librarian.list_processing, + # Chunked upload operations + "begin-upload": self.librarian.begin_upload, + "upload-chunk": self.librarian.upload_chunk, + "complete-upload": self.librarian.complete_upload, + "abort-upload": self.librarian.abort_upload, + "get-upload-status": self.librarian.get_upload_status, + "list-uploads": self.librarian.list_uploads, } if v.operation not in impls: