"""Distributed alias-lock primitive: concurrency, staleness and release.""" from __future__ import annotations import asyncio from collections.abc import AsyncIterator from datetime import UTC, datetime, timedelta from uuid import uuid4 import pytest from sqlalchemy.ext.asyncio import AsyncEngine from ragcore.application.indexes.models import IndexOperationStatus from ragcore.config import Settings from ragcore.infrastructure.database import create_database_engine, create_session_factory from ragcore.infrastructure.database.indexes import PostgresIndexOperationStore pytestmark = pytest.mark.integration @pytest.fixture async def engine() -> AsyncIterator[AsyncEngine]: engine = create_database_engine(Settings(env="test")) try: yield engine finally: await engine.dispose() @pytest.fixture def store(engine: AsyncEngine) -> PostgresIndexOperationStore: return PostgresIndexOperationStore(create_session_factory(engine)) def alias() -> str: return f"rag_dense_lock-test-{uuid4().hex}_active" @pytest.mark.asyncio async def test_begin_grants_lock_when_alias_is_free(store: PostgresIndexOperationStore) -> None: now = datetime.now(UTC) operation_id = await store.begin( alias_name=alias(), operation="build", collection_name="rag_dense_x_v1", owner_token=uuid4().hex, now=now, lease_duration=timedelta(seconds=30), ) assert operation_id is not None @pytest.mark.asyncio async def test_concurrent_begin_only_one_succeeds(store: PostgresIndexOperationStore) -> None: target = alias() now = datetime.now(UTC) async def attempt(owner: str) -> object: return await store.begin( alias_name=target, operation="build", collection_name="rag_dense_x_v1", owner_token=owner, now=now, lease_duration=timedelta(seconds=30), ) results = await asyncio.gather(attempt(uuid4().hex), attempt(uuid4().hex)) granted = [result for result in results if result is not None] assert len(granted) == 1 @pytest.mark.asyncio async def test_second_promotion_is_rejected_until_first_finishes( store: PostgresIndexOperationStore, ) -> None: target = alias() now = datetime.now(UTC) first = await store.begin( alias_name=target, operation="build", collection_name="rag_dense_x_v1", owner_token="owner-a", now=now, lease_duration=timedelta(seconds=30), ) assert first is not None blocked = await store.begin( alias_name=target, operation="rollback", collection_name="rag_dense_x_v1", owner_token="owner-b", now=now, lease_duration=timedelta(seconds=30), ) assert blocked is None released = await store.finish( first, "owner-a", IndexOperationStatus.SUCCEEDED, now=now ) assert released is True second = await store.begin( alias_name=target, operation="rollback", collection_name="rag_dense_x_v1", owner_token="owner-b", now=now, lease_duration=timedelta(seconds=30), ) assert second is not None @pytest.mark.asyncio async def test_expired_lease_is_reclaimed_as_needs_repair( store: PostgresIndexOperationStore, ) -> None: target = alias() started = datetime.now(UTC) crashed = await store.begin( alias_name=target, operation="build", collection_name="rag_dense_x_v1", owner_token="owner-crashed", now=started, lease_duration=timedelta(seconds=1), ) assert crashed is not None later = started + timedelta(seconds=5) resumed = await store.begin( alias_name=target, operation="build", collection_name="rag_dense_x_v1", owner_token="owner-resumed", now=later, lease_duration=timedelta(seconds=30), ) assert resumed is not None assert resumed != crashed # The crashed operation must stay detectable, not silently vanish. still_blocked_third_party = await store.begin( alias_name=target, operation="build", collection_name="rag_dense_x_v1", owner_token="owner-third", now=later, lease_duration=timedelta(seconds=30), ) assert still_blocked_third_party is None @pytest.mark.asyncio async def test_finish_is_idempotent_and_requires_terminal_status( store: PostgresIndexOperationStore, ) -> None: target = alias() now = datetime.now(UTC) operation_id = await store.begin( alias_name=target, operation="build", collection_name="rag_dense_x_v1", owner_token="owner-a", now=now, lease_duration=timedelta(seconds=30), ) assert operation_id is not None with pytest.raises(ValueError, match="terminal"): await store.finish(operation_id, "owner-a", IndexOperationStatus.IN_PROGRESS, now=now) first_finish = await store.finish( operation_id, "owner-a", IndexOperationStatus.FAILED_COMPENSATED, now=now ) assert first_finish is True second_finish = await store.finish( operation_id, "owner-a", IndexOperationStatus.FAILED_COMPENSATED, now=now ) assert second_finish is False @pytest.mark.asyncio async def test_lock_scope_is_specific_to_one_alias(store: PostgresIndexOperationStore) -> None: now = datetime.now(UTC) first_alias, second_alias = alias(), alias() first = await store.begin( alias_name=first_alias, operation="build", collection_name="rag_dense_x_v1", owner_token="owner-a", now=now, lease_duration=timedelta(seconds=30), ) second = await store.begin( alias_name=second_alias, operation="build", collection_name="rag_dense_y_v1", owner_token="owner-b", now=now, lease_duration=timedelta(seconds=30), ) assert first is not None assert second is not None @pytest.mark.asyncio async def test_record_step_appends_ordered_entries(store: PostgresIndexOperationStore) -> None: target = alias() now = datetime.now(UTC) operation_id = await store.begin( alias_name=target, operation="build", collection_name="rag_dense_x_v1", owner_token="owner-a", now=now, lease_duration=timedelta(seconds=30), ) assert operation_id is not None await store.record_step(operation_id, "owner-a", "created_shadow", now=now) await store.record_step( operation_id, "owner-a", "activated", now=now, details={"previous_collection": "rag_dense_x_v0"}, ) # A step recorded by a non-owner must not be applied. await store.record_step(operation_id, "owner-mismatched", "hijacked", now=now)