FabledScribe/tests/test_embedding_yields_to_a_delete.py

"""The embedding refresh must LOSE to a delete, not race it (#3262).

Both upserts replace a record's vectors as delete-then-insert. That takes two
row locks — the chunk rows, then the parent row via the insert's foreign key —
in the exact reverse of the order a cascading delete of the parent takes them.
Postgres calls that a deadlock and kills one side at random, which sometimes
means killing the user's delete.

These pin the ORDER, not the outcome: the claim on the parent goes first, and
when the claim fails nothing else in the transaction runs. Compiling the
statement is the only way to assert on a lock mode without a database — the
integration twin (test_integration_embedding_yields_to_delete.py) proves the
behaviour against a real one.
"""
from unittest.mock import AsyncMock, MagicMock, patch

from sqlalchemy.dialects import postgresql
from sqlalchemy.exc import OperationalError

from scribe.services import embeddings as emb
from tests.helpers import compiled_sql

ONE_VECTOR = [[0.0] * 384]

# The lock mode is a Postgres extension — the generic dialect renders a plain
# FOR UPDATE and would pass an assertion that proves nothing.
PG = postgresql.dialect()


def _mock_session(lock_result: object = 7, execute_side_effect=None):
    """A session stand-in whose first execute answers the parent-row claim."""
    session = MagicMock()
    claimed = MagicMock()
    claimed.scalar_one_or_none.return_value = lock_result
    # A task's document carries its work logs (#4251), so the refresh reads
    # task_logs before it builds chunks. Answered explicitly — left to
    # autovivify, `list(result.all())` would raise and be swallowed, and these
    # tests would be exercising the failure path without saying so.
    claimed.all.return_value = []
    if execute_side_effect is not None:
        session.execute = AsyncMock(side_effect=execute_side_effect)
    else:
        session.execute = AsyncMock(return_value=claimed)
    session.commit = AsyncMock()
    session.add = MagicMock()
    ctx = MagicMock()
    ctx.__aenter__ = AsyncMock(return_value=session)
    ctx.__aexit__ = AsyncMock(return_value=False)
    return session, ctx


def _lock_unavailable() -> OperationalError:
    """What asyncpg raises through SQLAlchemy when NOWAIT can't take the row."""
    return OperationalError("SELECT ...", {}, Exception("lock not available"))


async def test_a_note_refresh_claims_the_row_before_rewriting_its_vectors():
    """The claim is FIRST, and it is FOR KEY SHARE NOWAIT.

    FOR KEY SHARE because that is exactly the lock the insert's foreign key
    takes anyway — it conflicts with a delete of the note and with nothing
    else, so an ordinary edit is unaffected. NOWAIT because the whole point is
    to lose immediately rather than queue up behind the delete and hold the
    chunk rows while doing it.
    """
    session, ctx = _mock_session()
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)),
    ):
        await emb.upsert_note_embedding(7, 42, "T", "a short body")

    sqls = [compiled_sql(c.args[0], dialect=PG) for c in session.execute.call_args_list]
    # Located by what they ARE rather than by position: the work-log read a
    # task's document needs (#4251) runs before any of this, and an index is
    # not what the claim is about.
    claim = next(i for i, sql in enumerate(sqls) if "FOR KEY SHARE NOWAIT" in sql)
    replace = next(
        i for i, sql in enumerate(sqls) if sql.startswith("DELETE FROM note_embeddings")
    )
    assert sqls[claim].startswith("SELECT notes.id")
    assert claim < replace, "the claim goes first — that is the whole fix"
    assert not any("note_embeddings" in sql for sql in sqls[:claim]), (
        "nothing may touch the chunk rows before the parent is claimed"
    )
    session.add.assert_called()


async def test_a_rule_refresh_claims_the_row_before_rewriting_its_vectors():
    """The rule twin — the path the reported deadlock actually took."""
    session, ctx = _mock_session()
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)),
    ):
        await emb.upsert_rule_embedding(9, "T", "a short statement", "on write")

    claim, replace = [c.args[0] for c in session.execute.call_args_list][:2]
    assert compiled_sql(claim, dialect=PG).startswith("SELECT rules.id")
    assert "FOR KEY SHARE NOWAIT" in compiled_sql(claim, dialect=PG)
    assert compiled_sql(replace, dialect=PG).startswith("DELETE FROM rule_embeddings")
    session.add.assert_called()


async def test_a_record_being_deleted_is_left_alone_rather_than_raced():
    """The claim failing ends the write — it does not fall through to the
    delete-and-insert that would take the locks in the losing order."""
    session, ctx = _mock_session(execute_side_effect=_lock_unavailable())
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)),
    ):
        await emb.upsert_rule_embedding(9, "T", "a short statement", "on write")

    assert session.execute.await_count == 1, "it stopped at the claim"
    session.add.assert_not_called()
    session.commit.assert_not_awaited()


async def test_a_record_already_gone_is_not_re_embedded():
    """A vector inserted for a row that no longer exists is either a foreign
    key violation or, worse, a resurrected chunk. Nothing to refresh."""
    session, ctx = _mock_session(lock_result=None)
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)),
    ):
        await emb.upsert_note_embedding(7, 42, "T", "a short body")

    sqls = [compiled_sql(c.args[0], dialect=PG) for c in session.execute.call_args_list]
    assert any("FOR KEY SHARE NOWAIT" in sql for sql in sqls), "it did reach the claim"
    assert not any("note_embeddings" in sql for sql in sqls), "and stopped there"
    session.add.assert_not_called()
    session.commit.assert_not_awaited()