flwr-nightly 1.23.0.dev20250930__py3-none-any.whl → 1.26.0.dev20260121__py3-none-any.whl

flwr/supercore/sql_mixin.py

@@ -0,0 +1,292 @@
+# 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.
+# ==============================================================================
+"""Mixin providing common SQL connection and initialization logic via SQLAlchemy."""
+
+
+import re
+from abc import ABC
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+from contextvars import ContextVar
+from logging import DEBUG, ERROR
+from pathlib import Path
+from typing import Any
+
+from sqlalchemy import Engine, MetaData, create_engine, event, inspect, text
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import Session, sessionmaker
+
+from flwr.common.logger import log
+from flwr.supercore.constant import SQLITE_PRAGMAS
+
+_current_session: ContextVar[Session | None] = ContextVar(
+    "current_sqlalchemy_session",
+    default=None,
+)
+
+
+def _set_sqlite_pragmas(dbapi_conn: Any, _connection_record: Any) -> None:
+    """Set SQLite pragmas for performance and correctness."""
+    cursor = dbapi_conn.cursor()
+    for pragma, value in SQLITE_PRAGMAS:
+        cursor.execute(f"PRAGMA {pragma} = {value};")
+    cursor.close()
+
+
+def _log_query(  # pylint: disable=W0613,R0913,R0917
+    conn: Any,
+    cursor: Any,
+    statement: str,
+    parameters: Any,
+    context: Any,
+    executemany: bool,
+) -> None:
+    """Log SQL queries via Flower logger."""
+    log(DEBUG, {"query": statement, "params": parameters})
+
+
+class SqlMixin(ABC):
+    """Mixin providing common SQLite connection and initialization logic.
+
+    This mixin uses SQLAlchemy Core API for SQLite database access. It accepts either a
+    database file path or a SQLite URL, automatically converting file paths to SQLite
+    URLs.
+    """
+
+    def __init__(self, database_path: str) -> None:
+        """Initialize the SqlMixin.
+
+        Parameters
+        ----------
+        database_path : str
+            Either a file path or SQLite database URL. Examples:
+            - "path/to/db.db" or "/absolute/path/to/db.db"
+            - ":memory:" for in-memory SQLite
+            - "sqlite:///path/to/db.db" for explicit SQLite URL
+        """
+        # Auto-convert file path to SQLAlchemy SQLite URL if needed
+        if database_path == ":memory:":
+            self.database_url = "sqlite:///:memory:"
+        elif not database_path.startswith("sqlite://"):
+            # Treat as file path, convert to absolute and create SQLite URL
+            abs_path = Path(database_path).resolve()
+            self.database_url = f"sqlite:///{abs_path}"
+        else:
+            # Already a SQLite URL
+            self.database_url = database_path
+
+        self._engine: Engine | None = None
+        self._session_factory: sessionmaker[Session] | None = None
+
+    @contextmanager
+    def session(self) -> Iterator[Session]:
+        """Provide a transactional database session context.
+
+        Yields a SQLAlchemy Session that automatically commits on success or rolls
+        back on exceptions. Re-entrant: nested calls reuse the same session. Use for
+        multi-statement transactions; prefer `query()` for single statements.
+
+        Yields
+        ------
+        Session
+            SQLAlchemy session. Commits on context exit, rolls back on exceptions.
+
+        Examples
+        --------
+            with self.session() as session:
+                session.execute(text("DELETE FROM t WHERE id = :id"), {"id": 1})
+                session.execute(text("INSERT INTO t2 SELECT * FROM t"))
+        """
+        existing = _current_session.get()
+
+        # Re-entrant: reuse the session if in a session context, no begin()
+        if existing is not None:
+            yield existing
+            return
+
+        if self._session_factory is None:
+            raise AttributeError("Database not initialized. Call initialize() first.")
+
+        # Create new session; outermost scope owns the transaction
+        session = self._session_factory()
+        token = _current_session.set(session)
+
+        try:
+            with session.begin():
+                yield session
+        finally:
+            _current_session.reset(token)
+            session.close()
+
+    def get_metadata(self) -> MetaData | None:
+        """Return the MetaData object for this class.
+
+        Subclasses can override this to provide their SQLAlchemy MetaData.
+        The base implementation returns None.
+
+        Returns
+        -------
+        MetaData | None
+            SQLAlchemy MetaData object for this class.
+        """
+        return None
+
+    def initialize(self, log_queries: bool = False) -> list[str]:
+        """Connect to the DB and create tables if needed.
+
+        This method creates the SQLAlchemy engine and session factory,
+        and creates tables returned by `get_metadata()`.
+
+        Parameters
+        ----------
+        log_queries : bool
+            Log each query which is executed.
+
+        Returns
+        -------
+        list[str]
+            The list of all tables in the DB.
+        """
+        # Create engine with SQLite-specific settings
+        engine_kwargs: dict[str, Any] = {
+            # SQLite needs check_same_thread=False for multi-threaded access
+            "connect_args": {"check_same_thread": False}
+        }
+        self._engine = create_engine(self.database_url, **engine_kwargs)
+
+        # Set SQLite pragmas via event listener for optimal performance and correctness
+        event.listen(self._engine, "connect", _set_sqlite_pragmas)
+
+        if log_queries:
+            # Set up query logging via event listener
+            event.listen(self._engine, "before_cursor_execute", _log_query)
+
+        # 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:
+            metadata.create_all(self._engine)
+
+        # Get all table names using inspector
+        inspector = inspect(self._engine)
+        return inspector.get_table_names()
+
+    def query(
+        self,
+        query: str,
+        data: Sequence[dict[str, Any]] | dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Execute a SQL query and return the results as list of dicts.
+
+        TRANSACTION SEMANTICS:
+        ----------------------
+        If called outside a session context, each call to query() runs in its own
+        isolated transaction that is automatically committed. This is suitable for
+        single SQL statements.
+
+        If called within a session() context, query() reuses the existing session
+        and transaction. This enables atomic multi-query operations:
+
+            with self.session() as session:
+                self.query("UPDATE ...", {...})    # Shares same transaction
+                self.query("INSERT ...", {...})    # Shares same transaction
+                # Both succeed or fail together
+
+        You can also use session.execute() directly for the same effect:
+
+            with self.session() as session:
+                session.execute(text("UPDATE ..."), {...})
+                session.execute(text("INSERT ..."), {...})
+
+        Parameters
+        ----------
+        query : str
+            SQL query string with named parameter placeholders.
+            Use :name syntax for parameters: "SELECT * FROM t WHERE a = :a AND b = :b"
+        data : Sequence[dict[str, Any]] | dict[str, Any] | None
+            Query parameters using named parameter syntax:
+            - Single execution: pass dict, e.g., {"a": value1, "b": value2}
+            - Batch execution: pass sequence of dicts, e.g., [{"a": 1}, {"a": 2}]
+
+        Returns
+        -------
+        list[dict[str, Any]]
+            Query results as a list of dictionaries.
+
+        Examples
+        --------
+        # Single query with named parameters (auto-committed transaction)
+        rows = self.query(
+            "SELECT * FROM node WHERE node_id = :id AND status = :status",
+            {"id": node_id, "status": status}
+        )
+
+        # Batch insert with named parameters (auto-committed transaction)
+        rows = self.query(
+            "INSERT INTO node (node_id, status) VALUES (:id, :status)",
+            [{"id": 1, "status": "online"}, {"id": 2, "status": "offline"}]
+        )
+
+        # Multi-statement transaction - query() calls share the same session
+        with self.session():
+            # Both statements succeed or fail together
+            self.query("DELETE FROM token_store WHERE active_until < :time", {...})
+            self.query("UPDATE run SET status = :status WHERE id = :id", {...})
+
+        # Nested session() - query() calls share the same session
+        with self.session():
+            self.query("DELETE FROM token_store WHERE active_until < :time", {...})
+            with self.session():
+                self.query("UPDATE run SET status = :status WHERE id = :id", {...})
+        """
+        if self._engine is None:
+            raise AttributeError(
+                "LinkState is not initialized. Call initialize() first."
+            )
+
+        if data is None:
+            data = {}
+
+        # Clean up whitespace to make the logs nicer
+        query = re.sub(r"\s+", " ", query.strip())
+
+        try:
+            # Wrap query in text() to enable SQLAlchemy named parameter syntax (:param).
+            sql = text(query)
+
+            def execute_and_fetch(session: Session) -> list[dict[str, Any]]:
+                """Execute query and fetch results from the given session."""
+                # Execute query (results live in database cursor).
+                # There is no need to check for batch vs single execution;
+                # SQLAlchemy handles both cases automatically.
+                result = session.execute(sql, data)
+
+                # Fetch results into Python memory before commit.
+                # mappings() returns dict-like rows (works for SELECT and RETURNING).
+                if result.returns_rows:  # type: ignore
+                    return [dict(row) for row in result.mappings()]
+
+                # For statements without RETURNING (INSERT/UPDATE/DELETE),
+                # returns_rows is False, so we return empty list.
+                return []
+
+            # Not in a session context, create a new session context for this query
+            with self.session() as session:
+                return execute_and_fetch(session)
+
+        except SQLAlchemyError as exc:
+            log(ERROR, {"query": query, "data": data, "exception": exc})
+            raise

flwr/supercore/sqlite_mixin.py

@@ -0,0 +1,156 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Mixin providing common SQLite connection and initialization logic."""
+
+
+import re
+import sqlite3
+from abc import ABC
+from collections.abc import Sequence
+from logging import DEBUG, ERROR
+from typing import Any
+
+from flwr.common.logger import log
+from flwr.supercore.constant import SQLITE_PRAGMAS
+
+DictOrTuple = tuple[Any, ...] | dict[str, Any]
+
+
+class SqliteMixin(ABC):
+    """Mixin providing common SQLite connection and initialization logic."""
+
+    def __init__(self, database_path: str) -> None:
+        self.database_path = database_path
+        self._conn: sqlite3.Connection | None = None
+
+    @property
+    def conn(self) -> sqlite3.Connection:
+        """Get the SQLite connection."""
+        if self._conn is None:
+            raise AttributeError("Database not initialized. Call initialize() first.")
+        return self._conn
+
+    def get_sql_statements(self) -> tuple[str, ...]:
+        """Return SQL statements for this class.
+
+        Subclasses can override this to provide their SQL CREATE statements.
+        The base implementation returns an empty tuple.
+
+        Returns
+        -------
+        tuple[str, ...]
+            SQL CREATE TABLE/INDEX statements for this class.
+        """
+        return ()
+
+    def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
+        """Connect to the DB, enable FK support, and create tables if needed.
+
+        This method executes SQL statements returned by `get_sql_statements()`.
+
+        Parameters
+        ----------
+        log_queries : bool
+            Log each query which is executed.
+
+        Returns
+        -------
+        list[tuple[str]]
+            The list of all tables in the DB.
+
+        Examples
+        --------
+        Override `get_sql_statements()` in your subclass:
+
+        .. code:: python
+
+            def get_sql_statements(self) -> tuple[str, ...]:
+                return (
+                    SQL_CREATE_TABLE_FOO,
+                    SQL_CREATE_TABLE_BAR,
+                )
+
+        To include parent SQL statements, call super():
+
+        .. code:: python
+
+            def get_sql_statements(self) -> tuple[str, ...]:
+                return super().get_sql_statements() + (
+                    SQL_CREATE_TABLE_FOO,
+                    SQL_CREATE_TABLE_BAR,
+                )
+        """
+        self._conn = sqlite3.connect(self.database_path)
+        # Set SQLite pragmas for optimal performance and correctness
+        for pragma, value in SQLITE_PRAGMAS:
+            self._conn.execute(f"PRAGMA {pragma} = {value};")
+        self._conn.row_factory = dict_factory
+
+        if log_queries:
+            self._conn.set_trace_callback(lambda q: log(DEBUG, q))
+
+        # Create tables and indexes
+        cur = self._conn.cursor()
+        for sql in self.get_sql_statements():
+            cur.execute(sql)
+        res = cur.execute("SELECT name FROM sqlite_schema;")
+        return res.fetchall()
+
+    def query(
+        self,
+        query: str,
+        data: Sequence[DictOrTuple] | DictOrTuple | None = None,
+    ) -> list[dict[str, Any]]:
+        """Execute a SQL query and return the results as list of dicts."""
+        if self._conn is None:
+            raise AttributeError("LinkState is not initialized.")
+
+        if data is None:
+            data = []
+
+        # Clean up whitespace to make the logs nicer
+        query = re.sub(r"\s+", " ", query)
+
+        try:
+            with self._conn:
+                if (
+                    len(data) > 0
+                    and isinstance(data, (tuple | list))
+                    and isinstance(data[0], (tuple | dict))
+                ):
+                    rows = self._conn.executemany(query, data)
+                else:
+                    rows = self._conn.execute(query, data)
+
+                # Extract results before committing to support
+                #   INSERT/UPDATE ... RETURNING
+                # style queries
+                result = rows.fetchall()
+        except KeyError as exc:
+            log(ERROR, {"query": query, "data": data, "exception": exc})
+
+        return result
+
+
+def dict_factory(
+    cursor: sqlite3.Cursor,
+    row: sqlite3.Row,
+) -> dict[str, Any]:
+    """Turn SQLite results into dicts.
+
+    Less efficent for retrival of large amounts of data but easier to use.
+    """
+    fields = [column[0] for column in cursor.description]
+    return dict(zip(fields, row, strict=True))

flwr/{client/clientapp → supercore/state}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+# 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.
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower AppIO service."""
+"""Flower SuperCore state components."""

flwr/supercore/state/schema/README.md

@@ -0,0 +1,125 @@
+# State Entity Relationship Diagram
+
+## Schema
+
+```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
+  }
+
+<<<<<<< HEAD
+  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"
+  }
+
+=======
+>>>>>>> main
+  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
+  }
+
+<<<<<<< HEAD
+  run ||--o| context : run_id
+  run ||--o{ logs : run_id
+  run ||--o{ message_ins : run_id
+  run ||--o{ message_res : run_id
+=======
+>>>>>>> main
+  objects ||--o| object_children : parent_id
+  objects ||--o| object_children : child_id
+  objects ||--o| run_objects : object_id
+
+```

flwr/{cli/new/templates → supercore/state/schema}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+# 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.
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower CLI `new` command templates."""
+"""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