leocrm/app/plugins/builtins/unified_search/lifecycle.py

"""Derived-data lifecycle management for the Unified Search index.

Provides functions to remove, rebuild, and purge search index data
(embeddings, FTS vectors, document chunks) when entities are deleted,
restored, or corrected.
"""

from __future__ import annotations

import logging
import uuid
from typing import TYPE_CHECKING

from sqlalchemy import text

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

# Entity type -> (table_name, tsv_column, embedding_column)
_ENTITY_MAP: dict[str, tuple[str, str, str]] = {
    "contact": ("contacts", "search_tsv", "embedding"),
    "mail": ("mails", "body_tsv", "embedding"),
    "file": ("files", "content_tsv", "embedding"),
    "event": ("calendar_entries", "search_tsv", "embedding"),
}


def _resolve_entity(entity_type: str) -> tuple[str, str, str] | None:
    """Return (table, tsv_column, embedding_column) for an entity type."""
    return _ENTITY_MAP.get(entity_type)


async def remove_from_index(
    db: "AsyncSession",
    entity_type: str,
    entity_id: uuid.UUID,
    tenant_id: uuid.UUID,
) -> None:
    """Remove an entity from the search index (embedding=NULL, TSV=NULL).

    Does **not** delete document_chunks — use :func:`remove_chunks` for that.
    """
    mapping = _resolve_entity(entity_type)
    if mapping is None:
        logger.warning("remove_from_index: unknown entity_type=%s", entity_type)
        return

    table, tsv_col, emb_col = mapping
    await db.execute(
        text(
            f"UPDATE {table} "
            f"SET {emb_col} = NULL, {tsv_col} = NULL, indexed_at = NULL "
            f"WHERE id = :eid AND tenant_id = :tid"
        ),
        {"eid": entity_id, "tid": tenant_id},
    )
    await db.commit()
    logger.info("Removed %s/%s from search index", entity_type, entity_id)


async def remove_chunks(
    db: "AsyncSession",
    file_id: uuid.UUID,
    tenant_id: uuid.UUID,
) -> None:
    """Remove all document_chunks for a file."""
    await db.execute(
        text(
            "DELETE FROM document_chunks "
            "WHERE file_id = :fid AND tenant_id = :tid"
        ),
        {"fid": file_id, "tid": tenant_id},
    )
    await db.commit()
    logger.info("Removed chunks for file %s", file_id)


async def rebuild_index(
    db: "AsyncSession",
    entity_type: str,
    entity_id: uuid.UUID,
    tenant_id: uuid.UUID,
) -> bool:
    """Rebuild the search index for a single entity from the authoritative source.

    Re-generates the embedding and refreshes the TSV vector.
    Returns True on success, False on failure.
    """
    from app.plugins.builtins.unified_search.embedding import index_entity

    # Clear existing index data first
    await remove_from_index(db, entity_type, entity_id, tenant_id)

    # Re-generate embedding (index_entity also updates indexed_at)
    try:
        success = await index_entity(entity_type, entity_id, tenant_id, db)
        if success and entity_type == "file":
            # Re-index chunks for files
            from app.plugins.builtins.unified_search.jobs import index_file_chunks
            await index_file_chunks({"_lifecycle": True}, str(entity_id))
        return success
    except Exception:
        logger.exception("rebuild_index failed for %s/%s", entity_type, entity_id)
        return False


async def handle_entity_delete(
    db: "AsyncSession",
    entity_type: str,
    entity_id: uuid.UUID,
    tenant_id: uuid.UUID,
) -> None:
    """Handle entity deletion: remove from index + remove chunks if applicable."""
    await remove_from_index(db, entity_type, entity_id, tenant_id)
    if entity_type == "file":
        await remove_chunks(db, entity_id, tenant_id)


async def handle_entity_restore(
    db: "AsyncSession",
    entity_type: str,
    entity_id: uuid.UUID,
    tenant_id: uuid.UUID,
) -> None:
    """Handle entity restore from trash: rebuild search index."""
    await rebuild_index(db, entity_type, entity_id, tenant_id)


async def handle_entity_correction(
    db: "AsyncSession",
    entity_type: str,
    entity_id: uuid.UUID,
    tenant_id: uuid.UUID,
) -> None:
    """Handle entity data correction: rebuild search index."""
    await rebuild_index(db, entity_type, entity_id, tenant_id)