fix: remove hard-delete from critical path and make PgVectorStore deletions a no-op
This commit is contained in:
Rajat Ahuja
parent
42d4c01d47
commit
cf804ce790
|
|
@ -421,14 +421,10 @@ async def delete_document(
|
|||
session_name: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Delete a single document by ID using hybrid sync/soft delete pattern.
|
||||
Soft-delete a document by ID.
|
||||
|
||||
Soft deletes first (sets deleted_at), then tries to delete from vector store.
|
||||
If vector store delete succeeds, hard deletes from DB.
|
||||
If vector store delete fails, leaves soft-deleted for cleanup job.
|
||||
|
||||
This order ensures crash safety: if the process crashes at any point,
|
||||
the document is either fully deleted or soft-deleted (never orphaned).
|
||||
Sets deleted_at timestamp to mark the document as deleted. The reconciliation
|
||||
job handles vector store cleanup and hard deletion from the database.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
|
|
@ -441,56 +437,28 @@ async def delete_document(
|
|||
Raises:
|
||||
ResourceNotFoundException: If document not found or doesn't match criteria
|
||||
"""
|
||||
# Build base query conditions
|
||||
conditions = [
|
||||
models.Document.id == document_id,
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
models.Document.deleted_at.is_(None), # Only delete non-deleted docs
|
||||
models.Document.deleted_at.is_(None),
|
||||
]
|
||||
if session_name is not None:
|
||||
conditions.append(models.Document.session_name == session_name)
|
||||
|
||||
# Check document exists first
|
||||
check_stmt = select(models.Document).where(*conditions)
|
||||
result = await db.execute(check_stmt)
|
||||
doc = result.scalar_one_or_none()
|
||||
update_stmt = (
|
||||
update(models.Document).where(*conditions).values(deleted_at=func.now())
|
||||
)
|
||||
result = await db.execute(update_stmt)
|
||||
|
||||
if doc is None:
|
||||
if result.rowcount == 0:
|
||||
raise ResourceNotFoundException(
|
||||
f"Document {document_id} not found or does not belong to the specified collection/session"
|
||||
)
|
||||
|
||||
# Step 1: Soft delete first (crash-safe - ensures document is marked for deletion)
|
||||
update_stmt = (
|
||||
update(models.Document)
|
||||
.where(models.Document.id == document_id)
|
||||
.values(deleted_at=func.now())
|
||||
)
|
||||
await db.execute(update_stmt)
|
||||
await db.commit()
|
||||
|
||||
# Step 2: Try to delete from vector store
|
||||
vector_store = get_vector_store()
|
||||
namespace = vector_store.get_vector_namespace(
|
||||
"document", workspace_name, observer, observed
|
||||
)
|
||||
vector_deleted = False
|
||||
|
||||
try:
|
||||
await vector_store.delete_many(namespace, [document_id])
|
||||
vector_deleted = True
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to delete vector for document {document_id}: {e}")
|
||||
|
||||
# Step 3: If vector deleted successfully, hard delete from DB
|
||||
if vector_deleted:
|
||||
delete_stmt = delete(models.Document).where(models.Document.id == document_id)
|
||||
await db.execute(delete_stmt)
|
||||
await db.commit()
|
||||
# If vector delete failed, document stays soft-deleted for cleanup job
|
||||
|
||||
|
||||
async def delete_document_by_id(
|
||||
db: AsyncSession,
|
||||
|
|
@ -498,14 +466,10 @@ async def delete_document_by_id(
|
|||
document_id: str,
|
||||
) -> None:
|
||||
"""
|
||||
Delete a single document by ID and workspace using hybrid sync/soft delete pattern.
|
||||
Soft-delete a document by ID and workspace.
|
||||
|
||||
Soft deletes first (sets deleted_at), then tries to delete from vector store.
|
||||
If vector store delete succeeds, hard deletes from DB.
|
||||
If vector store delete fails, leaves soft-deleted for cleanup job.
|
||||
|
||||
This order ensures crash safety: if the process crashes at any point,
|
||||
the document is either fully deleted or soft-deleted (never orphaned).
|
||||
Sets deleted_at timestamp to mark the document as deleted. The reconciliation
|
||||
job handles vector store cleanup and hard deletion from the database.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
|
|
@ -515,52 +479,24 @@ async def delete_document_by_id(
|
|||
Raises:
|
||||
ResourceNotFoundException: If document not found or doesn't belong to the workspace
|
||||
"""
|
||||
# Fetch document to get observer/observed for namespace
|
||||
stmt = select(models.Document).where(
|
||||
models.Document.id == document_id,
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.deleted_at.is_(None), # Only delete non-deleted docs
|
||||
update_stmt = (
|
||||
update(models.Document)
|
||||
.where(
|
||||
models.Document.id == document_id,
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.deleted_at.is_(None),
|
||||
)
|
||||
.values(deleted_at=func.now())
|
||||
)
|
||||
result = await db.execute(stmt)
|
||||
doc = result.scalar_one_or_none()
|
||||
result = await db.execute(update_stmt)
|
||||
|
||||
if doc is None:
|
||||
if result.rowcount == 0:
|
||||
raise ResourceNotFoundException(
|
||||
f"Document {document_id} not found or does not belong to workspace {workspace_name}"
|
||||
)
|
||||
|
||||
# Step 1: Soft delete first (crash-safe - ensures document is marked for deletion)
|
||||
update_stmt = (
|
||||
update(models.Document)
|
||||
.where(models.Document.id == document_id)
|
||||
.values(deleted_at=func.now())
|
||||
)
|
||||
await db.execute(update_stmt)
|
||||
await db.commit()
|
||||
|
||||
# Step 2: Try to delete from vector store
|
||||
vector_store = get_vector_store()
|
||||
namespace = vector_store.get_vector_namespace(
|
||||
"document",
|
||||
workspace_name,
|
||||
doc.observer,
|
||||
doc.observed,
|
||||
)
|
||||
vector_deleted = False
|
||||
|
||||
try:
|
||||
await vector_store.delete_many(namespace, [document_id])
|
||||
vector_deleted = True
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to delete vector for document {document_id}: {e}")
|
||||
|
||||
# Step 3: If vector deleted successfully, hard delete from DB
|
||||
if vector_deleted:
|
||||
delete_stmt = delete(models.Document).where(models.Document.id == document_id)
|
||||
await db.execute(delete_stmt)
|
||||
await db.commit()
|
||||
# If vector delete failed, document stays soft-deleted for cleanup job
|
||||
|
||||
|
||||
async def create_observations(
|
||||
db: AsyncSession,
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ using the existing embedding columns on documents and message_embeddings tables.
|
|||
import logging
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import delete, select
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models
|
||||
|
|
@ -280,91 +280,35 @@ class PgVectorStore(VectorStore):
|
|||
|
||||
async def delete_many(self, namespace: str, ids: list[str]) -> None:
|
||||
"""
|
||||
Delete vectors by removing the rows from the database.
|
||||
No-op for pgvector. Vector deletion is handled by row deletion.
|
||||
|
||||
For pgvector, since the vector data is stored in the same table as the entities,
|
||||
deleting from the vector store means deleting the actual rows.
|
||||
For pgvector, vectors are stored in the same postgres rows as documents/
|
||||
message_embeddings. The reconciliation job handles hard-deleting rows
|
||||
after calling this method, which removes both the data and the embedding.
|
||||
|
||||
Args:
|
||||
namespace: The namespace containing the vectors
|
||||
ids: List of vector identifiers to delete
|
||||
"""
|
||||
if not ids:
|
||||
return
|
||||
|
||||
table_type, _ = self._parse_namespace(namespace)
|
||||
|
||||
async with tracked_db("pgvector_delete") as db:
|
||||
try:
|
||||
if table_type == "documents":
|
||||
stmt = delete(models.Document).where(models.Document.id.in_(ids))
|
||||
await db.execute(stmt)
|
||||
|
||||
elif table_type == "message_embeddings":
|
||||
for vector_id in ids:
|
||||
try:
|
||||
embedding_id = int(vector_id)
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"Invalid message vector id format: {vector_id}"
|
||||
) from exc
|
||||
|
||||
stmt = delete(models.MessageEmbedding).where(
|
||||
models.MessageEmbedding.id == embedding_id
|
||||
)
|
||||
await db.execute(stmt)
|
||||
|
||||
await db.commit()
|
||||
logger.debug(
|
||||
f"Deleted {len(ids)} rows from {table_type} in namespace {namespace}"
|
||||
)
|
||||
|
||||
except Exception:
|
||||
await db.rollback()
|
||||
logger.exception(
|
||||
f"Failed to delete {len(ids)} rows from namespace {namespace}"
|
||||
)
|
||||
raise
|
||||
if ids:
|
||||
logger.debug(
|
||||
f"PgVectorStore.delete_many() no-op for {len(ids)} vectors in {namespace} (row deletion handles embedding removal)"
|
||||
)
|
||||
|
||||
async def delete_namespace(self, namespace: str) -> None:
|
||||
"""
|
||||
Delete all vectors in a namespace by removing rows from the database.
|
||||
No-op for pgvector. Namespace deletion is handled by row deletion.
|
||||
|
||||
For pgvector, since the vector data is stored in the same table as the entities,
|
||||
deleting a namespace means deleting the actual rows.
|
||||
For pgvector, vectors are stored in the same postgres rows as documents/
|
||||
message_embeddings. Deleting a collection or workspace should delete
|
||||
the rows directly via ORM/SQL, which removes both data and embeddings.
|
||||
|
||||
Args:
|
||||
namespace: The namespace to delete
|
||||
"""
|
||||
table_type, context = self._parse_namespace(namespace)
|
||||
|
||||
async with tracked_db("pgvector_delete_namespace") as db:
|
||||
try:
|
||||
if table_type == "documents":
|
||||
stmt = (
|
||||
delete(models.Document)
|
||||
.where(
|
||||
models.Document.workspace_name == context["workspace_name"]
|
||||
)
|
||||
.where(models.Document.observer == context["observer"])
|
||||
.where(models.Document.observed == context["observed"])
|
||||
)
|
||||
await db.execute(stmt)
|
||||
|
||||
elif table_type == "message_embeddings":
|
||||
stmt = delete(models.MessageEmbedding).where(
|
||||
models.MessageEmbedding.workspace_name
|
||||
== context["workspace_name"]
|
||||
)
|
||||
await db.execute(stmt)
|
||||
|
||||
await db.commit()
|
||||
logger.debug(f"Deleted all rows from namespace {namespace}")
|
||||
|
||||
except Exception:
|
||||
await db.rollback()
|
||||
logger.exception(f"Failed to delete namespace {namespace}")
|
||||
raise
|
||||
logger.debug(
|
||||
f"PgVectorStore.delete_namespace() no-op for {namespace} (row deletion handles embedding removal)"
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the pgvector store (no-op for pgvector)"""
|
||||
|
|
|
|||
|
|
@ -313,9 +313,11 @@ class TestDocumentCRUD:
|
|||
observed=test_peer2.name,
|
||||
)
|
||||
|
||||
# Verify document is deleted
|
||||
# Verify document is soft-deleted
|
||||
result = await db_session.execute(stmt)
|
||||
assert result.scalar_one_or_none() is None
|
||||
doc = result.scalar_one_or_none()
|
||||
assert doc is not None
|
||||
assert doc.deleted_at is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_document_not_found(
|
||||
|
|
|
|||
|
|
@ -624,12 +624,14 @@ class TestObservationRoutes:
|
|||
data = response.json()
|
||||
assert data["message"] == "Observation deleted successfully"
|
||||
|
||||
# Verify observation is deleted
|
||||
# Verify observation is soft-deleted
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(models.Document).where(models.Document.id == observation_id)
|
||||
result = await db_session.execute(stmt)
|
||||
assert result.scalar_one_or_none() is None
|
||||
doc = result.scalar_one_or_none()
|
||||
assert doc is not None
|
||||
assert doc.deleted_at is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_observation_not_found(
|
||||
|
|
|
|||
Loading…
Reference in New Issue