move index query to models, tests

Author: monneyboi

poliloom/poliloom/cli.py

@@ -14,7 +14,7 @@
 from poliloom.database import get_engine
 from poliloom.logging import setup_logging
 from sqlalchemy.orm import Session
-from sqlalchemy import exists, text
+from sqlalchemy import exists, func, select
 from poliloom.models import (
     Country,
     CurrentImportEntity,
@@ -899,24 +899,6 @@ def clean_properties(dry_run):
             raise SystemExit(1)
 
 
-def _get_search_indexed_models():
-    """Get all models that use WikidataEntityMixin (which provides search indexing)."""
-    from poliloom.models.wikidata import WikidataEntityMixin
-    import poliloom.models as models_module
-
-    search_indexed_models = []
-    for name in dir(models_module):
-        obj = getattr(models_module, name)
-        if (
-            isinstance(obj, type)
-            and issubclass(obj, WikidataEntityMixin)
-            and obj is not WikidataEntityMixin
-            and hasattr(obj, "__tablename__")
-        ):
-            search_indexed_models.append(obj)
-    return search_indexed_models
-
-
 @main.command("index-create")
 def index_create():
     """Create the Meilisearch entities index.
@@ -1001,48 +983,16 @@ def index_build(batch_size, rebuild):
         if search_service.ensure_index():
             click.echo(f"   Created index '{INDEX_NAME}'")
 
-    # Get all searchable models
-    models = _get_search_indexed_models()
-    click.echo(f"   Types: {', '.join(m.__name__ for m in models)}")
-
-    # Build dynamic SQL query
-    # LEFT JOIN each model table and build types array from which ones match
-    left_joins = []
-    case_statements = []
-    group_by_columns = ["we.wikidata_id"]
-
-    for model in models:
-        table_name = model.__tablename__
-        left_joins.append(
-            f"LEFT JOIN {table_name} ON we.wikidata_id = {table_name}.wikidata_id"
-        )
-        case_statements.append(
-            f"CASE WHEN {table_name}.wikidata_id IS NOT NULL THEN '{model.__name__}' END"
-        )
-        group_by_columns.append(f"{table_name}.wikidata_id")
-
-    array_expr = f"array_remove(ARRAY[{', '.join(case_statements)}], NULL)"
-
-    base_sql = f"""
-        SELECT
-            we.wikidata_id,
-            array_agg(DISTINCT wel.label) as labels,
-            {array_expr} as types
-        FROM wikidata_entities we
-        JOIN wikidata_entity_labels wel ON we.wikidata_id = wel.entity_id
-        {chr(10).join(left_joins)}
-        WHERE we.deleted_at IS NULL
-        GROUP BY {", ".join(group_by_columns)}
-        HAVING array_length({array_expr}, 1) > 0
-    """
+    # Build query for search index documents
+    query = WikidataEntity.search_index_query()
 
     total_indexed = 0
     task_uids = []
 
     with Session(get_engine()) as session:
         # Count total
-        count_result = session.execute(text(f"SELECT COUNT(*) FROM ({base_sql}) subq"))
-        total = count_result.scalar()
+        count_query = select(func.count()).select_from(query.subquery())
+        total = session.execute(count_query).scalar()
 
         if total == 0:
             click.echo("   No entities to index")
@@ -1053,11 +1003,8 @@ def index_build(batch_size, rebuild):
         # Process in batches
         offset_val = 0
         while offset_val < total:
-            paginated_sql = f"{base_sql} OFFSET :offset LIMIT :limit"
-            rows = session.execute(
-                text(paginated_sql),
-                {"offset": offset_val, "limit": batch_size},
-            ).fetchall()
+            paginated_query = query.offset(offset_val).limit(batch_size)
+            rows = session.execute(paginated_query).fetchall()
 
             if not rows:
                 break

poliloom/poliloom/models/wikidata.py

@@ -19,10 +19,12 @@
     delete,
     exists,
     func,
+    literal,
     literal_column,
     or_,
     select,
     text,
+    union_all,
     update,
 )
 from sqlalchemy.dialects.postgresql import UUID
@@ -623,6 +625,51 @@ def cleanup_orphaned(cls, session: Session) -> int:
 
         return result.rowcount
 
+    @classmethod
+    def search_index_query(cls):
+        """Build query for search index documents.
+
+        Creates a query that returns all searchable entities with their
+        aggregated types and labels. Only includes non-deleted entities.
+
+        Returns:
+            SQLAlchemy select query with columns: wikidata_id, types, labels
+        """
+        models = WikidataEntityMixin.__subclasses__()
+
+        # Build UNION of all model tables with their type names
+        entity_unions = union_all(
+            *[
+                select(
+                    model.wikidata_id.label("wikidata_id"),
+                    literal(model.__name__).label("type"),
+                )
+                for model in models
+            ]
+        ).subquery("entity_types")
+
+        # Main query: aggregate types and labels per entity
+        return (
+            select(
+                entity_unions.c.wikidata_id,
+                func.array_agg(func.distinct(entity_unions.c.type)).label("types"),
+                func.array_agg(func.distinct(WikidataEntityLabel.label)).label(
+                    "labels"
+                ),
+            )
+            .select_from(entity_unions)
+            .join(
+                WikidataEntityLabel,
+                entity_unions.c.wikidata_id == WikidataEntityLabel.entity_id,
+            )
+            .join(
+                cls,
+                entity_unions.c.wikidata_id == cls.wikidata_id,
+            )
+            .where(cls.deleted_at.is_(None))
+            .group_by(entity_unions.c.wikidata_id)
+        )
+
 
 class WikidataEntityLabel(Base, TimestampMixin, UpsertMixin):
     """Normalized label storage for wikidata entities."""

poliloom/tests/models/test_wikidata.py

@@ -987,3 +987,175 @@ def test_dry_run_does_not_call_delete_documents(self, db_session):
 
         # Dry run reports what would be removed
         assert stats["entities_removed"] == 1
+
+
+class TestSearchIndexQuery:
+    """Test WikidataEntity.search_index_query functionality."""
+
+    def _create_entity_with_labels(self, db_session, wikidata_id, name, labels):
+        """Helper to create a WikidataEntity with labels."""
+        from poliloom.models import WikidataEntityLabel
+
+        stmt = insert(WikidataEntity).values(
+            [{"wikidata_id": wikidata_id, "name": name}]
+        )
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+
+        for label in labels:
+            stmt = insert(WikidataEntityLabel).values(
+                [{"entity_id": wikidata_id, "label": label}]
+            )
+            stmt = stmt.on_conflict_do_nothing(index_elements=["entity_id", "label"])
+            db_session.execute(stmt)
+
+        db_session.flush()
+
+    def _create_location(self, db_session, wikidata_id, name, labels):
+        """Helper to create a Location entity with labels."""
+        from poliloom.models import Location
+
+        self._create_entity_with_labels(db_session, wikidata_id, name, labels)
+
+        stmt = insert(Location.__table__).values([{"wikidata_id": wikidata_id}])
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+        db_session.flush()
+
+    def _create_country(self, db_session, wikidata_id, name, labels):
+        """Helper to create a Country entity with labels."""
+        from poliloom.models import Country
+
+        self._create_entity_with_labels(db_session, wikidata_id, name, labels)
+
+        stmt = insert(Country.__table__).values([{"wikidata_id": wikidata_id}])
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+        db_session.flush()
+
+    def _create_position(self, db_session, wikidata_id, name, labels):
+        """Helper to create a Position entity with labels."""
+        from poliloom.models import Position
+
+        self._create_entity_with_labels(db_session, wikidata_id, name, labels)
+
+        stmt = insert(Position.__table__).values([{"wikidata_id": wikidata_id}])
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+        db_session.flush()
+
+    def test_returns_entity_with_single_type(self, db_session):
+        """Test query returns entity with single type."""
+        self._create_location(db_session, "Q60", "New York City", ["New York", "NYC"])
+
+        query = WikidataEntity.search_index_query()
+        results = db_session.execute(query).fetchall()
+
+        assert len(results) == 1
+        row = results[0]
+        assert row.wikidata_id == "Q60"
+        assert "Location" in row.types
+        assert set(row.labels) == {"New York", "NYC"}
+
+    def test_returns_entity_with_multiple_types(self, db_session):
+        """Test query aggregates multiple types for same entity."""
+        from poliloom.models import Location, Country
+
+        # Germany is both a Location and a Country
+        self._create_entity_with_labels(
+            db_session, "Q183", "Germany", ["Germany", "Deutschland"]
+        )
+
+        # Add to both tables
+        stmt = insert(Location.__table__).values([{"wikidata_id": "Q183"}])
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+
+        stmt = insert(Country.__table__).values([{"wikidata_id": "Q183"}])
+        stmt = stmt.on_conflict_do_nothing(index_elements=["wikidata_id"])
+        db_session.execute(stmt)
+        db_session.flush()
+
+        query = WikidataEntity.search_index_query()
+        results = db_session.execute(query).fetchall()
+
+        assert len(results) == 1
+        row = results[0]
+        assert row.wikidata_id == "Q183"
+        assert "Location" in row.types
+        assert "Country" in row.types
+        assert set(row.labels) == {"Germany", "Deutschland"}
+
+    def test_returns_multiple_entities(self, db_session):
+        """Test query returns multiple entities."""
+        self._create_location(db_session, "Q60", "New York City", ["NYC"])
+        self._create_position(db_session, "Q30185", "Mayor", ["Mayor", "Bürgermeister"])
+
+        query = WikidataEntity.search_index_query()
+        results = db_session.execute(query).fetchall()
+
+        assert len(results) == 2
+        results_by_id = {r.wikidata_id: r for r in results}
+
+        assert "Location" in results_by_id["Q60"].types
+        assert set(results_by_id["Q60"].labels) == {"NYC"}
+
+        assert "Position" in results_by_id["Q30185"].types
+        assert set(results_by_id["Q30185"].labels) == {"Mayor", "Bürgermeister"}
+
+    def test_excludes_soft_deleted_entities(self, db_session):
+        """Test query excludes soft-deleted entities."""
+        from datetime import datetime, timezone
+
+        self._create_location(db_session, "Q60", "New York City", ["NYC"])
+        self._create_location(db_session, "Q84", "London", ["London"])
+
+        # Soft-delete London
+        db_session.execute(
+            WikidataEntity.__table__.update()
+            .where(WikidataEntity.wikidata_id == "Q84")
+            .values(deleted_at=datetime.now(timezone.utc))
+        )
+        db_session.flush()
+
+        query = WikidataEntity.search_index_query()
+        results = db_session.execute(query).fetchall()
+
+        assert len(results) == 1
+        assert results[0].wikidata_id == "Q60"
+
+    def test_excludes_entities_not_in_model_tables(self, db_session):
+        """Test query only returns entities that exist in model tables."""
+        # Create entity with labels but NOT in any model table
+        self._create_entity_with_labels(db_session, "Q999", "Orphan Entity", ["Orphan"])
+
+        # Create proper location
+        self._create_location(db_session, "Q60", "New York City", ["NYC"])
+
+        query = WikidataEntity.search_index_query()
+        results = db_session.execute(query).fetchall()
+
+        # Only the location should be returned
+        assert len(results) == 1
+        assert results[0].wikidata_id == "Q60"
+
+    def test_pagination_with_offset_and_limit(self, db_session):
+        """Test query supports pagination with offset and limit."""
+        # Create multiple locations
+        for i in range(5):
+            self._create_location(db_session, f"Q{i}", f"Location {i}", [f"Label {i}"])
+
+        query = WikidataEntity.search_index_query()
+
+        # Get first 2
+        results_page1 = db_session.execute(query.limit(2)).fetchall()
+        assert len(results_page1) == 2
+
+        # Get next 2
+        results_page2 = db_session.execute(query.offset(2).limit(2)).fetchall()
+        assert len(results_page2) == 2
+
+        # Verify no overlap
+        page1_ids = {r.wikidata_id for r in results_page1}
+        page2_ids = {r.wikidata_id for r in results_page2}
+        assert page1_ids.isdisjoint(page2_ids)