From 3a1a30944865eabca15d6d99444151f439ad105f Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 21 Dec 2024 19:35:56 -0600 Subject: [PATCH] document service gets id before write --- src/basic_memory/models/documents.py | 2 +- src/basic_memory/services/document_service.py | 55 +++++++++++------ tests/services/test_document_service.py | 61 +++++++++++++++++-- 3 files changed, 95 insertions(+), 23 deletions(-) diff --git a/src/basic_memory/models/documents.py b/src/basic_memory/models/documents.py index 11ffc6b3..cc5f18ac 100644 --- a/src/basic_memory/models/documents.py +++ b/src/basic_memory/models/documents.py @@ -22,7 +22,7 @@ class Document(Base): id: Mapped[int] = mapped_column(primary_key=True) path: Mapped[str] = mapped_column(String, unique=True, nullable=False) - checksum: Mapped[str] = mapped_column(String, nullable=False) + checksum: Mapped[str] = mapped_column(String, nullable=True) doc_metadata: Mapped[Optional[dict]] = mapped_column( JSON, nullable=True ) # renamed from metadata diff --git a/src/basic_memory/services/document_service.py b/src/basic_memory/services/document_service.py index 52fe6b96..ff3a58bc 100644 --- a/src/basic_memory/services/document_service.py +++ b/src/basic_memory/services/document_service.py @@ -1,9 +1,11 @@ """Service for managing documents in the system.""" import hashlib +from datetime import datetime, UTC from pathlib import Path from typing import Optional, Dict, Any, List +import yaml from loguru import logger from basic_memory.models import Document @@ -13,19 +15,16 @@ from basic_memory.services.service import BaseService class DocumentError(Exception): """Base exception for document operations.""" - pass class DocumentNotFoundError(DocumentError): """Raised when a document doesn't exist.""" - pass class DocumentWriteError(DocumentError): """Raised when document file operations fail.""" - pass @@ -66,6 +65,21 @@ class DocumentService(BaseService[DocumentRepository]): except Exception as e: raise DocumentWriteError(f"Directory not writable: {parent}: {e}") + async def add_frontmatter(self, content: str, doc_id: int, metadata: Optional[Dict[str, Any]] = None) -> str: + """Add frontmatter to document content.""" + # Generate frontmatter with timestamps + now = datetime.now(UTC).isoformat() + frontmatter = { + "id": doc_id, + "created": now, + "modified": now + } + if metadata: + frontmatter.update(metadata) + + yaml_fm = yaml.dump(frontmatter, sort_keys=False) + return f"---\n{yaml_fm}---\n\n{content}" + async def list_documents(self) -> List[Document]: """List all documents in the database.""" return await self.repository.find_all() @@ -93,23 +107,30 @@ class DocumentService(BaseService[DocumentRepository]): file_path = Path(path) await self.ensure_parent_directory(file_path) - # Write file first try: - file_path.write_text(content) - except Exception as e: - raise DocumentWriteError(f"Failed to write document file: {e}") + # 1. Create initial DB record to get ID + doc = await self.repository.create({ + "path": str(path), + "doc_metadata": metadata + }) - # After file is written, create database record - try: - checksum = await self.compute_checksum(content) - doc = await self.repository.create( - {"path": str(path), "checksum": checksum, "doc_metadata": metadata} - ) + # 2. Add frontmatter with DB-generated ID + content_with_frontmatter = await self.add_frontmatter(content, doc.id, metadata) + + # 3. Write complete file + file_path.write_text(content_with_frontmatter) + + # 4. Update DB with checksum to mark completion + checksum = await self.compute_checksum(content_with_frontmatter) + doc = await self.repository.update(doc.id, {"checksum": checksum}) return doc - except Exception: - # If database operation fails, clean up the file + + except Exception as e: + # Clean up on any failure + if 'doc' in locals(): # DB record was created + await self.repository.delete(doc.id) file_path.unlink(missing_ok=True) - raise + raise DocumentWriteError(f"Failed to create document: {e}") async def read_document(self, path: str) -> tuple[Document, str]: """ @@ -223,4 +244,4 @@ class DocumentService(BaseService[DocumentRepository]): # Delete database record if it exists doc = await self.repository.find_by_path(str(path)) if doc: - await self.repository.delete(doc.id) + await self.repository.delete(doc.id) \ No newline at end of file diff --git a/tests/services/test_document_service.py b/tests/services/test_document_service.py index 4083d6e2..d1c9ab68 100644 --- a/tests/services/test_document_service.py +++ b/tests/services/test_document_service.py @@ -3,6 +3,8 @@ import os import pytest import pytest_asyncio +import yaml +from datetime import datetime from pathlib import Path import stat @@ -28,8 +30,8 @@ async def test_doc_path(tmp_path) -> Path: @pytest.mark.asyncio -async def test_create_document_file_first(document_service, test_doc_path): - """Test that files are written before database records.""" +async def test_create_document_with_frontmatter(document_service, test_doc_path): + """Test that created documents have proper frontmatter.""" content = "# Test Document\n\nThis is a test." doc = await document_service.create_document( str(test_doc_path), @@ -37,11 +39,34 @@ async def test_create_document_file_first(document_service, test_doc_path): {"type": "test"} ) - # Verify both file and database record - assert test_doc_path.exists() - assert test_doc_path.read_text() == content + # Verify file content + file_content = test_doc_path.read_text() + + # Parse frontmatter + try: + # Split content at the second "---" marker + _, frontmatter, doc_content = file_content.split("---", 2) + metadata = yaml.safe_load(frontmatter) + except Exception as e: + pytest.fail(f"Failed to parse frontmatter: {e}") + + # Verify frontmatter contents + assert metadata["id"] == doc.id + assert metadata["type"] == "test" + assert "created" in metadata + assert "modified" in metadata + + # Verify timestamps are valid ISO format + datetime.fromisoformat(metadata["created"]) + datetime.fromisoformat(metadata["modified"]) + + # Verify original content is preserved + assert content in doc_content + + # Verify DB record assert doc.path == str(test_doc_path) assert doc.doc_metadata == {"type": "test"} + assert doc.checksum is not None @pytest.mark.asyncio @@ -68,6 +93,32 @@ async def test_create_document_unwriteable_directory(document_service, tmp_path) parent_dir.chmod(stat.S_IWRITE | stat.S_IREAD | stat.S_IEXEC) +@pytest.mark.asyncio +async def test_create_document_cleanup_on_failure(document_service, test_doc_path, monkeypatch): + """Test that failed document creation cleans up DB record.""" + + # Mock write_text to fail after DB record creation + def mock_write_text(*args): + raise PermissionError("Mock write failure") + + # Apply the mock to the specific test file + monkeypatch.setattr(Path, "write_text", mock_write_text) + + with pytest.raises(DocumentWriteError): + await document_service.create_document( + str(test_doc_path), + "content", + {"type": "test"} + ) + + # Verify no DB record remains + doc = await document_service.repository.find_by_path(str(test_doc_path)) + assert doc is None + + # Verify no file was created + assert not test_doc_path.exists() + + @pytest.mark.asyncio async def test_delete_nonexistent_file(document_service, test_doc_path): """Test deleting a file that doesn't exist."""