diff --git a/trustgraph-base/trustgraph/api/library.py b/trustgraph-base/trustgraph/api/library.py index 5e6a6c64..b61348fc 100644 --- a/trustgraph-base/trustgraph/api/library.py +++ b/trustgraph-base/trustgraph/api/library.py @@ -303,14 +303,16 @@ class Library: logger.warning(f"Failed to abort upload: {abort_error}") raise - def get_documents(self, user): + def get_documents(self, user, include_children=False): """ List all documents for a user. Retrieves metadata for all documents owned by the specified user. + By default, only returns top-level documents (not child/extracted documents). Args: user: User identifier + include_children: If True, also include child documents (default: False) Returns: list[DocumentMetadata]: List of document metadata objects @@ -321,18 +323,24 @@ class Library: Example: ```python library = api.library() + + # Get only top-level documents docs = library.get_documents(user="trustgraph") for doc in docs: print(f"{doc.id}: {doc.title} ({doc.kind})") print(f" Uploaded: {doc.time}") print(f" Tags: {', '.join(doc.tags)}") + + # Get all documents including extracted pages + all_docs = library.get_documents(user="trustgraph", include_children=True) ``` """ input = { "operation": "list-documents", "user": user, + "include-children": include_children, } object = self.request(input) @@ -354,7 +362,9 @@ class Library: for w in v["metadata"] ], user = v["user"], - tags = v["tags"] + tags = v["tags"], + parent_id = v.get("parent-id", ""), + document_type = v.get("document-type", "source"), ) for v in object["document-metadatas"] ] @@ -397,7 +407,7 @@ class Library: doc = object["document-metadata"] try: - DocumentMetadata( + return DocumentMetadata( id = doc["id"], time = datetime.datetime.fromtimestamp(doc["time"]), kind = doc["kind"], @@ -412,7 +422,9 @@ class Library: for w in doc["metadata"] ], user = doc["user"], - tags = doc["tags"] + tags = doc["tags"], + parent_id = doc.get("parent-id", ""), + document_type = doc.get("document-type", "source"), ) except Exception as e: logger.error("Failed to parse document response", exc_info=True) @@ -860,3 +872,258 @@ class Library: return self.request(complete_request) + # Child document methods + + def add_child_document( + self, document, id, parent_id, user, title, comments, + kind="text/plain", tags=[], metadata=None, + ): + """ + Add a child document linked to a parent document. + + Child documents are typically extracted content (e.g., pages from a PDF). + They are automatically marked with document_type="extracted" and linked + to their parent via parent_id. + + Args: + document: Document content as bytes + id: Document identifier (auto-generated if None) + parent_id: Parent document identifier (required) + user: User/owner identifier + title: Document title + comments: Document description or comments + kind: MIME type of the document (default: "text/plain") + tags: List of tags for categorization (default: []) + metadata: Optional metadata as list of Triple objects + + Returns: + dict: Response from the add operation + + Raises: + RuntimeError: If parent_id is not provided + + Example: + ```python + library = api.library() + + # Add extracted page from a PDF + library.add_child_document( + document=page_text.encode('utf-8'), + id="doc-123-page-1", + parent_id="doc-123", + user="trustgraph", + title="Page 1 of Research Paper", + comments="First page extracted from PDF", + kind="text/plain", + tags=["extracted", "page"] + ) + ``` + """ + if not parent_id: + raise RuntimeError("parent_id is required for child documents") + + if id is None: + id = hash(document) + + if not title: + title = "" + if not comments: + comments = "" + + triples = [] + if metadata: + if isinstance(metadata, list): + triples = [ + { + "s": from_value(t.s), + "p": from_value(t.p), + "o": from_value(t.o), + } + for t in metadata + ] + + input = { + "operation": "add-child-document", + "document-metadata": { + "id": id, + "time": int(time.time()), + "kind": kind, + "title": title, + "comments": comments, + "metadata": triples, + "user": user, + "tags": tags, + "parent-id": parent_id, + "document-type": "extracted", + }, + "content": base64.b64encode(document).decode("utf-8"), + } + + return self.request(input) + + def list_children(self, document_id, user): + """ + List all child documents for a given parent document. + + Args: + document_id: Parent document identifier + user: User identifier + + Returns: + list[DocumentMetadata]: List of child document metadata objects + + Example: + ```python + library = api.library() + children = library.list_children( + document_id="doc-123", + user="trustgraph" + ) + + for child in children: + print(f"{child.id}: {child.title}") + ``` + """ + input = { + "operation": "list-children", + "document-id": document_id, + "user": user, + } + + response = self.request(input) + + try: + return [ + DocumentMetadata( + id=v["id"], + time=datetime.datetime.fromtimestamp(v["time"]), + kind=v["kind"], + title=v["title"], + comments=v.get("comments", ""), + metadata=[ + Triple( + s=to_value(w["s"]), + p=to_value(w["p"]), + o=to_value(w["o"]) + ) + for w in v.get("metadata", []) + ], + user=v["user"], + tags=v.get("tags", []), + parent_id=v.get("parent-id", ""), + document_type=v.get("document-type", "source"), + ) + for v in response.get("document-metadatas", []) + ] + except Exception as e: + logger.error("Failed to parse children response", exc_info=True) + raise ProtocolException("Response not formatted correctly") + + def get_document_content(self, user, id): + """ + Get the content of a document. + + Retrieves the full content of a document as bytes. + + Args: + user: User identifier + id: Document identifier + + Returns: + bytes: Document content + + Example: + ```python + library = api.library() + content = library.get_document_content( + user="trustgraph", + id="doc-123" + ) + + # Write to file + with open("output.pdf", "wb") as f: + f.write(content) + ``` + """ + input = { + "operation": "get-document-content", + "user": user, + "document-id": id, + } + + response = self.request(input) + content_b64 = response.get("content", "") + + return base64.b64decode(content_b64) + + def stream_document_to_file(self, user, id, file_path, chunk_size=1024*1024, on_progress=None): + """ + Stream document content to a file. + + Downloads document content in chunks and writes directly to a file, + enabling memory-efficient handling of large documents. + + Args: + user: User identifier + id: Document identifier + file_path: Path to write the document content + chunk_size: Size of each chunk to download (default 1MB) + on_progress: Optional callback(bytes_received, total_bytes) for progress updates + + Returns: + int: Total bytes written + + Example: + ```python + library = api.library() + + def progress(received, total): + print(f"Downloaded {received}/{total} bytes") + + library.stream_document_to_file( + user="trustgraph", + id="large-doc-123", + file_path="/tmp/document.pdf", + on_progress=progress + ) + ``` + """ + chunk_index = 0 + total_bytes_written = 0 + total_bytes = None + + with open(file_path, "wb") as f: + while True: + input = { + "operation": "stream-document", + "user": user, + "document-id": id, + "chunk-index": chunk_index, + "chunk-size": chunk_size, + } + + response = self.request(input) + + content_b64 = response.get("content", "") + chunk_data = base64.b64decode(content_b64) + + if not chunk_data: + break + + f.write(chunk_data) + total_bytes_written += len(chunk_data) + + total_chunks = response.get("total-chunks", 1) + total_bytes = response.get("total-bytes", total_bytes_written) + + if on_progress: + on_progress(total_bytes_written, total_bytes) + + # Check if we've received all chunks + if chunk_index >= total_chunks - 1: + break + + chunk_index += 1 + + return total_bytes_written + diff --git a/trustgraph-base/trustgraph/api/types.py b/trustgraph-base/trustgraph/api/types.py index 3b4e476e..47aa5ae0 100644 --- a/trustgraph-base/trustgraph/api/types.py +++ b/trustgraph-base/trustgraph/api/types.py @@ -64,6 +64,8 @@ class DocumentMetadata: metadata: List of RDF triples providing structured metadata user: User/owner identifier tags: List of tags for categorization + parent_id: Parent document ID for child documents (empty for top-level docs) + document_type: "source" for uploaded documents, "extracted" for derived content """ id : str time : datetime.datetime @@ -73,6 +75,8 @@ class DocumentMetadata: metadata : List[Triple] user : str tags : List[str] + parent_id : str = "" + document_type : str = "source" @dataclasses.dataclass class ProcessingMetadata: diff --git a/trustgraph-base/trustgraph/schema/knowledge/document.py b/trustgraph-base/trustgraph/schema/knowledge/document.py index d8ce97b4..f3996136 100644 --- a/trustgraph-base/trustgraph/schema/knowledge/document.py +++ b/trustgraph-base/trustgraph/schema/knowledge/document.py @@ -10,6 +10,9 @@ from ..core.topic import topic class Document: metadata: Metadata | None = None data: bytes = b"" + # For large document streaming: if document_id is set, the receiver should + # fetch content from librarian instead of using inline data + document_id: str = "" ############################################################################ diff --git a/trustgraph-base/trustgraph/schema/services/library.py b/trustgraph-base/trustgraph/schema/services/library.py index 30f8088b..4025977a 100644 --- a/trustgraph-base/trustgraph/schema/services/library.py +++ b/trustgraph-base/trustgraph/schema/services/library.py @@ -89,6 +89,9 @@ class DocumentMetadata: metadata: list[Triple] = field(default_factory=list) user: str = "" tags: list[str] = field(default_factory=list) + # Child document support + parent_id: str = "" # Empty for top-level docs, set for children + document_type: str = "source" # "source" or "extracted" @dataclass class ProcessingMetadata: @@ -167,9 +170,12 @@ class LibrarianRequest: # upload-chunk, complete-upload, abort-upload, get-upload-status upload_id: str = "" - # upload-chunk + # upload-chunk, stream-document chunk_index: int = 0 + # list-documents - whether to include child documents (default False) + include_children: bool = False + @dataclass class LibrarianResponse: error: Error | None = None diff --git a/trustgraph-cli/trustgraph/cli/load_pdf.py b/trustgraph-cli/trustgraph/cli/load_pdf.py index d305cb4b..a770ec05 100644 --- a/trustgraph-cli/trustgraph/cli/load_pdf.py +++ b/trustgraph-cli/trustgraph/cli/load_pdf.py @@ -1,15 +1,30 @@ """ -Loads a PDF document into TrustGraph processing by directing to -the pdf-decoder queue. -Consider using tg-add-library-document to load -a document, followed by tg-start-library-processing to initiate processing. +DEPRECATED: This tool is deprecated and will be removed in a future version. + +Use tg-add-library-document to add a document to the library, followed by +tg-start-library-processing to initiate processing. + +Example: + # Add document to library + tg-add-library-document --id my-doc --title "My PDF" document.pdf + + # Start processing + tg-start-library-processing --document-id my-doc --collection default + +The library-based workflow provides: +- Progress feedback for large file uploads +- Resumable uploads +- Better error handling +- Centralized document management """ import hashlib import argparse import os +import sys import time import uuid +import warnings from trustgraph.api import Api from trustgraph.knowledge import hash, to_uri @@ -71,6 +86,18 @@ class Loader: def main(): + # Print deprecation warning + print("\n" + "=" * 70, file=sys.stderr) + print("DEPRECATION WARNING: tg-load-pdf is deprecated.", file=sys.stderr) + print("", file=sys.stderr) + print("Please use the library-based workflow instead:", file=sys.stderr) + print(" 1. tg-add-library-document to add the document", file=sys.stderr) + print(" 2. tg-start-library-processing to start processing", file=sys.stderr) + print("", file=sys.stderr) + print("This provides progress feedback, resumable uploads, and", file=sys.stderr) + print("better handling of large documents.", file=sys.stderr) + print("=" * 70 + "\n", file=sys.stderr) + parser = argparse.ArgumentParser( prog='tg-load-pdf', description=__doc__, diff --git a/trustgraph-cli/trustgraph/cli/load_text.py b/trustgraph-cli/trustgraph/cli/load_text.py index 594d1c04..3533a7ae 100644 --- a/trustgraph-cli/trustgraph/cli/load_text.py +++ b/trustgraph-cli/trustgraph/cli/load_text.py @@ -1,8 +1,21 @@ """ -Loads a text document into TrustGraph processing by directing to a text -loader queue. -Consider using tg-add-library-document to load -a document, followed by tg-start-library-processing to initiate processing. +DEPRECATED: This tool is deprecated and will be removed in a future version. + +Use tg-add-library-document to add a document to the library, followed by +tg-start-library-processing to initiate processing. + +Example: + # Add document to library + tg-add-library-document --id my-doc --title "My Text" --kind text/plain document.txt + + # Start processing + tg-start-library-processing --document-id my-doc --collection default + +The library-based workflow provides: +- Progress feedback for large file uploads +- Resumable uploads +- Better error handling +- Centralized document management """ import pulsar @@ -10,8 +23,10 @@ from pulsar.schema import JsonSchema import hashlib import argparse import os +import sys import time import uuid +import warnings from trustgraph.api import Api from trustgraph.knowledge import hash, to_uri @@ -73,6 +88,18 @@ class Loader: def main(): + # Print deprecation warning + print("\n" + "=" * 70, file=sys.stderr) + print("DEPRECATION WARNING: tg-load-text is deprecated.", file=sys.stderr) + print("", file=sys.stderr) + print("Please use the library-based workflow instead:", file=sys.stderr) + print(" 1. tg-add-library-document to add the document", file=sys.stderr) + print(" 2. tg-start-library-processing to start processing", file=sys.stderr) + print("", file=sys.stderr) + print("This provides progress feedback, resumable uploads, and", file=sys.stderr) + print("better handling of large documents.", file=sys.stderr) + print("=" * 70 + "\n", file=sys.stderr) + parser = argparse.ArgumentParser( prog='tg-load-text', description=__doc__, diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index bb641a26..d60e9b77 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -2,8 +2,12 @@ """ Simple decoder, accepts PDF documents on input, outputs pages from the PDF document as text as separate output objects. + +Supports both inline document data and streaming from librarian API +for large documents. """ +import os import tempfile import base64 import logging @@ -17,6 +21,10 @@ logger = logging.getLogger(__name__) default_ident = "pdf-decoder" +# Default API URL for fetching documents from librarian +default_api_url = os.environ.get("TRUSTGRAPH_API_URL", "http://api:8088") + + class Processor(FlowProcessor): def __init__(self, **params): @@ -29,6 +37,8 @@ class Processor(FlowProcessor): } ) + self.api_url = params.get("api_url", default_api_url) + self.register_specification( ConsumerSpec( name = "input", @@ -46,6 +56,59 @@ class Processor(FlowProcessor): logger.info("PDF decoder initialized") + def _fetch_document_to_file(self, document_id, user, file_path, chunk_size=1024*1024): + """ + Fetch document content from librarian API and stream to file. + + This avoids loading the entire document into memory at once. + """ + import requests + + logger.info(f"Streaming document {document_id} to temp file...") + + # Use chunk-based streaming to minimize memory usage + chunk_index = 0 + total_bytes = 0 + + with open(file_path, 'wb') as f: + while True: + url = f"{self.api_url}/api/v1/librarian" + payload = { + "operation": "stream-document", + "user": user, + "document-id": document_id, + "chunk-index": chunk_index, + "chunk-size": chunk_size, + } + + try: + response = requests.post(url, json=payload, timeout=60) + response.raise_for_status() + data = response.json() + except requests.RequestException as e: + logger.error(f"Failed to fetch chunk {chunk_index}: {e}") + raise + + if "error" in data and data["error"]: + raise RuntimeError(f"API error: {data['error']}") + + content_b64 = data.get("content", "") + if not content_b64: + break + + chunk_data = base64.b64decode(content_b64) + f.write(chunk_data) + total_bytes += len(chunk_data) + + total_chunks = data.get("total-chunks", 1) + if chunk_index >= total_chunks - 1: + break + + chunk_index += 1 + + logger.info(f"Downloaded {total_bytes} bytes to temp file") + return total_bytes + async def on_message(self, msg, consumer, flow): logger.debug("PDF message received") @@ -54,26 +117,43 @@ class Processor(FlowProcessor): logger.info(f"Decoding PDF {v.metadata.id}...") - with tempfile.NamedTemporaryFile(delete_on_close=False) as fp: + with tempfile.NamedTemporaryFile(delete_on_close=False, suffix='.pdf') as fp: + temp_path = fp.name - fp.write(base64.b64decode(v.data)) - fp.close() + # Check if we should fetch from librarian or use inline data + if v.document_id: + # Stream from librarian API to temp file + logger.info(f"Fetching document {v.document_id} from librarian...") + fp.close() + self._fetch_document_to_file( + document_id=v.document_id, + user=v.metadata.user, + file_path=temp_path, + ) + else: + # Use inline data (backward compatibility) + fp.write(base64.b64decode(v.data)) + fp.close() - with open(fp.name, mode='rb') as f: + loader = PyPDFLoader(temp_path) + pages = loader.load() - loader = PyPDFLoader(fp.name) - pages = loader.load() + for ix, page in enumerate(pages): - for ix, page in enumerate(pages): + logger.debug(f"Processing page {ix}") - logger.debug(f"Processing page {ix}") + r = TextDocument( + metadata=v.metadata, + text=page.page_content.encode("utf-8"), + ) - r = TextDocument( - metadata=v.metadata, - text=page.page_content.encode("utf-8"), - ) + await flow("output").send(r) - await flow("output").send(r) + # Clean up temp file + try: + os.unlink(temp_path) + except OSError: + pass logger.debug("PDF decoding complete") @@ -81,6 +161,12 @@ class Processor(FlowProcessor): def add_args(parser): FlowProcessor.add_args(parser) + parser.add_argument( + '--api-url', + default=default_api_url, + help=f'TrustGraph API URL for document streaming (default: {default_api_url})', + ) + def run(): Processor.launch(default_ident, __doc__) diff --git a/trustgraph-flow/trustgraph/librarian/librarian.py b/trustgraph-flow/trustgraph/librarian/librarian.py index fba89875..831e015a 100644 --- a/trustgraph-flow/trustgraph/librarian/librarian.py +++ b/trustgraph-flow/trustgraph/librarian/librarian.py @@ -91,6 +91,21 @@ class Librarian: ): raise RuntimeError("Document does not exist") + # First, cascade delete all child documents + children = await self.table_store.list_children(request.document_id) + for child in children: + logger.debug(f"Cascade deleting child document {child.id}") + try: + child_object_id = await self.table_store.get_document_object_id( + child.user, + child.id + ) + await self.blob_store.remove(child_object_id) + await self.table_store.remove_document(child.user, child.id) + except Exception as e: + logger.warning(f"Failed to delete child document {child.id}: {e}") + + # Now remove the parent document object_id = await self.table_store.get_document_object_id( request.user, request.document_id @@ -262,6 +277,14 @@ class Librarian: docs = await self.table_store.list_documents(request.user) + # Filter out child documents by default unless include_children is True + include_children = getattr(request, 'include_children', False) + if not include_children: + docs = [ + doc for doc in docs + if not doc.parent_id # Only include top-level documents + ] + return LibrarianResponse( error = None, document_metadata = None, @@ -593,3 +616,124 @@ class Librarian: upload_sessions=upload_sessions, ) + # Child document operations + + async def add_child_document(self, request): + """ + Add a child document linked to a parent document. + + Child documents are typically extracted content (e.g., pages from a PDF). + They have a parent_id pointing to the source document and document_type + set to "extracted". + """ + logger.info(f"Adding child document {request.document_metadata.id} " + f"for parent {request.document_metadata.parent_id}") + + if not request.document_metadata.parent_id: + raise RequestError("parent_id is required for child documents") + + # Verify parent exists + if not await self.table_store.document_exists( + request.document_metadata.user, + request.document_metadata.parent_id + ): + raise RequestError( + f"Parent document {request.document_metadata.parent_id} does not exist" + ) + + if await self.table_store.document_exists( + request.document_metadata.user, + request.document_metadata.id + ): + raise RequestError("Document already exists") + + # Ensure document_type is set to "extracted" + request.document_metadata.document_type = "extracted" + + # Create object ID for blob + object_id = uuid.uuid4() + + logger.debug("Adding blob...") + + await self.blob_store.add( + object_id, base64.b64decode(request.content), + request.document_metadata.kind + ) + + logger.debug("Adding to table...") + + await self.table_store.add_document( + request.document_metadata, object_id + ) + + logger.debug("Add child document complete") + + return LibrarianResponse( + error=None, + document_id=request.document_metadata.id, + ) + + async def list_children(self, request): + """ + List all child documents for a given parent document. + """ + logger.debug(f"Listing children for parent {request.document_id}") + + children = await self.table_store.list_children(request.document_id) + + return LibrarianResponse( + error=None, + document_metadatas=children, + ) + + async def stream_document(self, request): + """ + Stream document content in chunks. + + This operation returns document content in smaller chunks, allowing + memory-efficient processing of large documents. The response includes + chunk information for reassembly. + + Note: This operation returns a single chunk at a time. Clients should + call repeatedly with increasing chunk_index until all chunks are received. + """ + logger.debug(f"Streaming document {request.document_id}, chunk {request.chunk_index}") + + object_id = await self.table_store.get_document_object_id( + request.user, + request.document_id + ) + + # Default chunk size of 1MB + chunk_size = request.chunk_size if request.chunk_size > 0 else 1024 * 1024 + + # Get the full content and slice out the requested chunk + # Note: This is a simple implementation. For true streaming, we'd need + # range requests on the object storage. + content = await self.blob_store.get(object_id) + total_size = len(content) + total_chunks = math.ceil(total_size / chunk_size) + + if request.chunk_index >= total_chunks: + raise RequestError( + f"Invalid chunk index {request.chunk_index}, " + f"document has {total_chunks} chunks" + ) + + start = request.chunk_index * chunk_size + end = min(start + chunk_size, total_size) + chunk_content = content[start:end] + + logger.debug(f"Returning chunk {request.chunk_index}/{total_chunks}, " + f"bytes {start}-{end} of {total_size}") + + return LibrarianResponse( + error=None, + content=base64.b64encode(chunk_content), + chunk_index=request.chunk_index, + chunks_received=1, # Using as "current chunk" indicator + total_chunks=total_chunks, + bytes_received=end, + total_bytes=total_size, + ) + diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index 9ef449e8..217da4f1 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -271,6 +271,9 @@ class Processor(AsyncProcessor): pass + # Threshold for sending document_id instead of inline content (2MB) + STREAMING_THRESHOLD = 2 * 1024 * 1024 + async def load_document(self, document, processing, content): logger.debug("Ready for document processing...") @@ -303,15 +306,31 @@ class Processor(AsyncProcessor): ) schema = TextDocument else: - doc = Document( - metadata = Metadata( - id = document.id, - metadata = document.metadata, - user = processing.user, - collection = processing.collection - ), - data = base64.b64encode(content).decode("utf-8") - ) + # For large PDF documents, send document_id for streaming retrieval + # instead of embedding the entire content in the message + if len(content) >= self.STREAMING_THRESHOLD: + logger.info(f"Document {document.id} is large ({len(content)} bytes), " + f"sending document_id for streaming retrieval") + doc = Document( + metadata = Metadata( + id = document.id, + metadata = document.metadata, + user = processing.user, + collection = processing.collection + ), + document_id = document.id, + data = b"", # Empty data, receiver will fetch via API + ) + else: + doc = Document( + metadata = Metadata( + id = document.id, + metadata = document.metadata, + user = processing.user, + collection = processing.collection + ), + data = base64.b64encode(content).decode("utf-8") + ) schema = Document logger.debug(f"Submitting to queue {q}...") @@ -368,6 +387,10 @@ class Processor(AsyncProcessor): "abort-upload": self.librarian.abort_upload, "get-upload-status": self.librarian.get_upload_status, "list-uploads": self.librarian.list_uploads, + # Child document and streaming operations + "add-child-document": self.librarian.add_child_document, + "list-children": self.librarian.list_children, + "stream-document": self.librarian.stream_document, } if v.operation not in impls: diff --git a/trustgraph-flow/trustgraph/tables/library.py b/trustgraph-flow/trustgraph/tables/library.py index aa73c950..354a66e3 100644 --- a/trustgraph-flow/trustgraph/tables/library.py +++ b/trustgraph-flow/trustgraph/tables/library.py @@ -112,6 +112,34 @@ class LibraryTableStore: ON document (object_id) """); + # Add parent_id and document_type columns for child document support + logger.debug("document table parent_id column...") + + try: + self.cassandra.execute(""" + ALTER TABLE document ADD parent_id text + """); + except Exception as e: + # Column may already exist + if "already exists" not in str(e).lower() and "Invalid column name" not in str(e): + logger.debug(f"parent_id column may already exist: {e}") + + try: + self.cassandra.execute(""" + ALTER TABLE document ADD document_type text + """); + except Exception as e: + # Column may already exist + if "already exists" not in str(e).lower() and "Invalid column name" not in str(e): + logger.debug(f"document_type column may already exist: {e}") + + logger.debug("document parent index...") + + self.cassandra.execute(""" + CREATE INDEX IF NOT EXISTS document_parent + ON document (parent_id) + """); + logger.debug("processing table...") self.cassandra.execute(""" @@ -162,9 +190,10 @@ class LibraryTableStore: ( id, user, time, kind, title, comments, - metadata, tags, object_id + metadata, tags, object_id, + parent_id, document_type ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """) self.update_document_stmt = self.cassandra.prepare(""" @@ -175,7 +204,8 @@ class LibraryTableStore: """) self.get_document_stmt = self.cassandra.prepare(""" - SELECT time, kind, title, comments, metadata, tags, object_id + SELECT time, kind, title, comments, metadata, tags, object_id, + parent_id, document_type FROM document WHERE user = ? AND id = ? """) @@ -194,14 +224,16 @@ class LibraryTableStore: self.list_document_stmt = self.cassandra.prepare(""" SELECT - id, time, kind, title, comments, metadata, tags, object_id + id, time, kind, title, comments, metadata, tags, object_id, + parent_id, document_type FROM document WHERE user = ? """) self.list_document_by_tag_stmt = self.cassandra.prepare(""" SELECT - id, time, kind, title, comments, metadata, tags, object_id + id, time, kind, title, comments, metadata, tags, object_id, + parent_id, document_type FROM document WHERE user = ? AND tags CONTAINS ? ALLOW FILTERING @@ -277,6 +309,16 @@ class LibraryTableStore: WHERE user = ? """) + # Child document queries + self.list_children_stmt = self.cassandra.prepare(""" + SELECT + id, user, time, kind, title, comments, metadata, tags, + object_id, parent_id, document_type + FROM document + WHERE parent_id = ? + ALLOW FILTERING + """) + async def document_exists(self, user, id): resp = self.cassandra.execute( @@ -303,6 +345,10 @@ class LibraryTableStore: for v in document.metadata ] + # Get parent_id and document_type from document, defaulting if not set + parent_id = getattr(document, 'parent_id', '') or '' + document_type = getattr(document, 'document_type', 'source') or 'source' + while True: try: @@ -312,7 +358,8 @@ class LibraryTableStore: ( document.id, document.user, int(document.time * 1000), document.kind, document.title, document.comments, - metadata, document.tags, object_id + metadata, document.tags, object_id, + parent_id, document_type ) ) @@ -419,6 +466,55 @@ class LibraryTableStore: for m in row[5] ], tags = row[6] if row[6] else [], + parent_id = row[8] if row[8] else "", + document_type = row[9] if row[9] else "source", + ) + for row in resp + ] + + logger.debug("Done") + + return lst + + async def list_children(self, parent_id): + """List all child documents for a given parent document ID.""" + + logger.debug(f"List children for parent {parent_id}") + + while True: + + try: + + resp = self.cassandra.execute( + self.list_children_stmt, + (parent_id,) + ) + + break + + except Exception as e: + logger.error("Exception occurred", exc_info=True) + raise e + + lst = [ + DocumentMetadata( + id = row[0], + user = row[1], + time = int(time.mktime(row[2].timetuple())), + kind = row[3], + title = row[4], + comments = row[5], + metadata = [ + Triple( + s=tuple_to_term(m[0], m[1]), + p=tuple_to_term(m[2], m[3]), + o=tuple_to_term(m[4], m[5]) + ) + for m in row[6] + ], + tags = row[7] if row[7] else [], + parent_id = row[9] if row[9] else "", + document_type = row[10] if row[10] else "source", ) for row in resp ] @@ -464,6 +560,8 @@ class LibraryTableStore: for m in row[4] ], tags = row[5] if row[5] else [], + parent_id = row[7] if row[7] else "", + document_type = row[8] if row[8] else "source", ) logger.debug("Done")