Refine runtime migration runner integration

Author: Saravana Kumar Rajendran

docs/final-diff.md

@@ -1,22 +1,22 @@
 # Final Diff
 
 ## Summary
-- Introduced a reusable runtime migration plan so one function governs self-healing schema fixes for the primary database and every organisation database.【F:src/backend/base/langflow/services/database/service.py†L50-L122】【F:src/backend/base/langflow/services/database/runtime_migrations.py†L1-L44】
-- Updated `DatabaseService` to execute the shared plan once per engine, ensuring migrations run on first access without per-query overhead.【F:src/backend/base/langflow/services/database/service.py†L50-L154】
+- Introduced a reusable runtime migration plan so one function governs self-healing schema fixes for the primary database and every organisation database.【F:src/backend/base/langflow/services/database/runtime_migrations.py†L1-L86】
+- Updated `DatabaseService` to invoke the shared plan once per engine, ensuring migrations run on first access without per-query overhead.【F:src/backend/base/langflow/services/database/service.py†L186-L254】
 - Hooked the runtime plan into bootstrap (`initialize_database`) so the base database aligns before Alembic checks execute.【F:src/backend/base/langflow/services/database/utils.py†L16-L57】
-- Refreshed the schema guide to point to the new runtime-migration framework and its tenant coverage.【F:docs/schema-diff.md†L68-L111】
+- Refreshed the schema guide to point to the new runtime-migration framework and its tenant coverage.【F:docs/schema-diff.md†L41-L49】
 
 ## Key Changes
 
 ### `src/backend/base/langflow/services/database/runtime_migrations.py`
-- Added a dedicated module that lists runtime migrations and exposes `apply_runtime_migrations` for the service layer to call. New migrations now land in one place and run everywhere automatically.【F:src/backend/base/langflow/services/database/runtime_migrations.py†L1-L44】
+- Added a dedicated module that lists runtime migrations and exposes helpers to apply them and ensure they run exactly once per engine. New migrations now land in one place and run everywhere automatically.【F:src/backend/base/langflow/services/database/runtime_migrations.py†L1-L86】
 
 ### `src/backend/base/langflow/services/database/service.py`
-- Replaced the ad-hoc `drop_legacy_message_context_id` helper with a general `run_runtime_migrations` gate that wraps execution in an async lock and caches completion per service instance.【F:src/backend/base/langflow/services/database/service.py†L50-L154】
-- Ensured `with_session` blocks on the runtime plan only once per database, covering organisation-specific engines without affecting later queries.【F:src/backend/base/langflow/services/database/service.py†L92-L122】
+- Replaced inline guards with a thin `run_runtime_migrations` wrapper that delegates to the shared runtime-migration framework.【F:src/backend/base/langflow/services/database/service.py†L250-L254】
+- Ensured `with_session` blocks on the runtime plan only once per database, covering organisation-specific engines without affecting later queries.【F:src/backend/base/langflow/services/database/service.py†L186-L199】
 
 ### `src/backend/base/langflow/services/database/utils.py`
 - Adjusted bootstrap to call the new `run_runtime_migrations` entrypoint and improved error messaging around runtime fixes.【F:src/backend/base/langflow/services/database/utils.py†L16-L57】
 
 ### `docs/schema-diff.md`
-- Documented the runtime migration plan so future schema tweaks only require edits to the central function to reach every tenant database.【F:docs/schema-diff.md†L68-L111】
+- Documented the runtime migration plan so future schema tweaks only require edits to the central function to reach every tenant database.【F:docs/schema-diff.md†L41-L49】

docs/schema-diff.md

@@ -40,10 +40,10 @@ All other tables share the same schema definitions across both branches.
 
 ## Deployment guidance for `context_id`
 
-The runtime initialisation path now removes the obsolete `context_id` column automatically before Alembic compares metadata with the live database. During startup `initialize_database` calls `DatabaseService.run_runtime_migrations`, which executes the shared plan in `runtime_migrations.apply_runtime_migrations` to issue `ALTER TABLE "message" DROP COLUMN context_id` whenever the column is still present. In addition, each `DatabaseService` (including organisation-scoped instances) performs the same plan the first time it opens a session, so every tenant database self-heals on first use after deployment.【F:src/backend/base/langflow/services/database/utils.py†L16-L68】【F:src/backend/base/langflow/services/database/service.py†L57-L308】【F:src/backend/base/langflow/services/database/runtime_migrations.py†L1-L44】
+The runtime initialisation path now removes the obsolete `context_id` column automatically before Alembic compares metadata with the live database. During startup `initialize_database` calls `DatabaseService.run_runtime_migrations`, which delegates to `runtime_migrations.ensure_runtime_migrations` to issue `ALTER TABLE "message" DROP COLUMN context_id` whenever the column is still present. In addition, each `DatabaseService` (including organisation-scoped instances) performs the same plan the first time it opens a session, so every tenant database self-heals on first use after deployment.【F:src/backend/base/langflow/services/database/utils.py†L16-L68】【F:src/backend/base/langflow/services/database/service.py†L186-L254】【F:src/backend/base/langflow/services/database/runtime_migrations.py†L1-L86】
 
 - **When deploying the merged branch**: no manual SQL is required; the service will drop the column the first time it sees an upgraded database—both for the primary database at boot and for each organisation database when it is next accessed. This keeps forward compatibility by aligning the schema with the code that no longer references `context_id`.
 - **If you must roll back to the old `main` image**: reintroduce the column with `ALTER TABLE "message" ADD COLUMN context_id TEXT NULL;` before starting the old image, otherwise SQLModel will fail to hydrate the `Message` records because the expected column is missing.
 - **Fresh installations**: databases created after the merge never create the column, so nothing extra runs.
 
-Because the drop happens ahead of Alembic’s `command.check`, the migration runner no longer sees a schema mismatch, ensuring the new container starts cleanly without needing the destructive `--fix` flag.【F:src/backend/base/langflow/services/database/service.py†L300-L386】【F:src/backend/base/langflow/services/database/utils.py†L18-L60】
+Because the drop happens ahead of Alembic’s `command.check`, the migration runner no longer sees a schema mismatch, ensuring the new container starts cleanly without needing the destructive `--fix` flag.【F:src/backend/base/langflow/services/database/service.py†L327-L384】【F:src/backend/base/langflow/services/database/utils.py†L16-L68】

src/backend/base/langflow/services/database/runtime_migrations.py

@@ -8,17 +8,38 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Callable
+from dataclasses import dataclass
 from typing import Final
+from weakref import WeakKeyDictionary
 
 from lfx.log.logger import logger
 from sqlalchemy import inspect
 from sqlalchemy.engine import Connection
+from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlmodel import text
 
 RuntimeMigration = Callable[[Connection], None]
 
 
+@dataclass
+class _RuntimeMigrationState:
+    lock: asyncio.Lock
+    done: bool = False
+
+
+_ENGINE_STATES: "WeakKeyDictionary[AsyncEngine, _RuntimeMigrationState]" = WeakKeyDictionary()
+
+
+def _state_for_engine(engine: AsyncEngine) -> _RuntimeMigrationState:
+    state = _ENGINE_STATES.get(engine)
+    if state is None:
+        state = _RuntimeMigrationState(lock=asyncio.Lock())
+        _ENGINE_STATES[engine] = state
+    return state
+
+
 def _drop_legacy_message_context_id(connection: Connection) -> None:
     inspector = inspect(connection)
 
@@ -50,3 +71,28 @@ def apply_runtime_migrations(connection: Connection) -> None:
 
     for migration in _RUNTIME_MIGRATIONS:
         migration(connection)
+
+
+async def ensure_runtime_migrations(
+    engine: AsyncEngine, *, use_noop_database: bool = False
+) -> None:
+    """Run the runtime migrations exactly once for ``engine``.
+
+    ``use_noop_database`` short-circuits the plan for in-memory/disabled
+    database configurations.
+    """
+
+    state = _state_for_engine(engine)
+
+    if use_noop_database:
+        state.done = True
+        return
+
+    async with state.lock:
+        if state.done:
+            return
+
+        async with engine.begin() as conn:
+            await conn.run_sync(apply_runtime_migrations)
+
+        state.done = True

src/backend/base/langflow/services/database/service.py

@@ -28,7 +28,7 @@
 from langflow.services.database import models
 from langflow.services.database.models.user.crud import get_user_by_username
 from langflow.services.database.session import NoopSession
-from langflow.services.database.runtime_migrations import apply_runtime_migrations
+from langflow.services.database.runtime_migrations import ensure_runtime_migrations
 from langflow.services.database.utils import Result, TableResults
 from langflow.services.deps import get_settings_service
 from langflow.services.utils import teardown_superuser
@@ -49,9 +49,6 @@ def __init__(self, settings_service: SettingsService):
         self.database_url: str = settings_service.settings.database_url
         self._sanitize_database_url()
 
-        self._runtime_migration_lock = asyncio.Lock()
-        self._runtime_migration_done = False
-
         # This file is in langflow.services.database.manager.py
         # the ini is in langflow
         langflow_dir = Path(__file__).parent.parent.parent
@@ -190,10 +187,9 @@ def on_connection(self, dbapi_connection, _connection_record) -> None:
     @asynccontextmanager
     async def with_session(self):
         if self.settings_service.settings.use_noop_database:
-            self._runtime_migration_done = True
             yield NoopSession()
         else:
-            await self._ensure_runtime_migrations()
+            await self.run_runtime_migrations()
             async with AsyncSession(self.engine, expire_on_commit=False) as session:
                 # Start of Selection
                 try:
@@ -253,30 +249,10 @@ async def assign_orphaned_flows_to_superuser(self) -> None:
             await logger.adebug("Successfully assigned orphaned flows to the default superuser")
 
     async def run_runtime_migrations(self) -> None:
-        """Execute the runtime migration plan once for this service."""
-
-        async with self._runtime_migration_lock:
-            if self._runtime_migration_done:
-                return
-
-            if self.settings_service.settings.use_noop_database:
-                self._runtime_migration_done = True
-                return
-
-            async with self.engine.begin() as conn:
-                await conn.run_sync(self._apply_runtime_migrations)
-
-            self._runtime_migration_done = True
-
-    @staticmethod
-    def _apply_runtime_migrations(connection) -> None:
-        apply_runtime_migrations(connection)
-
-    async def _ensure_runtime_migrations(self) -> None:
-        if self._runtime_migration_done:
-            return
-
-        await self.run_runtime_migrations()
+        await ensure_runtime_migrations(
+            self.engine,
+            use_noop_database=self.settings_service.settings.use_noop_database,
+        )
 
     @staticmethod
     def _generate_unique_flow_name(original_name: str, existing_names: set[str]) -> str: