flwr 1.25.0__py3-none-any.whl → 1.26.0__py3-none-any.whl

flwr/supercore/state/alembic/utils.py

@@ -0,0 +1,239 @@
+# Copyright 2026 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Helpers for running and validating Alembic migrations."""
+
+
+from logging import INFO
+from pathlib import Path
+
+from alembic import command
+from alembic.config import Config
+from sqlalchemy import MetaData, create_engine, inspect, pool
+from sqlalchemy.engine import Engine
+
+from flwr.common.exit import ExitCode, flwr_exit
+from flwr.common.logger import log
+from flwr.supercore.state.schema.corestate_tables import create_corestate_metadata
+from flwr.supercore.state.schema.linkstate_tables import create_linkstate_metadata
+from flwr.supercore.state.schema.objectstore_tables import create_objectstore_metadata
+
+ALEMBIC_DIR = Path(__file__).resolve().parent
+ALEMBIC_VERSION_TABLE = "alembic_version"
+FLWR_STATE_BASELINE_REVISION = "8e65d8ae60b0"
+
+
+def get_combined_metadata() -> MetaData:
+    """Combine all Flower state metadata objects into a single MetaData instance.
+
+    This ensures Alembic can track all tables across CoreState, LinkState, and
+    ObjectStore.
+
+    Returns
+    -------
+    MetaData
+        Combined SQLAlchemy MetaData with all Flower state tables.
+    """
+    # Start with linkstate tables
+    metadata = create_linkstate_metadata()
+
+    # Add corestate tables
+    corestate_metadata = create_corestate_metadata()
+    for table in corestate_metadata.tables.values():
+        table.to_metadata(metadata)
+
+    # Add objectstore tables
+    objectstore_metadata = create_objectstore_metadata()
+    for table in objectstore_metadata.tables.values():
+        table.to_metadata(metadata)
+
+    return metadata
+
+
+def run_migrations(engine: Engine) -> None:
+    """Run pending Alembic migrations, handling pre-Alembic legacy databases.
+
+    Expected scenarios:
+    - If alembic_version exists: run migrations normally.
+    - If DB is empty, e.g. when newly created: run migrations normally.
+    - If DB is pre-Alembic and schema is mismatched: exit with guidance.
+    - If DB is pre-Alembic and schema matches baseline: stamp, then upgrade.
+    """
+    config = build_alembic_config(engine)
+    has_version_table = _has_alembic_version_table(engine)
+
+    # Standard database with version tracking: just upgrade.
+    if has_version_table:
+        command.upgrade(config, "head")
+        return
+
+    table_names = _get_user_table_names(engine)
+
+    # Empty/new database: run all migrations from scratch.
+    if not table_names:
+        command.upgrade(config, "head")
+        return
+
+    # Pre-Alembic database detected without version tracking: verify database matches
+    # baseline schema before stamping version and upgrading.
+    is_valid, error_msg = _verify_legacy_schema_matches_baseline(engine)
+
+    # This is an edge case and unlikely to happen since SuperLink requires a specific
+    # schema to operate normally.
+    if not is_valid:
+        flwr_exit(
+            ExitCode.SUPERLINK_DATABASE_SCHEMA_MISMATCH,
+            "Detected a pre-Alembic Flower state database, but its schema does not "
+            f"match the baseline migration (revision {FLWR_STATE_BASELINE_REVISION}). "
+            "Back up the database and either migrate it manually to the baseline "
+            "schema or start with a fresh database. "
+            f"{error_msg}",
+        )
+
+    log(
+        INFO,
+        "Detected pre-Alembic state database without alembic_version; stamping to %s "
+        "before upgrading.",
+        FLWR_STATE_BASELINE_REVISION,
+    )
+    stamp_existing_database(engine, FLWR_STATE_BASELINE_REVISION)
+    command.upgrade(config, "head")
+    log(INFO, "Flower state database stamped and upgraded successfully!")
+
+
+def build_alembic_config(engine: Engine) -> Config:
+    """Create Alembic config with script location and DB URL."""
+    config = Config()
+    config.set_main_option("script_location", str(ALEMBIC_DIR))
+    config.set_main_option("sqlalchemy.url", str(engine.url))
+    return config
+
+
+def _get_user_table_names(engine: Engine) -> set[str]:
+    """Return non-internal table names for the given engine."""
+    inspector = inspect(engine)
+    table_names = set(inspector.get_table_names())
+    # Exclude SQLite internal tables (for example, sqlite_sequence)
+    return {name for name in table_names if not name.startswith("sqlite_")}
+
+
+def _has_alembic_version_table(engine: Engine) -> bool:
+    """Return True if the Alembic version table exists."""
+    inspector = inspect(engine)
+    return inspector.has_table(ALEMBIC_VERSION_TABLE)
+
+
+def _get_baseline_metadata() -> MetaData:
+    """Create an in-memory DB at baseline revision and reflect its schema.
+
+    Uses an in-memory SQLite database instead of a temporary file to avoid requiring
+    filesystem write access. Note that this function is only invoked for pre-Alembic
+    databases.
+
+    The implementation uses StaticPool and passes an active connection via
+    config.attributes to Alembic's env.py. This ensures the same in-memory database
+    instance is used throughout migration and reflection, since each new connection to
+    sqlite:///:memory: creates a separate empty database.
+    """
+    # Create an in-memory SQLite database with StaticPool to ensure connection reuse.
+    # This is needed because in-memory databases are instance-specific per connection.
+    engine = create_engine(
+        "sqlite:///:memory:",
+        connect_args={"check_same_thread": False},
+        poolclass=pool.StaticPool,
+    )
+    try:
+        # Open a connection and pass it to Alembic to ensure the in-memory database
+        # persists throughout the migration process. Without this, Alembic would
+        # create a new connection (and thus a new empty database) from the URL.
+        with engine.begin() as connection:
+            config = build_alembic_config(engine)
+            # Store the connection in config.attributes so env.py can use it directly.
+            # This prevents Alembic from creating a new connection and losing our data.
+            config.attributes["connection"] = connection
+            command.upgrade(config, FLWR_STATE_BASELINE_REVISION)
+
+        # Reflect the baseline schema from the in-memory database.
+        # At this point, the StaticPool ensures we're still connected to the same
+        # database instance that contains the migrated tables.
+        baseline_metadata = MetaData()
+        baseline_metadata.reflect(
+            bind=engine,
+            only=lambda table_name, _: table_name != ALEMBIC_VERSION_TABLE,
+        )
+    finally:
+        engine.dispose()
+
+    return baseline_metadata
+
+
+def _verify_legacy_schema_matches_baseline(engine: Engine) -> tuple[bool, str]:
+    """Verify legacy schema matches baseline tables and columns.
+
+    Only missing tables/columns are reported as errors.
+
+    Returns
+    -------
+    tuple[bool, str]
+        (is_valid, error_message). If valid, error_message is empty.
+    """
+    inspector = inspect(engine)
+
+    # Get the baseline schema by running migrations up to the baseline revision
+    # in a temporary database and reflecting its schema
+    baseline_metadata = _get_baseline_metadata()
+    existing_tables = set(inspector.get_table_names())
+
+    # Filter out SQLite internal tables and alembic_version
+    existing_tables = {
+        t
+        for t in existing_tables
+        if not t.startswith("sqlite_") and t != ALEMBIC_VERSION_TABLE
+    }
+
+    # Exclude alembic_version from baseline comparison
+    expected_tables = {
+        t for t in baseline_metadata.tables.keys() if t != ALEMBIC_VERSION_TABLE
+    }
+    missing_tables = expected_tables - existing_tables
+
+    if missing_tables:
+        table_list = ", ".join(sorted(existing_tables))
+        missing_str = ", ".join(sorted(missing_tables))
+        return False, (
+            f"Detected tables: [{table_list}]. "
+            f"Missing baseline tables: [{missing_str}]."
+        )
+
+    # Verify columns for each expected table
+    for table_name in expected_tables:
+        table = baseline_metadata.tables[table_name]
+        expected_columns = {col.name for col in table.columns}
+        actual_columns = {col["name"] for col in inspector.get_columns(table_name)}
+
+        missing_cols = expected_columns - actual_columns
+        if missing_cols:
+            missing_cols_str = ", ".join(sorted(missing_cols))
+            return False, (
+                f"Table '{table_name}' missing columns: [{missing_cols_str}]."
+            )
+
+    return True, ""
+
+
+def stamp_existing_database(
+    engine: Engine, revision: str = FLWR_STATE_BASELINE_REVISION
+) -> None:
+    """Stamp an existing legacy database to the baseline Alembic revision."""
+    command.stamp(build_alembic_config(engine), revision)

flwr/supercore/state/alembic/versions/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Alembic migration versions."""

flwr/supercore/state/alembic/versions/rev_2026_01_28_initialize_migration_of_state_tables.py

@@ -0,0 +1,200 @@
+# Copyright 2026 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Initialize migration of state tables.
+
+Revision ID: 8e65d8ae60b0
+Revises:
+Create Date: 2026-01-28 11:03:18.038794
+"""
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+# pylint: disable=no-member
+
+# revision identifiers, used by Alembic.
+revision: str = "8e65d8ae60b0"
+down_revision: str | Sequence[str] | None = None
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    # LinkState tables
+    op.create_table(
+        "node",
+        sa.Column("node_id", sa.Integer(), nullable=True),
+        sa.Column("owner_aid", sa.String(), nullable=True),
+        sa.Column("owner_name", sa.String(), nullable=True),
+        sa.Column("status", sa.String(), nullable=True),
+        sa.Column("registered_at", sa.String(), nullable=True),
+        sa.Column("last_activated_at", sa.String(), nullable=True),
+        sa.Column("last_deactivated_at", sa.String(), nullable=True),
+        sa.Column("unregistered_at", sa.String(), nullable=True),
+        sa.Column("online_until", sa.TIMESTAMP(), nullable=True),
+        sa.Column("heartbeat_interval", sa.Float(), nullable=True),
+        sa.Column("public_key", sa.LargeBinary(), nullable=True),
+        sa.UniqueConstraint("node_id"),
+        sa.UniqueConstraint("public_key"),
+    )
+    op.create_index("idx_node_owner_aid", "node", ["owner_aid"], unique=False)
+    op.create_index("idx_node_status", "node", ["status"], unique=False)
+    op.create_index("idx_online_until", "node", ["online_until"], unique=False)
+    op.create_table(
+        "run",
+        sa.Column("run_id", sa.Integer(), nullable=True),
+        sa.Column("fab_id", sa.String(), nullable=True),
+        sa.Column("fab_version", sa.String(), nullable=True),
+        sa.Column("fab_hash", sa.String(), nullable=True),
+        sa.Column("override_config", sa.String(), nullable=True),
+        sa.Column("pending_at", sa.String(), nullable=True),
+        sa.Column("starting_at", sa.String(), nullable=True),
+        sa.Column("running_at", sa.String(), nullable=True),
+        sa.Column("finished_at", sa.String(), nullable=True),
+        sa.Column("sub_status", sa.String(), nullable=True),
+        sa.Column("details", sa.String(), nullable=True),
+        sa.Column("federation", sa.String(), nullable=True),
+        sa.Column("federation_options", sa.LargeBinary(), nullable=True),
+        sa.Column("flwr_aid", sa.String(), nullable=True),
+        sa.Column("bytes_sent", sa.Integer(), server_default="0", nullable=True),
+        sa.Column("bytes_recv", sa.Integer(), server_default="0", nullable=True),
+        sa.Column("clientapp_runtime", sa.Float(), server_default="0.0", nullable=True),
+        sa.UniqueConstraint("run_id"),
+    )
+    op.create_table(
+        "logs",
+        sa.Column("timestamp", sa.Float(), nullable=True),
+        sa.Column("run_id", sa.Integer(), nullable=True),
+        sa.Column("node_id", sa.Integer(), nullable=True),
+        sa.Column("log", sa.String(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["run_id"],
+            ["run.run_id"],
+        ),
+        sa.UniqueConstraint("timestamp", "run_id", "node_id"),
+    )
+    op.create_table(
+        "context",
+        sa.Column("run_id", sa.Integer(), nullable=True),
+        sa.Column("context", sa.LargeBinary(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["run_id"],
+            ["run.run_id"],
+        ),
+        sa.UniqueConstraint("run_id"),
+    )
+    op.create_table(
+        "message_ins",
+        sa.Column("message_id", sa.String(), nullable=True),
+        sa.Column("group_id", sa.String(), nullable=True),
+        sa.Column("run_id", sa.Integer(), nullable=True),
+        sa.Column("src_node_id", sa.Integer(), nullable=True),
+        sa.Column("dst_node_id", sa.Integer(), nullable=True),
+        sa.Column("reply_to_message_id", sa.String(), nullable=True),
+        sa.Column("created_at", sa.Float(), nullable=True),
+        sa.Column("delivered_at", sa.String(), nullable=True),
+        sa.Column("ttl", sa.Float(), nullable=True),
+        sa.Column("message_type", sa.String(), nullable=True),
+        sa.Column("content", sa.LargeBinary(), nullable=True),
+        sa.Column("error", sa.LargeBinary(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["run_id"],
+            ["run.run_id"],
+        ),
+        sa.UniqueConstraint("message_id"),
+    )
+    op.create_table(
+        "message_res",
+        sa.Column("message_id", sa.String(), nullable=True),
+        sa.Column("group_id", sa.String(), nullable=True),
+        sa.Column("run_id", sa.Integer(), nullable=True),
+        sa.Column("src_node_id", sa.Integer(), nullable=True),
+        sa.Column("dst_node_id", sa.Integer(), nullable=True),
+        sa.Column("reply_to_message_id", sa.String(), nullable=True),
+        sa.Column("created_at", sa.Float(), nullable=True),
+        sa.Column("delivered_at", sa.String(), nullable=True),
+        sa.Column("ttl", sa.Float(), nullable=True),
+        sa.Column("message_type", sa.String(), nullable=True),
+        sa.Column("content", sa.LargeBinary(), nullable=True),
+        sa.Column("error", sa.LargeBinary(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["run_id"],
+            ["run.run_id"],
+        ),
+        sa.UniqueConstraint("message_id"),
+    )
+    # CoreState tables
+    op.create_table(
+        "token_store",
+        sa.Column("run_id", sa.Integer(), nullable=True),
+        sa.Column("token", sa.String(), nullable=False),
+        sa.Column("active_until", sa.Float(), nullable=True),
+        sa.PrimaryKeyConstraint("run_id"),
+        sa.UniqueConstraint("token"),
+    )
+    # ObjectStore tables
+    op.create_table(
+        "objects",
+        sa.Column("object_id", sa.String(), nullable=True),
+        sa.Column("content", sa.LargeBinary(), nullable=True),
+        sa.Column("is_available", sa.Integer(), server_default="0", nullable=False),
+        sa.Column("ref_count", sa.Integer(), server_default="0", nullable=False),
+        sa.CheckConstraint("is_available IN (0, 1)", name="ck_objects_is_available"),
+        sa.PrimaryKeyConstraint("object_id"),
+    )
+    op.create_table(
+        "object_children",
+        sa.Column("parent_id", sa.String(), nullable=False),
+        sa.Column("child_id", sa.String(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["child_id"], ["objects.object_id"], ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(
+            ["parent_id"], ["objects.object_id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("parent_id", "child_id"),
+    )
+    op.create_table(
+        "run_objects",
+        sa.Column("run_id", sa.Integer(), nullable=False),
+        sa.Column("object_id", sa.String(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["object_id"], ["objects.object_id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("run_id", "object_id"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("run_objects")
+    op.drop_table("object_children")
+    op.drop_table("objects")
+    op.drop_table("token_store")
+    op.drop_table("message_res")
+    op.drop_table("message_ins")
+    op.drop_table("context")
+    op.drop_table("logs")
+    op.drop_table("run")
+    op.drop_index("idx_online_until", table_name="node")
+    op.drop_index("idx_node_status", table_name="node")
+    op.drop_index("idx_node_owner_aid", table_name="node")
+    op.drop_table("node")
+    # ### end Alembic commands ###

flwr/supercore/state/schema/README.md

@@ -0,0 +1,121 @@
+# State Entity Relationship Diagram
+
+## Schema
+
+<!-- BEGIN_SQLALCHEMY_DOCS -->
+```mermaid
+
+---
+    config:
+        layout: elk
+---
+erDiagram
+  context {
+    INTEGER run_id FK "nullable"
+    BLOB context "nullable"
+  }
+
+  logs {
+    INTEGER run_id FK "nullable"
+    VARCHAR log "nullable"
+    INTEGER node_id "nullable"
+    FLOAT timestamp "nullable"
+  }
+
+  message_ins {
+    INTEGER run_id FK "nullable"
+    BLOB content "nullable"
+    FLOAT created_at "nullable"
+    VARCHAR delivered_at "nullable"
+    INTEGER dst_node_id "nullable"
+    BLOB error "nullable"
+    VARCHAR group_id "nullable"
+    VARCHAR message_id UK "nullable"
+    VARCHAR message_type "nullable"
+    VARCHAR reply_to_message_id "nullable"
+    INTEGER src_node_id "nullable"
+    FLOAT ttl "nullable"
+  }
+
+  message_res {
+    INTEGER run_id FK "nullable"
+    BLOB content "nullable"
+    FLOAT created_at "nullable"
+    VARCHAR delivered_at "nullable"
+    INTEGER dst_node_id "nullable"
+    BLOB error "nullable"
+    VARCHAR group_id "nullable"
+    VARCHAR message_id UK "nullable"
+    VARCHAR message_type "nullable"
+    VARCHAR reply_to_message_id "nullable"
+    INTEGER src_node_id "nullable"
+    FLOAT ttl "nullable"
+  }
+
+  node {
+    FLOAT heartbeat_interval "nullable"
+    VARCHAR last_activated_at "nullable"
+    VARCHAR last_deactivated_at "nullable"
+    INTEGER node_id UK "nullable"
+    TIMESTAMP online_until "nullable"
+    VARCHAR owner_aid "nullable"
+    VARCHAR owner_name "nullable"
+    BLOB public_key UK "nullable"
+    VARCHAR registered_at "nullable"
+    VARCHAR status "nullable"
+    VARCHAR unregistered_at "nullable"
+  }
+
+  object_children {
+    VARCHAR child_id PK,FK
+    VARCHAR parent_id PK,FK
+  }
+
+  objects {
+    VARCHAR object_id PK "nullable"
+    BLOB content "nullable"
+    INTEGER is_available
+    INTEGER ref_count
+  }
+
+  run {
+    INTEGER bytes_recv "nullable"
+    INTEGER bytes_sent "nullable"
+    FLOAT clientapp_runtime "nullable"
+    VARCHAR details "nullable"
+    VARCHAR fab_hash "nullable"
+    VARCHAR fab_id "nullable"
+    VARCHAR fab_version "nullable"
+    VARCHAR federation "nullable"
+    BLOB federation_options "nullable"
+    VARCHAR finished_at "nullable"
+    VARCHAR flwr_aid "nullable"
+    VARCHAR override_config "nullable"
+    VARCHAR pending_at "nullable"
+    INTEGER run_id UK "nullable"
+    VARCHAR running_at "nullable"
+    VARCHAR starting_at "nullable"
+    VARCHAR sub_status "nullable"
+  }
+
+  run_objects {
+    VARCHAR object_id PK,FK
+    INTEGER run_id PK
+  }
+
+  token_store {
+    INTEGER run_id PK "nullable"
+    FLOAT active_until "nullable"
+    VARCHAR token UK
+  }
+
+  run ||--o| context : run_id
+  run ||--o{ logs : run_id
+  run ||--o{ message_ins : run_id
+  run ||--o{ message_res : run_id
+  objects ||--o| object_children : parent_id
+  objects ||--o| object_children : child_id
+  objects ||--o| run_objects : object_id
+
+```
+<!-- END_SQLALCHEMY_DOCS -->

flwr/supercore/state/schema/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SQLAlchemy database schema."""

flwr/supercore/state/schema/corestate_tables.py

@@ -0,0 +1,36 @@
+# Copyright 2026 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SQLAlchemy Core Table definitions for CoreState."""
+
+
+from sqlalchemy import Column, Float, Integer, MetaData, String, Table
+
+
+def create_corestate_metadata() -> MetaData:
+    """Create and return MetaData with CoreState table definitions."""
+    metadata = MetaData()
+
+    # --------------------------------------------------------------------------
+    #  Table: token_store
+    # --------------------------------------------------------------------------
+    Table(
+        "token_store",
+        metadata,
+        Column("run_id", Integer, primary_key=True, nullable=True),
+        Column("token", String, unique=True, nullable=False),
+        Column("active_until", Float),
+    )
+
+    return metadata