from __future__ import annotations import os from types import SimpleNamespace import unittest import uuid from sqlalchemy import create_engine, delete, inspect, select from sqlalchemy.orm import Session from govoplan_core.core.search import SearchDocument, SearchQuery from govoplan_search.backend.db.models import ( SearchIndexAclToken, SearchIndexChangeQueue, SearchIndexDocument, SearchIndexState, ) from govoplan_search.backend.service import SearchIndexService DATABASE_URL = os.environ.get( "GOVOPLAN_SEARCH_POSTGRES_URL", "", ).strip() class _Registry: def manifests(self): return ( SimpleNamespace(id="search"), SimpleNamespace(id="cases"), ) def search_sources(self): return () @unittest.skipUnless( DATABASE_URL.startswith("postgresql"), "GOVOPLAN_SEARCH_POSTGRES_URL is not configured", ) class PostgreSqlSearchTests(unittest.TestCase): @classmethod def setUpClass(cls) -> None: cls.engine = create_engine(DATABASE_URL) table_names = set(inspect(cls.engine).get_table_names()) required = { SearchIndexDocument.__tablename__, SearchIndexAclToken.__tablename__, SearchIndexState.__tablename__, SearchIndexChangeQueue.__tablename__, } if not required.issubset(table_names): raise RuntimeError( "Search migrations were not applied before the " "PostgreSQL integration test." ) @classmethod def tearDownClass(cls) -> None: cls.engine.dispose() def setUp(self) -> None: self.tenant_id = str(uuid.uuid4()) self.other_tenant_id = str(uuid.uuid4()) self.session = Session(self.engine) self.service = SearchIndexService(_Registry()) self.principal = SimpleNamespace( tenant_id=self.tenant_id, account_id="account-allowed", membership_id="membership-allowed", identity_id=None, group_ids=frozenset(), role_ids=frozenset(), function_assignment_ids=frozenset(), scopes=frozenset({"cases:case:read"}), ) def tearDown(self) -> None: document_ids = select(SearchIndexDocument.id).where( SearchIndexDocument.tenant_id.in_( (self.tenant_id, self.other_tenant_id) ) ) self.session.execute( delete(SearchIndexAclToken).where( SearchIndexAclToken.document_id.in_(document_ids) ) ) self.session.execute( delete(SearchIndexChangeQueue).where( SearchIndexChangeQueue.tenant_id.in_( (self.tenant_id, self.other_tenant_id) ) ) ) self.session.execute( delete(SearchIndexState).where( SearchIndexState.tenant_id.in_( (self.tenant_id, self.other_tenant_id) ) ) ) self.session.execute( delete(SearchIndexDocument).where( SearchIndexDocument.tenant_id.in_( (self.tenant_id, self.other_tenant_id) ) ) ) self.session.commit() self.session.close() def test_postgres_full_text_query_is_tenant_and_acl_bounded( self, ) -> None: self.service.upsert_document( self.session, self.principal, document=SearchDocument( tenant_id=self.tenant_id, module_id="cases", provider_id="cases.records", resource_type="case", resource_id="case-visible", title="Monthly parking permit review", url="/cases/case-visible", body="Compare the permit data with the payment register.", visibility="tenant", ), ) self.service.upsert_document( self.session, self.principal, document=SearchDocument( tenant_id=self.tenant_id, module_id="cases", provider_id="cases.records", resource_type="case", resource_id="case-denied", title="Restricted parking permit", url="/cases/case-denied", visibility="restricted", acl_tokens=("account:someone-else",), ), ) other_principal = SimpleNamespace( **{ **vars(self.principal), "tenant_id": self.other_tenant_id, } ) self.service.upsert_document( self.session, other_principal, document=SearchDocument( tenant_id=self.other_tenant_id, module_id="cases", provider_id="cases.records", resource_type="case", resource_id="case-other-tenant", title="Other tenant parking permit", url="/cases/case-other-tenant", visibility="tenant", ), ) self.session.commit() page = self.service.search_page( self.session, self.principal, query=SearchQuery( text="parking permit", tenant_id=self.tenant_id, ), ) self.assertEqual( ("case-visible",), tuple(item.resource_id for item in page.results), ) self.assertEqual( "postgresql", self.service.diagnostics( self.session, self.principal, )["backend"], ) if __name__ == "__main__": unittest.main()