534 lines
16 KiB
Python
534 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
from govoplan_core.core.search import (
|
|
SearchBackfillPage,
|
|
SearchDocument,
|
|
SearchIndexChange,
|
|
SearchQuery,
|
|
SearchResourceType,
|
|
SearchResult,
|
|
)
|
|
from govoplan_core.db.base import Base
|
|
from govoplan_search.backend.db.models import (
|
|
SearchIndexAclToken,
|
|
SearchIndexChangeQueue,
|
|
SearchIndexDocument,
|
|
SearchIndexState,
|
|
)
|
|
from govoplan_search.backend.router import _search_resource_catalogue
|
|
from govoplan_search.backend.service import (
|
|
SearchIndexService,
|
|
aggregate_search_page,
|
|
)
|
|
|
|
|
|
class _Registry:
|
|
def __init__(self, source=None, active_modules=("search", "cases")):
|
|
self.source = source
|
|
self.active_modules = active_modules
|
|
|
|
def manifests(self):
|
|
return tuple(
|
|
SimpleNamespace(id=module_id)
|
|
for module_id in self.active_modules
|
|
)
|
|
|
|
def search_sources(self):
|
|
if self.source is None:
|
|
return ()
|
|
return (
|
|
(
|
|
SimpleNamespace(
|
|
registration=SimpleNamespace(
|
|
id="cases.records",
|
|
order=25,
|
|
)
|
|
),
|
|
self.source,
|
|
),
|
|
)
|
|
|
|
|
|
class _Source:
|
|
def __init__(self, pages=()):
|
|
self.pages = list(pages)
|
|
self.authorized_ids = {"case-1"}
|
|
|
|
def resource_types(self):
|
|
return (
|
|
SearchResourceType(
|
|
provider_id="cases.records",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
label="Cases",
|
|
requires_authorization_recheck=True,
|
|
),
|
|
)
|
|
|
|
def backfill(self, session, *, request):
|
|
del session, request
|
|
return self.pages.pop(0)
|
|
|
|
def authorize(self, session, principal, *, requests):
|
|
del session, principal
|
|
return {
|
|
request.reference.key: (
|
|
request.reference.resource_id
|
|
in self.authorized_ids
|
|
)
|
|
for request in requests
|
|
}
|
|
|
|
|
|
class _ResultProvider:
|
|
def search(self, session, principal, *, query):
|
|
del session, principal
|
|
return tuple(
|
|
SearchResult(
|
|
provider_id="ignored",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id=f"case-{number}",
|
|
title=f"Permit {number}",
|
|
url=f"/cases/case-{number}",
|
|
score=float(10 - number),
|
|
)
|
|
for number in range(1, 7)
|
|
)[: query.limit]
|
|
|
|
|
|
class _AggregateRegistry:
|
|
def search_providers(self):
|
|
return (
|
|
(
|
|
SimpleNamespace(
|
|
registration=SimpleNamespace(id="cases.live")
|
|
),
|
|
_ResultProvider(),
|
|
),
|
|
)
|
|
|
|
|
|
class SearchServiceTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(
|
|
self.engine,
|
|
tables=(
|
|
SearchIndexDocument.__table__,
|
|
SearchIndexAclToken.__table__,
|
|
SearchIndexState.__table__,
|
|
SearchIndexChangeQueue.__table__,
|
|
),
|
|
)
|
|
self.session = Session(self.engine)
|
|
self.service = SearchIndexService(_Registry())
|
|
self.principal = SimpleNamespace(
|
|
tenant_id="tenant-1",
|
|
account_id="account-1",
|
|
membership_id="membership-1",
|
|
identity_id=None,
|
|
group_ids=frozenset({"group-1"}),
|
|
role_ids=frozenset(),
|
|
function_assignment_ids=frozenset(),
|
|
scopes=frozenset({"cases:case:read"}),
|
|
)
|
|
|
|
def tearDown(self) -> None:
|
|
self.session.close()
|
|
self.engine.dispose()
|
|
|
|
def test_index_search_and_acl_filtering(self) -> None:
|
|
self.service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=SearchDocument(
|
|
tenant_id="tenant-1",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id="case-1",
|
|
title="Resident parking permit",
|
|
summary="Monthly permit review",
|
|
url="/cases/case-1",
|
|
acl_tokens=("group:group-1",),
|
|
),
|
|
)
|
|
self.service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=SearchDocument(
|
|
tenant_id="tenant-1",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id="case-2",
|
|
title="Restricted permit",
|
|
url="/cases/case-2",
|
|
acl_tokens=("account:someone-else",),
|
|
),
|
|
)
|
|
self.session.flush()
|
|
|
|
results = self.service.search(
|
|
self.session,
|
|
self.principal,
|
|
query=SearchQuery(text="permit", tenant_id="tenant-1"),
|
|
)
|
|
|
|
self.assertEqual(["case-1"], [result.resource_id for result in results])
|
|
|
|
def test_search_resource_catalogue_exposes_source_filters(self) -> None:
|
|
catalogue = _search_resource_catalogue(_Registry(_Source()))
|
|
|
|
self.assertEqual(1, len(catalogue))
|
|
self.assertEqual("cases", catalogue[0].module_id)
|
|
self.assertEqual("case", catalogue[0].resource_type)
|
|
self.assertEqual("Cases", catalogue[0].label)
|
|
self.assertEqual(25, catalogue[0].order)
|
|
|
|
def test_upsert_replaces_acl_tokens_and_delete_is_idempotent(self) -> None:
|
|
document = SearchDocument(
|
|
tenant_id="tenant-1",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id="case-1",
|
|
title="Permit",
|
|
url="/cases/case-1",
|
|
acl_tokens=("group:group-1",),
|
|
)
|
|
self.service.upsert_document(
|
|
self.session, self.principal, document=document
|
|
)
|
|
self.session.flush()
|
|
self.service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=SearchDocument(
|
|
**{
|
|
**{
|
|
field: getattr(document, field)
|
|
for field in (
|
|
"tenant_id",
|
|
"module_id",
|
|
"resource_type",
|
|
"resource_id",
|
|
"title",
|
|
"url",
|
|
)
|
|
},
|
|
"acl_tokens": ("account:account-1",),
|
|
}
|
|
),
|
|
)
|
|
self.session.flush()
|
|
|
|
self.assertEqual(
|
|
["account:account-1"],
|
|
[
|
|
token.token
|
|
for token in self.session.query(SearchIndexAclToken).all()
|
|
],
|
|
)
|
|
arguments = {
|
|
"tenant_id": "tenant-1",
|
|
"module_id": "cases",
|
|
"resource_type": "case",
|
|
"resource_id": "case-1",
|
|
}
|
|
self.assertTrue(
|
|
self.service.delete_document(
|
|
self.session, self.principal, **arguments
|
|
)
|
|
)
|
|
self.assertFalse(
|
|
self.service.delete_document(
|
|
self.session, self.principal, **arguments
|
|
)
|
|
)
|
|
|
|
def test_index_rejects_cross_tenant_write(self) -> None:
|
|
with self.assertRaises(PermissionError):
|
|
self.service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=SearchDocument(
|
|
tenant_id="tenant-2",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id="case-1",
|
|
title="Permit",
|
|
url="/cases/case-1",
|
|
visibility="tenant",
|
|
),
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "secret field"):
|
|
self.service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=SearchDocument(
|
|
tenant_id="tenant-1",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id="case-secret",
|
|
title="Permit",
|
|
url="/cases/case-secret",
|
|
acl_tokens=("account:account-1",),
|
|
metadata={"access_token": "must-not-index"},
|
|
),
|
|
)
|
|
|
|
def test_query_rejects_cross_tenant_principal(self) -> None:
|
|
with self.assertRaises(PermissionError):
|
|
self.service.search(
|
|
self.session,
|
|
self.principal,
|
|
query=SearchQuery(
|
|
text="permit",
|
|
tenant_id="tenant-2",
|
|
),
|
|
)
|
|
|
|
def test_source_authorization_recheck_is_fail_closed(self) -> None:
|
|
source = _Source()
|
|
service = SearchIndexService(_Registry(source))
|
|
for resource_id in ("case-1", "case-2"):
|
|
service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=_source_document(resource_id),
|
|
)
|
|
self.session.flush()
|
|
|
|
results = service.search(
|
|
self.session,
|
|
self.principal,
|
|
query=SearchQuery(
|
|
text="permit",
|
|
tenant_id="tenant-1",
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
["case-1"],
|
|
[item.resource_id for item in results],
|
|
)
|
|
unavailable_results = SearchIndexService(
|
|
_Registry()
|
|
).search(
|
|
self.session,
|
|
self.principal,
|
|
query=SearchQuery(
|
|
text="permit",
|
|
tenant_id="tenant-1",
|
|
),
|
|
)
|
|
self.assertEqual((), unavailable_results)
|
|
|
|
def test_durable_change_queue_is_idempotent(self) -> None:
|
|
document = _source_document("case-1")
|
|
change = SearchIndexChange(
|
|
change_id="change-1",
|
|
provider_id="cases.records",
|
|
kind="upsert",
|
|
reference=document.reference,
|
|
source_revision=document.source_revision,
|
|
cursor=document.change_cursor or "",
|
|
document=document,
|
|
)
|
|
|
|
self.assertTrue(
|
|
self.service.enqueue_change(
|
|
self.session,
|
|
change=change,
|
|
)
|
|
)
|
|
self.assertFalse(
|
|
self.service.enqueue_change(
|
|
self.session,
|
|
change=change,
|
|
)
|
|
)
|
|
result = self.service.process_changes(self.session)
|
|
self.assertEqual(1, result["applied"])
|
|
self.assertEqual(
|
|
"case-1",
|
|
self.session.query(SearchIndexDocument).one().resource_id,
|
|
)
|
|
self.assertEqual(
|
|
"applied",
|
|
self.session.query(SearchIndexChangeQueue).one().status,
|
|
)
|
|
|
|
delete_change = SearchIndexChange(
|
|
change_id="change-2",
|
|
provider_id="cases.records",
|
|
kind="delete",
|
|
reference=document.reference,
|
|
source_revision="3",
|
|
cursor="cursor-2",
|
|
)
|
|
self.service.enqueue_change(
|
|
self.session,
|
|
change=delete_change,
|
|
)
|
|
self.service.process_changes(self.session)
|
|
self.assertEqual(
|
|
0,
|
|
self.session.query(SearchIndexDocument).count(),
|
|
)
|
|
|
|
def test_rebuild_resumes_and_removes_stale_documents(self) -> None:
|
|
stale = _source_document("stale-case")
|
|
self.service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=stale,
|
|
)
|
|
source = _Source(
|
|
pages=(
|
|
SearchBackfillPage(
|
|
documents=(_source_document("case-1"),),
|
|
next_cursor="page-2",
|
|
complete=False,
|
|
high_watermark="changes-9",
|
|
),
|
|
SearchBackfillPage(
|
|
documents=(_source_document("case-2"),),
|
|
next_cursor=None,
|
|
complete=True,
|
|
high_watermark="changes-9",
|
|
),
|
|
)
|
|
)
|
|
service = SearchIndexService(_Registry(source))
|
|
|
|
started = service.start_rebuild(
|
|
self.session,
|
|
self.principal,
|
|
provider_id="cases.records",
|
|
resource_type="case",
|
|
)
|
|
first_page = service.continue_rebuild(
|
|
self.session,
|
|
self.principal,
|
|
provider_id="cases.records",
|
|
resource_type="case",
|
|
)
|
|
self.assertEqual(started.id, first_page.id)
|
|
self.assertEqual("backfilling", first_page.status)
|
|
self.assertEqual("page-2", first_page.checkpoint_cursor)
|
|
|
|
complete = service.continue_rebuild(
|
|
self.session,
|
|
self.principal,
|
|
provider_id="cases.records",
|
|
resource_type="case",
|
|
)
|
|
self.assertEqual("ready", complete.status)
|
|
self.assertEqual("changes-9", complete.high_watermark)
|
|
self.assertEqual(
|
|
["case-1", "case-2"],
|
|
sorted(
|
|
item.resource_id
|
|
for item in self.session.query(
|
|
SearchIndexDocument
|
|
).all()
|
|
),
|
|
)
|
|
|
|
def test_module_reconciliation_disables_derived_rows(self) -> None:
|
|
self.service.upsert_document(
|
|
self.session,
|
|
self.principal,
|
|
document=SearchDocument(
|
|
tenant_id="tenant-1",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id="case-1",
|
|
title="Permit",
|
|
url="/cases/case-1",
|
|
acl_tokens=("account:account-1",),
|
|
),
|
|
)
|
|
disabled_service = SearchIndexService(
|
|
_Registry(active_modules=("search",))
|
|
)
|
|
result = disabled_service.reconcile_active_modules(
|
|
self.session
|
|
)
|
|
self.assertEqual(1, result["disabled_documents"])
|
|
self.assertFalse(
|
|
self.session.query(SearchIndexDocument).one().active
|
|
)
|
|
|
|
def test_aggregate_cursor_is_stable_and_non_overlapping(self) -> None:
|
|
registry = _AggregateRegistry()
|
|
first = aggregate_search_page(
|
|
registry,
|
|
self.session,
|
|
self.principal,
|
|
query=SearchQuery(
|
|
text="permit",
|
|
tenant_id="tenant-1",
|
|
limit=2,
|
|
),
|
|
)
|
|
second = aggregate_search_page(
|
|
registry,
|
|
self.session,
|
|
self.principal,
|
|
query=SearchQuery(
|
|
text="permit",
|
|
tenant_id="tenant-1",
|
|
limit=2,
|
|
cursor=first.next_cursor,
|
|
),
|
|
)
|
|
|
|
self.assertIsNotNone(first.next_cursor)
|
|
self.assertEqual(
|
|
{"case-1", "case-2"},
|
|
{item.resource_id for item in first.results},
|
|
)
|
|
self.assertEqual(
|
|
{"case-3", "case-4"},
|
|
{item.resource_id for item in second.results},
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "cursor"):
|
|
aggregate_search_page(
|
|
registry,
|
|
self.session,
|
|
self.principal,
|
|
query=SearchQuery(
|
|
text="different",
|
|
tenant_id="tenant-1",
|
|
limit=2,
|
|
cursor=first.next_cursor,
|
|
),
|
|
)
|
|
|
|
|
|
def _source_document(resource_id: str) -> SearchDocument:
|
|
return SearchDocument(
|
|
tenant_id="tenant-1",
|
|
module_id="cases",
|
|
resource_type="case",
|
|
resource_id=resource_id,
|
|
title=f"Permit {resource_id}",
|
|
url=f"/cases/{resource_id}",
|
|
acl_tokens=("account:account-1",),
|
|
provider_id="cases.records",
|
|
source_revision="2",
|
|
change_cursor=f"cursor-{resource_id}",
|
|
requires_authorization_recheck=True,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|