Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e
This commit is contained in:
Public Source Publisher committed
commit
7ca726caf2
659 files changed
+98171
No files matched your search
@@ -0,0 +1,337 @@
|
||||
"""End-to-end empty-environment restore gate: backup a real source, restore
|
||||
it into disposable Postgres/Qdrant, and prove counts, golden queries and
|
||||
cross-space denial all reproduce -- without ever touching the source."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import sys
|
||||
from collections.abc import AsyncIterator
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
from qdrant_client.http.exceptions import UnexpectedResponse
|
||||
|
||||
from ragcore.application.backup import BackupService
|
||||
from ragcore.application.restore import GoldenQueryProbe, RestoreGateService, SecurityDenialProbe
|
||||
from ragcore.config import Settings
|
||||
from ragcore.domain.citations import Citation
|
||||
from ragcore.domain.control import DataPolicy, KnowledgeSpace, Role, UserGrant, Workspace
|
||||
from ragcore.infrastructure.backup import (
|
||||
LocalFileCopier,
|
||||
PgDumpPostgresDumper,
|
||||
QdrantHttpSnapshotter,
|
||||
)
|
||||
from ragcore.infrastructure.database import create_database_engine, create_session_factory
|
||||
from ragcore.infrastructure.database.control.repositories import ControlPlaneRepository
|
||||
from ragcore.infrastructure.restore import (
|
||||
QdrantHttpRestorer,
|
||||
SqlAlchemyControlPlaneReader,
|
||||
SqlAlchemyPostgresRestorer,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
QUERY_VECTOR = (0.1, 0.2, 0.3, 0.4)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def event_loop_policy():
|
||||
if sys.platform == "win32":
|
||||
return asyncio.WindowsSelectorEventLoopPolicy()
|
||||
return asyncio.get_event_loop_policy()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def settings() -> Settings:
|
||||
return Settings(env="test")
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def session_factory(settings: Settings):
|
||||
engine = create_database_engine(settings)
|
||||
try:
|
||||
yield create_session_factory(engine)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def qdrant_client(settings: Settings) -> AsyncIterator[AsyncQdrantClient]:
|
||||
client = AsyncQdrantClient(
|
||||
url=settings.qdrant_url, api_key=settings.qdrant_api_key.get_secret_value(), timeout=300
|
||||
)
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
class _SourceFixture:
|
||||
"""Real control-plane rows and a real Qdrant collection, torn down after."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.workspace_id = uuid4()
|
||||
self.granted_space_id = uuid4()
|
||||
self.denied_space_id = uuid4()
|
||||
self.principal_id = uuid4()
|
||||
self.document_id = uuid4()
|
||||
self.document_version_id = uuid4()
|
||||
self.chunk_id = uuid4()
|
||||
self.collection_name = f"restore-gate-{uuid4().hex[:12]}"
|
||||
self.alias = f"{self.collection_name}-live"
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def source(
|
||||
session_factory, qdrant_client: AsyncQdrantClient
|
||||
) -> AsyncIterator[_SourceFixture]:
|
||||
fixture = _SourceFixture()
|
||||
async with session_factory() as session:
|
||||
repo = ControlPlaneRepository(session)
|
||||
await repo.create_workspace(
|
||||
Workspace(id=fixture.workspace_id, name="Restore Gate WS", slug=f"rg-{uuid4().hex[:8]}")
|
||||
)
|
||||
await repo.create_space(
|
||||
KnowledgeSpace(
|
||||
id=fixture.granted_space_id,
|
||||
workspace_id=fixture.workspace_id,
|
||||
name="Granted",
|
||||
slug=f"granted-{uuid4().hex[:8]}",
|
||||
data_policy=DataPolicy.LOCAL_ONLY,
|
||||
)
|
||||
)
|
||||
await repo.create_space(
|
||||
KnowledgeSpace(
|
||||
id=fixture.denied_space_id,
|
||||
workspace_id=fixture.workspace_id,
|
||||
name="Denied",
|
||||
slug=f"denied-{uuid4().hex[:8]}",
|
||||
data_policy=DataPolicy.LOCAL_ONLY,
|
||||
)
|
||||
)
|
||||
await repo.create_user_grant(
|
||||
UserGrant(
|
||||
id=uuid4(),
|
||||
workspace_id=fixture.workspace_id,
|
||||
user_id=fixture.principal_id,
|
||||
knowledge_space_id=fixture.granted_space_id,
|
||||
role=Role.VIEWER,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
await qdrant_client.create_collection(
|
||||
collection_name=fixture.collection_name,
|
||||
vectors_config={
|
||||
"dense": qdrant_models.VectorParams(size=4, distance=qdrant_models.Distance.COSINE)
|
||||
},
|
||||
)
|
||||
await qdrant_client.upsert(
|
||||
collection_name=fixture.collection_name,
|
||||
points=[
|
||||
qdrant_models.PointStruct(
|
||||
id=str(fixture.chunk_id),
|
||||
vector={"dense": list(QUERY_VECTOR)},
|
||||
payload={
|
||||
"workspace_id": str(fixture.workspace_id),
|
||||
"space_id": str(fixture.granted_space_id),
|
||||
"document_id": str(fixture.document_id),
|
||||
"document_version_id": str(fixture.document_version_id),
|
||||
"chunk_id": str(fixture.chunk_id),
|
||||
"location": {"page": 1},
|
||||
"text": "The restore gate must reproduce this exact passage.",
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
await qdrant_client.update_collection_aliases(
|
||||
change_aliases_operations=[
|
||||
qdrant_models.CreateAliasOperation(
|
||||
create_alias=qdrant_models.CreateAlias(
|
||||
collection_name=fixture.collection_name, alias_name=fixture.alias
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
try:
|
||||
yield fixture
|
||||
finally:
|
||||
await qdrant_client.delete_collection(fixture.collection_name)
|
||||
|
||||
|
||||
async def _snapshot_source(
|
||||
fixture: _SourceFixture, settings: Settings, session_factory, tmp_path: Path
|
||||
) -> Path:
|
||||
qdrant_client = AsyncQdrantClient(
|
||||
url=settings.qdrant_url, api_key=settings.qdrant_api_key.get_secret_value(), timeout=300
|
||||
)
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=300) as http_client:
|
||||
backup_service = BackupService(
|
||||
postgres=PgDumpPostgresDumper(
|
||||
database_url=settings.database_url, session_factory=session_factory
|
||||
),
|
||||
qdrant=QdrantHttpSnapshotter(
|
||||
client=qdrant_client,
|
||||
http_client=http_client,
|
||||
base_url=settings.qdrant_url,
|
||||
api_key=settings.qdrant_api_key.get_secret_value(),
|
||||
),
|
||||
files=LocalFileCopier(),
|
||||
)
|
||||
manifest = await backup_service.run(
|
||||
backup_root=tmp_path / "backups",
|
||||
blob_root=tmp_path / "blobs-src",
|
||||
parsed_root=tmp_path / "parsed-src",
|
||||
collection_names=[fixture.collection_name],
|
||||
active_profiles={},
|
||||
)
|
||||
finally:
|
||||
await qdrant_client.close()
|
||||
return tmp_path / "backups" / manifest.backup_id
|
||||
|
||||
|
||||
async def test_restore_gate_reproduces_counts_citations_and_denies_cross_space_access(
|
||||
source: _SourceFixture, settings: Settings, session_factory, tmp_path: Path
|
||||
) -> None:
|
||||
backup_dir = await _snapshot_source(source, settings, session_factory, tmp_path)
|
||||
|
||||
admin_qdrant = AsyncQdrantClient(
|
||||
url=settings.qdrant_url, api_key=settings.qdrant_api_key.get_secret_value(), timeout=300
|
||||
)
|
||||
postgres_restorer = SqlAlchemyPostgresRestorer(admin_database_url=settings.database_url)
|
||||
disposable_database_name = f"restore_gate_{uuid4().hex[:12]}"
|
||||
disposable_collection_prefix = f"restored-{uuid4().hex[:8]}"
|
||||
try:
|
||||
service = RestoreGateService(
|
||||
postgres=postgres_restorer,
|
||||
qdrant=QdrantHttpRestorer(client=admin_qdrant),
|
||||
control_plane=SqlAlchemyControlPlaneReader(),
|
||||
files=LocalFileCopier(),
|
||||
)
|
||||
|
||||
report = await service.run(
|
||||
backup_dir=backup_dir,
|
||||
source_database_url=settings.database_url,
|
||||
disposable_database_name=disposable_database_name,
|
||||
disposable_collection_prefix=disposable_collection_prefix,
|
||||
disposable_blob_root=tmp_path / "restored-blobs",
|
||||
disposable_parsed_root=tmp_path / "restored-parsed",
|
||||
golden_queries=(
|
||||
GoldenQueryProbe(
|
||||
collection_alias=source.alias,
|
||||
point_id=source.chunk_id,
|
||||
query_vector=QUERY_VECTOR,
|
||||
expected_excerpt_checksum=Citation.create(
|
||||
workspace_id=source.workspace_id,
|
||||
space_id=source.granted_space_id,
|
||||
document_id=source.document_id,
|
||||
document_version_id=source.document_version_id,
|
||||
chunk_id=source.chunk_id,
|
||||
parent_chunk_id=None,
|
||||
location={"page": 1},
|
||||
excerpt="The restore gate must reproduce this exact passage.",
|
||||
source_uri="ragcore://source/restore-gate",
|
||||
acquired_at=datetime.now(UTC),
|
||||
retrieval_run_id=uuid4(),
|
||||
index_alias=source.alias,
|
||||
profile_version="restore-gate",
|
||||
).excerpt_checksum,
|
||||
),
|
||||
),
|
||||
security_probe=SecurityDenialProbe(
|
||||
principal_id=source.principal_id,
|
||||
workspace_id=source.workspace_id,
|
||||
granted_space_id=source.granted_space_id,
|
||||
denied_space_id=source.denied_space_id,
|
||||
),
|
||||
)
|
||||
|
||||
assert report.passed
|
||||
assert report.reconciliation[0].matches
|
||||
assert report.golden_queries[0].passed
|
||||
assert report.security_denial is not None
|
||||
assert report.security_denial.passed
|
||||
assert source.granted_space_id in report.security_denial.effective_space_ids
|
||||
assert source.denied_space_id not in report.security_denial.effective_space_ids
|
||||
assert (
|
||||
report.postgres_row_counts["source_spaces"]
|
||||
== report.postgres_row_counts["restored_spaces"]
|
||||
)
|
||||
finally:
|
||||
await postgres_restorer.drop_database(database_name=disposable_database_name)
|
||||
with contextlib.suppress(UnexpectedResponse):
|
||||
await admin_qdrant.delete_collection(
|
||||
f"{disposable_collection_prefix}-{source.collection_name}"
|
||||
)
|
||||
await admin_qdrant.close()
|
||||
|
||||
|
||||
async def test_source_data_is_byte_identical_after_a_restore_gate_run(
|
||||
source: _SourceFixture,
|
||||
settings: Settings,
|
||||
session_factory,
|
||||
qdrant_client: AsyncQdrantClient,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
backup_dir = await _snapshot_source(source, settings, session_factory, tmp_path)
|
||||
|
||||
before_points = await qdrant_client.scroll(
|
||||
collection_name=source.collection_name, limit=100, with_payload=True, with_vectors=True
|
||||
)
|
||||
async with session_factory() as session:
|
||||
before_space = await ControlPlaneRepository(session).get_workspace(source.workspace_id)
|
||||
|
||||
admin_qdrant = AsyncQdrantClient(
|
||||
url=settings.qdrant_url, api_key=settings.qdrant_api_key.get_secret_value(), timeout=300
|
||||
)
|
||||
postgres_restorer = SqlAlchemyPostgresRestorer(admin_database_url=settings.database_url)
|
||||
disposable_database_name = f"restore_gate_{uuid4().hex[:12]}"
|
||||
disposable_collection_prefix = f"restored-{uuid4().hex[:8]}"
|
||||
try:
|
||||
service = RestoreGateService(
|
||||
postgres=postgres_restorer,
|
||||
qdrant=QdrantHttpRestorer(client=admin_qdrant),
|
||||
control_plane=SqlAlchemyControlPlaneReader(),
|
||||
files=LocalFileCopier(),
|
||||
)
|
||||
await service.run(
|
||||
backup_dir=backup_dir,
|
||||
source_database_url=settings.database_url,
|
||||
disposable_database_name=disposable_database_name,
|
||||
disposable_collection_prefix=disposable_collection_prefix,
|
||||
disposable_blob_root=tmp_path / "restored-blobs",
|
||||
disposable_parsed_root=tmp_path / "restored-parsed",
|
||||
golden_queries=(),
|
||||
security_probe=SecurityDenialProbe(
|
||||
principal_id=source.principal_id,
|
||||
workspace_id=source.workspace_id,
|
||||
granted_space_id=source.granted_space_id,
|
||||
denied_space_id=source.denied_space_id,
|
||||
),
|
||||
)
|
||||
finally:
|
||||
await postgres_restorer.drop_database(database_name=disposable_database_name)
|
||||
with contextlib.suppress(UnexpectedResponse):
|
||||
await admin_qdrant.delete_collection(
|
||||
f"{disposable_collection_prefix}-{source.collection_name}"
|
||||
)
|
||||
await admin_qdrant.close()
|
||||
|
||||
after_points = await qdrant_client.scroll(
|
||||
collection_name=source.collection_name, limit=100, with_payload=True, with_vectors=True
|
||||
)
|
||||
async with session_factory() as session:
|
||||
after_space = await ControlPlaneRepository(session).get_workspace(source.workspace_id)
|
||||
|
||||
assert before_points == after_points
|
||||
assert before_space == after_space
|
||||
Reference in new issue
Block a user