feat(framework): Enable SqlMixin to automigrate database

framework/py/flwr/server/superlink/linkstate/linkstate_factory.py

@@ -56,20 +56,28 @@ def __init__(
 
     def state(self) -> LinkState:
         """Return a State instance and create it, if necessary."""
+        # Return cached state if it exists
+        if self.state_instance is not None:
+            if self.database == FLWR_IN_MEMORY_DB_NAME:
+                log(DEBUG, "Using InMemoryState")
+            else:
+                log(DEBUG, "Using SqlLinkState")
+            return self.state_instance
+
         # Get the ObjectStore instance
         object_store = self.objectstore_factory.store()
 
         # InMemoryState
         if self.database == FLWR_IN_MEMORY_DB_NAME:
-            if self.state_instance is None:
-                self.state_instance = InMemoryLinkState(
-                    self.federation_manager, object_store
-                )
+            self.state_instance = InMemoryLinkState(
+                self.federation_manager, object_store
+            )
             log(DEBUG, "Using InMemoryState")
             return self.state_instance
 
         # SqlLinkState
         state = SqlLinkState(self.database, self.federation_manager, object_store)
         state.initialize()
+        self.state_instance = state
         log(DEBUG, "Using SqlLinkState")
         return state

framework/py/flwr/server/superlink/linkstate/linkstate_test.py

@@ -1823,17 +1823,6 @@ def state_factory(self) -> SqlLinkState:
         state.initialize()
         return state
 
-    def test_initialize(self) -> None:
-        """Test initialization."""
-        # Prepare
-        state = self.state_factory()
-
-        # Execute
-        result = state.query("SELECT name FROM sqlite_schema;")
-
-        # Assert - 7 tables + 11 indexes (3 explicit + 8 auto from UNIQUE constraints)
-        assert len(result) == 18
-
 
 class SqlFileBasedTest(StateTest, unittest.TestCase):
     """Test SqlLinkState implementation with file-based database."""
@@ -1852,17 +1841,6 @@ def state_factory(self) -> SqlLinkState:
         state.initialize()
         return state
 
-    def test_initialize(self) -> None:
-        """Test initialization."""
-        # Prepare
-        state = self.state_factory()
-
-        # Execute
-        result = state.query("SELECT name FROM sqlite_schema;")
-
-        # Assert - 7 tables + 11 indexes (3 explicit + 8 auto from UNIQUE constraints)
-        assert len(result) == 18
-
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

framework/py/flwr/supercore/object_store/object_store_factory.py

@@ -49,15 +49,23 @@ def store(self) -> ObjectStore:
         ObjectStore
             An ObjectStore instance for storing objects by object_id.
         """
+        # Return cached store if it exists
+        if self.store_instance is not None:
+            if self.database == FLWR_IN_MEMORY_DB_NAME:
+                log(DEBUG, "Using InMemoryObjectStore")
+            else:
+                log(DEBUG, "Using SqlObjectStore")
+            return self.store_instance
+
         # InMemoryObjectStore
         if self.database == FLWR_IN_MEMORY_DB_NAME:
-            if self.store_instance is None:
-                self.store_instance = InMemoryObjectStore()
+            self.store_instance = InMemoryObjectStore()
             log(DEBUG, "Using InMemoryObjectStore")
             return self.store_instance
 
         # SqlObjectStore
         store = SqlObjectStore(self.database)
         store.initialize()
+        self.store_instance = store
         log(DEBUG, "Using SqlObjectStore")
         return store

framework/py/flwr/supercore/object_store/object_store_test.py

@@ -18,8 +18,10 @@
 import tempfile
 import unittest
 from abc import abstractmethod
+from typing import cast
 
 from parameterized import parameterized
+from sqlalchemy import Engine, inspect
 
 from flwr.common.inflatable import get_object_id, get_object_tree, iterate_object_tree
 from flwr.common.inflatable_test import CustomDataClass
@@ -377,12 +379,20 @@ class SqlInMemoryObjectStoreTest(ObjectStoreTest):
 
     __test__ = True
 
-    def object_store_factory(self) -> ObjectStore:
+    def object_store_factory(self) -> SqlObjectStore:
         """Return SqlObjectStore."""
         store = SqlObjectStore(":memory:")
         store.initialize()
         return store
 
+    def test_in_memory_does_not_create_alembic_version(self) -> None:
+        """Ensure in-memory DB uses create_all without Alembic versioning."""
+        store = self.object_store_factory()
+        table_names = inspect(
+            cast(Engine, store._engine)  # pylint: disable=W0212
+        ).get_table_names()
+        self.assertNotIn("alembic_version", table_names)
+
 
 class SqlFileBasedObjectStoreTest(ObjectStoreTest):
     """Test SqlObjectStore implementation with file-based database."""
@@ -399,8 +409,16 @@ def tearDown(self) -> None:
         super().tearDown()
         self.temp_file.close()
 
-    def object_store_factory(self) -> ObjectStore:
+    def object_store_factory(self) -> SqlObjectStore:
         """Return SqlObjectStore."""
         store = SqlObjectStore(self.temp_file.name)
         store.initialize()
         return store
+
+    def test_file_db_creates_alembic_version(self) -> None:
+        """Ensure file-based DBs run Alembic migrations."""
+        store = self.object_store_factory()
+        table_names = inspect(
+            cast(Engine, store._engine)  # pylint: disable=W0212
+        ).get_table_names()
+        self.assertIn("alembic_version", table_names)

framework/py/flwr/supercore/sql_mixin.py

@@ -30,6 +30,7 @@
 
 from flwr.common.logger import log
 from flwr.supercore.constant import FLWR_IN_MEMORY_SQLITE_DB_URL, SQLITE_PRAGMAS
+from flwr.supercore.state.alembic.utils import run_migrations
 
 _current_session: ContextVar[Session | None] = ContextVar(
     "current_sqlalchemy_session",
@@ -193,9 +194,14 @@ def initialize(self, log_queries: bool = False) -> list[str]:
         # Create session factory
         self._session_factory = sessionmaker(bind=self._engine)
 
-        # Create tables defined in metadata (idempotent - only creates missing tables)
-        if (metadata := self.get_metadata()) is not None:
+        # Create database
+        metadata: MetaData | None = self.get_metadata()
+        if metadata and self.database_url == FLWR_IN_MEMORY_SQLITE_DB_URL:
+            # In-memory databases: create tables directly from SQLAlchemy metadata
             metadata.create_all(self._engine)
+        else:
+            # File-based databases: use Alembic migrations for schema versioning
+            run_migrations(self._engine)
 
         # Get all table names using inspector
         inspector = inspect(self._engine)