fix: Start reprovisioning search index if a migration requires it (#779)

Author: eikek

bases/renku_data_services/data_api/app.py

@@ -25,6 +25,7 @@
 from renku_data_services.project.blueprints import ProjectsBP, ProjectSessionSecretBP
 from renku_data_services.repositories.blueprints import RepositoriesBP
 from renku_data_services.search.blueprints import SearchBP
+from renku_data_services.search.reprovision import SearchReprovision
 from renku_data_services.session.blueprints import BuildsBP, EnvironmentsBP, SessionLaunchersBP
 from renku_data_services.storage.blueprints import StorageBP, StorageSchemaBP
 from renku_data_services.users.blueprints import KCUsersBP, UserPreferencesBP, UserSecretsBP
@@ -207,12 +208,15 @@ def register_all_handlers(app: Sanic, config: Config) -> Sanic:
         name="search2",
         url_prefix=url_prefix,
         authenticator=config.authenticator,
-        reprovisioning_repo=config.reprovisioning_repo,
-        user_repo=config.kc_user_repo,
-        group_repo=config.group_repo,
+        search_reprovision=SearchReprovision(
+            search_updates_repo=config.search_updates_repo,
+            reprovisioning_repo=config.reprovisioning_repo,
+            solr_config=config.solr_config,
+            user_repo=config.kc_user_repo,
+            group_repo=config.group_repo,
+            project_repo=config.project_repo,
+        ),
         solr_config=config.solr_config,
-        project_repo=config.project_repo,
-        search_updates_repo=config.search_updates_repo,
         authz=config.authz,
     )
     data_connectors = DataConnectorsBP(

bases/renku_data_services/data_api/main.py

@@ -16,8 +16,10 @@
 from sentry_sdk.integrations.sanic import SanicIntegration, _context_enter, _context_exit, _set_transaction
 
 import renku_data_services.search.core as search_core
+import renku_data_services.solr.entity_schema as entity_schema
 from renku_data_services.app_config import Config
 from renku_data_services.authz.admin_sync import sync_admins_from_keycloak
+from renku_data_services.base_models.core import APIUser
 from renku_data_services.data_api.app import register_all_handlers
 from renku_data_services.data_api.prometheus import collect_system_metrics, setup_app_metrics, setup_prometheus
 from renku_data_services.errors.errors import (
@@ -27,7 +29,6 @@
     ValidationError,
 )
 from renku_data_services.migrations.core import run_migrations_for_app
-from renku_data_services.solr import entity_schema
 from renku_data_services.solr.solr_client import DefaultSolrClient
 from renku_data_services.solr.solr_migrate import SchemaMigrator
 from renku_data_services.storage.rclone import RCloneValidator
@@ -72,6 +73,17 @@ async def _update_search(app: Sanic) -> None:
             raise
 
 
+async def _solr_reindex(app: Sanic) -> None:
+    """Run a solr reindex of all data.
+
+    This might be required after migrating the solr schema.
+    """
+    config = Config.from_env()
+    reprovision = config.search_reprovisioning
+    admin = APIUser(is_admin=True)
+    await reprovision.run_reprovision(admin)
+
+
 def send_pending_events(app_name: str) -> None:
     """Send pending messages to redis."""
     app = Sanic(app_name)
@@ -94,6 +106,16 @@ def update_search(app_name: str) -> None:
     asyncio.run(_update_search(app))
 
 
+def solr_reindex(app_name: str) -> None:
+    """Runs a solr reindex."""
+    app = Sanic(app_name)
+    setup_app_metrics(app)
+
+    logger.info("Running SOLR reindex triggered by a migration")
+    asyncio.set_event_loop(uvloop.new_event_loop())
+    asyncio.run(_solr_reindex(app))
+
+
 def create_app() -> Sanic:
     """Create a Sanic application."""
     config = Config.from_env()
@@ -171,11 +193,13 @@ async def do_migrations(_: Sanic) -> None:
         await config.rp_repo.initialize(config.db.conn_url(async_client=False), config.default_resource_pool)
 
     @app.main_process_start
-    async def do_solr_migrations(_: Sanic) -> None:
+    async def do_solr_migrations(app: Sanic) -> None:
         logger.info(f"Running SOLR migrations at: {config.solr_config}")
         migrator = SchemaMigrator(config.solr_config)
         await migrator.ensure_core()
         result = await migrator.migrate(entity_schema.all_migrations)
+        # starting background tasks can only be done in `main_process_ready`
+        app.ctx.solr_reindex = result.requires_reindex
         logger.info(f"SOLR migration done: {result}")
 
     @app.before_server_start
@@ -186,10 +210,11 @@ async def setup_rclone_validator(app: Sanic) -> None:
     @app.main_process_ready
     async def ready(app: Sanic) -> None:
         """Application ready event handler."""
-        config = Config.from_env()
         logger.info("starting events background job.")
-        app.manager.manage("SendEvents", send_pending_events, {"app_name": config.app_name}, transient=True)
-        app.manager.manage("UpdateSearch", update_search, {"app_name": config.app_name}, transient=True)
+        app.manager.manage("SendEvents", send_pending_events, {"app_name": app.name}, transient=True)
+        app.manager.manage("UpdateSearch", update_search, {"app_name": app.name}, transient=True)
+        if getattr(app.ctx, "solr_reindex", False):
+            app.manager.manage("SolrReindex", solr_reindex, {"app_name": app.name}, transient=True)
 
     return app
 

components/renku_data_services/app_config/config.py

@@ -72,6 +72,7 @@
 )
 from renku_data_services.repositories.db import GitRepositoriesRepository
 from renku_data_services.search.db import SearchUpdatesRepo
+from renku_data_services.search.reprovision import SearchReprovision
 from renku_data_services.secrets.db import LowLevelUserSecretsRepo, UserSecretsRepo
 from renku_data_services.session import crs as session_crs
 from renku_data_services.session.db import SessionRepository
@@ -285,6 +286,7 @@ class Config:
     _event_repo: EventRepository | None = field(default=None, repr=False, init=False)
     _reprovisioning_repo: ReprovisioningRepository | None = field(default=None, repr=False, init=False)
     _search_updates_repo: SearchUpdatesRepo | None = field(default=None, repr=False, init=False)
+    _search_reprovisioning: SearchReprovision | None = field(default=None, repr=False, init=False)
     _session_repo: SessionRepository | None = field(default=None, repr=False, init=False)
     _user_preferences_repo: UserPreferencesRepository | None = field(default=None, repr=False, init=False)
     _kc_user_repo: KcUserRepo | None = field(default=None, repr=False, init=False)
@@ -400,6 +402,20 @@ def search_updates_repo(self) -> SearchUpdatesRepo:
             self._search_updates_repo = SearchUpdatesRepo(session_maker=self.db.async_session_maker)
         return self._search_updates_repo
 
+    @property
+    def search_reprovisioning(self) -> SearchReprovision:
+        """The SearchReprovisioning class."""
+        if not self._search_reprovisioning:
+            self._search_reprovisioning = SearchReprovision(
+                search_updates_repo=self.search_updates_repo,
+                reprovisioning_repo=self.reprovisioning_repo,
+                solr_config=self.solr_config,
+                user_repo=self.kc_user_repo,
+                group_repo=self.group_repo,
+                project_repo=self.project_repo,
+            )
+        return self._search_reprovisioning
+
     @property
     def project_repo(self) -> ProjectRepository:
         """The DB adapter for Renku native projects."""

components/renku_data_services/search/blueprints.py

@@ -11,27 +11,19 @@
 from renku_data_services.base_api.auth import authenticate, only_admins
 from renku_data_services.base_api.blueprint import BlueprintFactoryResponse, CustomBlueprint
 from renku_data_services.base_api.misc import validate_query
-from renku_data_services.message_queue.db import ReprovisioningRepository
-from renku_data_services.namespace.db import GroupRepository
-from renku_data_services.project.db import ProjectRepository
 from renku_data_services.search.apispec import SearchQuery
-from renku_data_services.search.db import SearchUpdatesRepo
+from renku_data_services.search.reprovision import SearchReprovision
 from renku_data_services.search.user_query_parser import QueryParser
 from renku_data_services.solr.solr_client import SolrClientConfig
-from renku_data_services.users.db import UserRepo
 
 
 @dataclass(kw_only=True)
 class SearchBP(CustomBlueprint):
     """Handlers for search."""
 
     authenticator: base_models.Authenticator
-    reprovisioning_repo: ReprovisioningRepository
-    user_repo: UserRepo
-    group_repo: GroupRepository
-    project_repo: ProjectRepository
-    search_updates_repo: SearchUpdatesRepo
     solr_config: SolrClientConfig
+    search_reprovision: SearchReprovision
     authz: Authz
 
     def post(self) -> BlueprintFactoryResponse:
@@ -40,19 +32,10 @@ def post(self) -> BlueprintFactoryResponse:
         @authenticate(self.authenticator)
         @only_admins
         async def _post(request: Request, user: base_models.APIUser) -> HTTPResponse | JSONResponse:
-            reprovisioning = await self.reprovisioning_repo.start()
+            reprovisioning = await self.search_reprovision.acquire_reprovision()
 
             request.app.add_task(
-                core.reprovision(
-                    requested_by=user,
-                    reprovisioning=reprovisioning,
-                    search_updates_repo=self.search_updates_repo,
-                    reprovisioning_repo=self.reprovisioning_repo,
-                    solr_config=self.solr_config,
-                    user_repo=self.user_repo,
-                    group_repo=self.group_repo,
-                    project_repo=self.project_repo,
-                ),
+                self.search_reprovision.init_reprovision(requested_by=user, reprovisioning=reprovisioning),
                 name=f"reprovisioning-{reprovisioning.id}",
             )
 
@@ -65,7 +48,7 @@ def get_status(self) -> BlueprintFactoryResponse:
 
         @authenticate(self.authenticator)
         async def _get_status(_: Request, __: base_models.APIUser) -> JSONResponse | HTTPResponse:
-            reprovisioning = await self.reprovisioning_repo.get_active_reprovisioning()
+            reprovisioning = await self.search_reprovision.get_current_reprovision()
             if not reprovisioning:
                 return HTTPResponse(status=404)
             return json({"id": str(reprovisioning.id), "start_date": reprovisioning.start_date.isoformat()})
@@ -78,7 +61,7 @@ def delete(self) -> BlueprintFactoryResponse:
         @authenticate(self.authenticator)
         @only_admins
         async def _delete(_: Request, __: base_models.APIUser) -> HTTPResponse:
-            await self.reprovisioning_repo.stop()
+            await self.search_reprovision.kill_reprovision_lock()
             return HTTPResponse(status=204)
 
         return "/search/reprovision", ["DELETE"], _delete

components/renku_data_services/search/core.py

@@ -9,10 +9,6 @@
 import renku_data_services.search.solr_token as st
 from renku_data_services.authz.models import Role
 from renku_data_services.base_models import APIUser
-from renku_data_services.message_queue.db import ReprovisioningRepository
-from renku_data_services.message_queue.models import Reprovisioning
-from renku_data_services.namespace.db import GroupRepository
-from renku_data_services.project.db import ProjectRepository
 from renku_data_services.search import authz, converters
 from renku_data_services.search.db import SearchUpdatesRepo
 from renku_data_services.search.models import DeleteDoc
@@ -37,57 +33,6 @@
     SolrQuery,
     SubQuery,
 )
-from renku_data_services.users.db import UserRepo
-
-
-async def reprovision(
-    requested_by: APIUser,
-    reprovisioning: Reprovisioning,
-    search_updates_repo: SearchUpdatesRepo,
-    reprovisioning_repo: ReprovisioningRepository,
-    solr_config: SolrClientConfig,
-    user_repo: UserRepo,
-    group_repo: GroupRepository,
-    project_repo: ProjectRepository,
-) -> None:
-    """Initiates reprovisioning by inserting documents into the staging table."""
-
-    def log_counter(c: int) -> None:
-        if c % 50 == 0:
-            logger.info(f"Inserted {c}. entities into staging table...")
-
-    try:
-        logger.info(f"Starting reprovisioning with ID {reprovisioning.id}")
-        started = datetime.now()
-        await search_updates_repo.clear_all()
-        async with DefaultSolrClient(solr_config) as client:
-            await client.delete("_type:*")
-        counter = 0
-        all_users = user_repo.get_all_users(requested_by=requested_by)
-        async for user_entity in all_users:
-            await search_updates_repo.insert(user_entity, started)
-            counter += 1
-            log_counter(counter)
-
-        all_groups = group_repo.get_all_groups(requested_by=requested_by)
-        async for group_entity in all_groups:
-            await search_updates_repo.insert(group_entity, started)
-            counter += 1
-            log_counter(counter)
-
-        all_projects = project_repo.get_all_projects(requested_by=requested_by)
-        async for project_entity in all_projects:
-            await search_updates_repo.insert(project_entity, started)
-            counter += 1
-            log_counter(counter)
-
-        logger.info(f"Inserted {counter} entities into the staging table.")
-
-    except Exception as e:
-        logger.error("Error while reprovisioning entities!", exc_info=e)
-        ## TODO error handling. skip or fail?
-    finally:
-        await reprovisioning_repo.stop()
 
 
 async def update_solr(search_updates_repo: SearchUpdatesRepo, solr_client: SolrClient, batch_size: int) -> None: