498 lines
18 KiB
Python
498 lines
18 KiB
Python
"""Real-Postgres coverage for the source registry and its document linkage:
|
|
ensure-on-first-use idempotency, list/filter/count, status and sync-state
|
|
bookkeeping, sync runs and remembered items, and the documents.source_id
|
|
foreign key the content schema now enforces."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
from collections.abc import AsyncIterator
|
|
from datetime import UTC, datetime, timedelta
|
|
from uuid import UUID, uuid4
|
|
|
|
import pytest
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
from ragcore.application.documents import DocumentListFilter, DocumentRecord
|
|
from ragcore.application.sources import (
|
|
SourceItemRecord,
|
|
SourceItemSummary,
|
|
SourceListFilter,
|
|
SourceOwnershipConflict,
|
|
SourceRecord,
|
|
SourceStatus,
|
|
SourceSyncOutcome,
|
|
SourceSyncRunRecord,
|
|
implicit_source_record,
|
|
uploads_source_record,
|
|
)
|
|
from ragcore.config import Settings
|
|
from ragcore.domain.sources import ConnectorKind, SourceSyncState, upload_source_id
|
|
from ragcore.infrastructure.database import create_database_engine, create_session_factory
|
|
from ragcore.infrastructure.database.content import PostgresDocumentRepository
|
|
from ragcore.infrastructure.database.sources import PostgresSourceRegistry
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def event_loop_policy() -> asyncio.AbstractEventLoopPolicy:
|
|
if sys.platform == "win32":
|
|
return asyncio.WindowsSelectorEventLoopPolicy()
|
|
return asyncio.get_event_loop_policy()
|
|
|
|
|
|
@pytest.fixture
|
|
async def db_engine() -> AsyncIterator[AsyncEngine]:
|
|
engine = create_database_engine(Settings(env="test"))
|
|
try:
|
|
yield engine
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
async def test_ensure_source_is_idempotent_and_upload_sources_are_builtin(
|
|
db_engine: AsyncEngine,
|
|
) -> None:
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
workspace_id, space_id = uuid4(), uuid4()
|
|
record = uploads_source_record(workspace_id=workspace_id, space_id=space_id)
|
|
|
|
first = await registry.ensure_source(record)
|
|
second = await registry.ensure_source(
|
|
SourceRecord(
|
|
id=record.id,
|
|
workspace_id=workspace_id,
|
|
space_id=space_id,
|
|
name="Renamed by a later caller",
|
|
kind=ConnectorKind.UPLOAD,
|
|
)
|
|
)
|
|
assert first.id == upload_source_id(space_id)
|
|
assert second.name == "Uploads", "ensure never overwrites an existing source"
|
|
assert second.is_builtin_uploads
|
|
assert not second.syncable
|
|
|
|
|
|
async def test_documents_require_a_registered_source(db_engine: AsyncEngine) -> None:
|
|
factory = create_session_factory(db_engine)
|
|
documents = PostgresDocumentRepository(factory)
|
|
orphan = DocumentRecord(
|
|
id=uuid4(),
|
|
workspace_id=uuid4(),
|
|
space_id=uuid4(),
|
|
source_key=str(uuid4()),
|
|
external_id="page",
|
|
title="Orphan",
|
|
)
|
|
with pytest.raises(IntegrityError):
|
|
await documents.create_document(orphan)
|
|
|
|
|
|
async def test_list_filter_count_status_and_sync_bookkeeping(db_engine: AsyncEngine) -> None:
|
|
factory = create_session_factory(db_engine)
|
|
registry = PostgresSourceRegistry(factory)
|
|
documents = PostgresDocumentRepository(factory)
|
|
workspace_id, space_id = uuid4(), uuid4()
|
|
folder = await registry.create_source(
|
|
SourceRecord(
|
|
id=uuid4(),
|
|
workspace_id=workspace_id,
|
|
space_id=space_id,
|
|
name="Infra runbooks",
|
|
kind=ConnectorKind.FOLDER,
|
|
config={"root": "/mnt/runbooks", "trust_level": "trusted_internal"},
|
|
)
|
|
)
|
|
git = await registry.create_source(
|
|
SourceRecord(
|
|
id=uuid4(),
|
|
workspace_id=workspace_id,
|
|
space_id=space_id,
|
|
name="Platform repo",
|
|
kind=ConnectorKind.GIT,
|
|
config={"repository": "platform", "ref": "main", "checkout_root": "/srv/platform"},
|
|
)
|
|
)
|
|
for index in range(3):
|
|
document = DocumentRecord(
|
|
id=uuid4(),
|
|
workspace_id=workspace_id,
|
|
space_id=space_id,
|
|
source_key=str(folder.id),
|
|
external_id=f"runbooks/{index}.md",
|
|
title=f"Runbook {index}",
|
|
source_id=folder.id,
|
|
)
|
|
await documents.create_document(document)
|
|
|
|
page, *_ = (await registry.list_sources(workspace_id, SourceListFilter()),)
|
|
assert {source.name for source in page.items} == {"Infra runbooks", "Platform repo"}
|
|
only_git = await registry.list_sources(workspace_id, SourceListFilter(kind=ConnectorKind.GIT))
|
|
assert [source.id for source in only_git.items] == [git.id]
|
|
searched = await registry.list_sources(workspace_id, SourceListFilter(search="infra"))
|
|
assert [source.id for source in searched.items] == [folder.id]
|
|
assert (await registry.list_sources(uuid4(), SourceListFilter())).total == 0
|
|
|
|
counts = await documents.count_documents_by_source((folder.id, git.id))
|
|
assert counts == {folder.id: 3, git.id: 0}
|
|
listed = await documents.list_documents(
|
|
workspace_id, DocumentListFilter(source_id=folder.id, search="Runbook 1")
|
|
)
|
|
assert [item.document.title for item in listed.items] == ["Runbook 1"]
|
|
assert listed.total == 1
|
|
|
|
now = datetime.now(UTC)
|
|
paused = await registry.set_status(folder.id, SourceStatus.PAUSED, now=now)
|
|
assert paused is not None and paused.status is SourceStatus.PAUSED
|
|
deleted = await registry.set_status(git.id, SourceStatus.DELETED, now=now)
|
|
assert deleted is not None
|
|
assert (await registry.list_sources(workspace_id, SourceListFilter())).total == 1, (
|
|
"deleted sources are hidden unless asked for"
|
|
)
|
|
assert (
|
|
await registry.list_sources(workspace_id, SourceListFilter(status=SourceStatus.DELETED))
|
|
).total == 1
|
|
|
|
run = SourceSyncRunRecord(
|
|
id=uuid4(),
|
|
source_id=folder.id,
|
|
started_at=now,
|
|
outcome=SourceSyncOutcome.RUNNING,
|
|
)
|
|
await registry.record_sync_run(run)
|
|
finished = SourceSyncRunRecord(
|
|
id=run.id,
|
|
source_id=folder.id,
|
|
started_at=now,
|
|
finished_at=now,
|
|
outcome=SourceSyncOutcome.COMPLETED,
|
|
discovered=3,
|
|
new=3,
|
|
revision="abc1234",
|
|
)
|
|
await registry.record_sync_run(finished)
|
|
[stored_run] = await registry.list_sync_runs(folder.id)
|
|
assert stored_run.outcome is SourceSyncOutcome.COMPLETED
|
|
assert stored_run.new == 3 and stored_run.revision == "abc1234"
|
|
synced = await registry.set_sync_state(
|
|
folder.id, SourceSyncState.COMPLETED, now=now, revision="abc1234", successful=True
|
|
)
|
|
assert synced is not None
|
|
assert synced.last_successful_sync_at is not None
|
|
assert synced.last_revision == "abc1234"
|
|
|
|
await registry.upsert_items(
|
|
(
|
|
SourceItemRecord(
|
|
source_id=folder.id,
|
|
external_id="runbooks/0.md",
|
|
fingerprint="f" * 64,
|
|
locator="runbooks/0.md",
|
|
last_seen_at=now,
|
|
),
|
|
)
|
|
)
|
|
await registry.mark_items_removed(folder.id, ("runbooks/0.md",), now=now)
|
|
[item] = await registry.list_items(folder.id)
|
|
assert item.removed is True
|
|
|
|
implied = implicit_source_record(
|
|
source_id=folder.id, workspace_id=workspace_id, space_id=space_id
|
|
)
|
|
assert implied.kind is ConnectorKind.UPLOAD # an implied record never guesses a connector
|
|
assert (await registry.ensure_source(implied)).kind is ConnectorKind.FOLDER
|
|
|
|
|
|
# -- ownership: a source belongs to exactly one workspace and space ----------
|
|
|
|
|
|
def _source(source_id, workspace_id, space_id, *, name: str = "Runbooks") -> SourceRecord:
|
|
return SourceRecord(
|
|
id=source_id,
|
|
workspace_id=workspace_id,
|
|
space_id=space_id,
|
|
name=name,
|
|
kind=ConnectorKind.FOLDER,
|
|
config={"root": "/mnt/runbooks", "trust_level": "trusted_internal"},
|
|
)
|
|
|
|
|
|
async def test_ensure_source_refuses_an_id_owned_by_another_workspace(
|
|
db_engine: AsyncEngine,
|
|
) -> None:
|
|
"""``ON CONFLICT DO NOTHING`` alone would silently hand the caller back
|
|
another tenant's row, and every document of this ingest would attach to
|
|
it. The ownership assertion is what turns "some source with this id
|
|
exists" into "this source, in this scope, exists"."""
|
|
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
source_id, space_id = uuid4(), uuid4()
|
|
owner = uuid4()
|
|
await registry.ensure_source(_source(source_id, owner, space_id))
|
|
|
|
with pytest.raises(SourceOwnershipConflict) as raised:
|
|
await registry.ensure_source(_source(source_id, uuid4(), space_id, name="Theirs"))
|
|
|
|
conflict = raised.value
|
|
assert conflict.cross_workspace is True
|
|
assert conflict.audit_details()["boundary"] == "workspace"
|
|
assert "another workspace" in str(conflict)
|
|
stored = await registry.get_source(source_id)
|
|
assert stored is not None and stored.workspace_id == owner
|
|
|
|
|
|
async def test_ensure_source_refuses_an_id_owned_by_another_space(db_engine: AsyncEngine) -> None:
|
|
"""The space boundary rejects too: a space is what grants and policies
|
|
attach to, so a source shared between two spaces would make one space's
|
|
retrieval filter answer for content the other owns."""
|
|
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
source_id, workspace_id = uuid4(), uuid4()
|
|
await registry.ensure_source(_source(source_id, workspace_id, uuid4()))
|
|
|
|
with pytest.raises(SourceOwnershipConflict) as raised:
|
|
await registry.ensure_source(_source(source_id, workspace_id, uuid4()))
|
|
|
|
assert raised.value.cross_workspace is False
|
|
assert raised.value.audit_details()["boundary"] == "space"
|
|
|
|
|
|
async def test_a_document_cannot_reference_a_source_from_another_space(
|
|
db_engine: AsyncEngine,
|
|
) -> None:
|
|
"""Enforced by PostgreSQL, not by remembering to check: the composite
|
|
foreign key makes the cross-scope reference unrepresentable."""
|
|
|
|
factory = create_session_factory(db_engine)
|
|
registry = PostgresSourceRegistry(factory)
|
|
documents = PostgresDocumentRepository(factory)
|
|
workspace_id, space_id = uuid4(), uuid4()
|
|
source = await registry.create_source(_source(uuid4(), workspace_id, space_id))
|
|
|
|
with pytest.raises(IntegrityError):
|
|
await documents.create_document(
|
|
DocumentRecord(
|
|
id=uuid4(),
|
|
workspace_id=workspace_id,
|
|
space_id=uuid4(), # a different space than the source's
|
|
source_key=str(source.id),
|
|
external_id="page",
|
|
title="Smuggled",
|
|
source_id=source.id,
|
|
)
|
|
)
|
|
|
|
|
|
async def test_get_owned_source_hides_a_foreign_source_rather_than_refusing_it(
|
|
db_engine: AsyncEngine,
|
|
) -> None:
|
|
"""A 404/403 difference would itself tell a caller which source ids are
|
|
real in someone else's tenant."""
|
|
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
workspace_id, space_id = uuid4(), uuid4()
|
|
source = await registry.create_source(_source(uuid4(), workspace_id, space_id))
|
|
|
|
assert await registry.get_owned_source(source.id, workspace_id=workspace_id) is not None
|
|
assert await registry.get_owned_source(source.id, workspace_id=uuid4()) is None
|
|
assert (
|
|
await registry.get_owned_source(source.id, workspace_id=workspace_id, space_id=uuid4())
|
|
is None
|
|
)
|
|
|
|
|
|
async def test_update_source_cannot_move_a_source_between_spaces(db_engine: AsyncEngine) -> None:
|
|
"""Ownership is fixed at creation: an update may change what a source
|
|
*is*, never which tenant or space it belongs to."""
|
|
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
workspace_id, space_id = uuid4(), uuid4()
|
|
source = await registry.create_source(_source(uuid4(), workspace_id, space_id))
|
|
|
|
with pytest.raises(SourceOwnershipConflict):
|
|
await registry.update_source(
|
|
_source(source.id, workspace_id, uuid4(), name="Moved elsewhere")
|
|
)
|
|
|
|
unchanged = await registry.get_source(source.id)
|
|
assert unchanged is not None
|
|
assert unchanged.space_id == space_id and unchanged.name == "Runbooks"
|
|
|
|
|
|
# -- convergence: discovered versus applied ----------------------------------
|
|
|
|
|
|
async def test_items_converge_only_when_an_ingestion_reports_the_same_fingerprint(
|
|
db_engine: AsyncEngine,
|
|
) -> None:
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
source = await registry.create_source(_source(uuid4(), uuid4(), uuid4()))
|
|
now = datetime.now(UTC)
|
|
await registry.upsert_items(
|
|
(
|
|
SourceItemRecord(
|
|
source_id=source.id,
|
|
external_id="a.md",
|
|
fingerprint="a" * 64,
|
|
locator="a.md",
|
|
last_seen_at=now,
|
|
),
|
|
SourceItemRecord(
|
|
source_id=source.id,
|
|
external_id="b.md",
|
|
fingerprint="b" * 64,
|
|
locator="b.md",
|
|
last_seen_at=now,
|
|
),
|
|
)
|
|
)
|
|
|
|
assert await registry.summarize_items(source.id) == SourceItemSummary(total=2, pending=2)
|
|
|
|
# A report for content that has since changed applies nothing.
|
|
assert (
|
|
await registry.mark_item_ingested(
|
|
source.id, "a.md", fingerprint="f" * 64, document_id=None, now=now
|
|
)
|
|
is False
|
|
)
|
|
assert await registry.summarize_items(source.id) == SourceItemSummary(total=2, pending=2)
|
|
|
|
assert (
|
|
await registry.mark_item_ingested(
|
|
source.id, "a.md", fingerprint="a" * 64, document_id=None, now=now
|
|
)
|
|
is True
|
|
)
|
|
assert (
|
|
await registry.mark_item_failed(
|
|
source.id,
|
|
"b.md",
|
|
fingerprint="b" * 64,
|
|
failure_code="security_quarantined",
|
|
retryable=False,
|
|
now=now,
|
|
)
|
|
is True
|
|
)
|
|
|
|
summary = await registry.summarize_items(source.id)
|
|
assert summary == SourceItemSummary(total=2, applied=1, failed=1)
|
|
assert summary.converged is False
|
|
assert summary.outcome() is SourceSyncOutcome.COMPLETED_WITH_FAILURES
|
|
|
|
[applied] = [item for item in await registry.list_items(source.id) if item.applied]
|
|
assert applied.external_id == "a.md" and applied.applied_at is not None
|
|
|
|
|
|
async def test_a_retired_item_counts_as_converged(db_engine: AsyncEngine) -> None:
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
source = await registry.create_source(_source(uuid4(), uuid4(), uuid4()))
|
|
now = datetime.now(UTC)
|
|
await registry.upsert_items(
|
|
(
|
|
SourceItemRecord(
|
|
source_id=source.id,
|
|
external_id="gone.md",
|
|
fingerprint="c" * 64,
|
|
locator="gone.md",
|
|
last_seen_at=now,
|
|
),
|
|
)
|
|
)
|
|
|
|
await registry.mark_items_removed(source.id, ("gone.md",), now=now)
|
|
|
|
summary = await registry.summarize_items(source.id)
|
|
assert summary == SourceItemSummary(total=1, removed=1)
|
|
assert summary.converged is True
|
|
|
|
|
|
async def test_convergence_never_overwrites_a_sync_that_is_still_running(
|
|
db_engine: AsyncEngine,
|
|
) -> None:
|
|
"""A child job finishing while a *new* sync is already scanning must not
|
|
stamp its stale conclusion over the live run."""
|
|
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
source = await registry.create_source(_source(uuid4(), uuid4(), uuid4()))
|
|
now = datetime.now(UTC)
|
|
|
|
await registry.set_sync_state(source.id, SourceSyncState.AWAITING_INGEST, now=now)
|
|
settled = await registry.converge_sync_state(
|
|
source.id, SourceSyncState.COMPLETED, now=now, successful=True
|
|
)
|
|
assert settled is not None
|
|
assert settled.sync_state is SourceSyncState.COMPLETED
|
|
assert settled.last_successful_sync_at is not None
|
|
|
|
await registry.set_sync_state(source.id, SourceSyncState.SCANNING, now=now)
|
|
assert (
|
|
await registry.converge_sync_state(
|
|
source.id, SourceSyncState.COMPLETED, now=now, successful=True
|
|
)
|
|
is None
|
|
)
|
|
live = await registry.get_source(source.id)
|
|
assert live is not None and live.sync_state is SourceSyncState.SCANNING
|
|
|
|
|
|
async def test_list_due_for_sync_finds_only_what_should_run_now(db_engine: AsyncEngine) -> None:
|
|
"""The scheduler's query against real PostgreSQL, including the interval
|
|
arithmetic -- which is done in SQL, so a Python-side rewrite of the same
|
|
rule would not prove it."""
|
|
|
|
registry = PostgresSourceRegistry(create_session_factory(db_engine))
|
|
workspace_id, space_id = uuid4(), uuid4()
|
|
now = datetime.now(UTC)
|
|
|
|
async def _register(
|
|
name: str,
|
|
*,
|
|
interval: int | None,
|
|
last_sync_at: datetime | None,
|
|
status: SourceStatus = SourceStatus.ACTIVE,
|
|
sync_state: SourceSyncState = SourceSyncState.COMPLETED,
|
|
kind: ConnectorKind = ConnectorKind.FOLDER,
|
|
) -> UUID:
|
|
record = await registry.create_source(
|
|
SourceRecord(
|
|
id=uuid4(),
|
|
workspace_id=workspace_id,
|
|
space_id=space_id,
|
|
name=name,
|
|
kind=kind,
|
|
status=status,
|
|
sync_state=sync_state,
|
|
last_sync_at=last_sync_at,
|
|
sync_interval_seconds=interval,
|
|
)
|
|
)
|
|
return record.id
|
|
|
|
overdue = await _register("overdue", interval=3600, last_sync_at=now - timedelta(hours=2))
|
|
never = await _register("never synced", interval=3600, last_sync_at=None)
|
|
await _register("recent", interval=3600, last_sync_at=now - timedelta(minutes=5))
|
|
await _register("manual", interval=None, last_sync_at=now - timedelta(days=9))
|
|
await _register(
|
|
"paused", interval=3600, last_sync_at=now - timedelta(days=1), status=SourceStatus.PAUSED
|
|
)
|
|
await _register(
|
|
"already running",
|
|
interval=3600,
|
|
last_sync_at=now - timedelta(days=1),
|
|
sync_state=SourceSyncState.AWAITING_INGEST,
|
|
)
|
|
await _register(
|
|
"upload", interval=3600, last_sync_at=None, kind=ConnectorKind.UPLOAD
|
|
)
|
|
|
|
due = await registry.list_due_for_sync(now=now, limit=50)
|
|
|
|
assert {source.id for source in due} == {overdue, never}
|
|
# Longest-neglected first, so a small batch cannot starve one source.
|
|
assert due[0].id == never
|