flwr-nightly 1.26.0.dev20260122__py3-none-any.whl → 1.26.0.dev20260126__py3-none-any.whl

flwr/supercore/corestate/sqlite_corestate.py

@@ -1,157 +0,0 @@
-# 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.
-# ==============================================================================
-"""SQLite-based CoreState implementation."""
-
-
-import secrets
-import sqlite3
-from typing import cast
-
-from flwr.common import now
-from flwr.common.constant import (
-    FLWR_APP_TOKEN_LENGTH,
-    HEARTBEAT_DEFAULT_INTERVAL,
-    HEARTBEAT_PATIENCE,
-)
-from flwr.supercore.sqlite_mixin import SqliteMixin
-from flwr.supercore.utils import int64_to_uint64, uint64_to_int64
-
-from ..object_store import ObjectStore
-from .corestate import CoreState
-
-SQL_CREATE_TABLE_TOKEN_STORE = """
-CREATE TABLE IF NOT EXISTS token_store (
-    run_id                  INTEGER PRIMARY KEY,
-    token                   TEXT UNIQUE NOT NULL,
-    active_until            REAL
-);
-"""
-
-
-class SqliteCoreState(CoreState, SqliteMixin):
-    """SQLite-based CoreState implementation."""
-
-    def __init__(self, database_path: str, object_store: ObjectStore) -> None:
-        super().__init__(database_path)
-        self._object_store = object_store
-
-    @property
-    def object_store(self) -> ObjectStore:
-        """Return the ObjectStore instance used by this CoreState."""
-        return self._object_store
-
-    def get_sql_statements(self) -> tuple[str, ...]:
-        """Return SQL statements needed for CoreState tables."""
-        return (SQL_CREATE_TABLE_TOKEN_STORE,)
-
-    def create_token(self, run_id: int) -> str | None:
-        """Create a token for the given run ID."""
-        token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
-        current = now().timestamp()
-        active_until = current + HEARTBEAT_DEFAULT_INTERVAL
-        query = """
-            INSERT INTO token_store (run_id, token, active_until)
-            VALUES (:run_id, :token, :active_until);
-        """
-        data = {
-            "run_id": uint64_to_int64(run_id),
-            "token": token,
-            "active_until": active_until,
-        }
-        try:
-            self.query(query, data)
-        except sqlite3.IntegrityError:
-            return None  # Token already created for this run ID
-        return token
-
-    def verify_token(self, run_id: int, token: str) -> bool:
-        """Verify a token for the given run ID."""
-        self._cleanup_expired_tokens()
-        query = "SELECT token FROM token_store WHERE run_id = :run_id;"
-        data = {"run_id": uint64_to_int64(run_id)}
-        rows = self.query(query, data)
-        if not rows:
-            return False
-        return cast(str, rows[0]["token"]) == token
-
-    def delete_token(self, run_id: int) -> None:
-        """Delete the token for the given run ID."""
-        query = "DELETE FROM token_store WHERE run_id = :run_id;"
-        data = {"run_id": uint64_to_int64(run_id)}
-        self.query(query, data)
-
-    def get_run_id_by_token(self, token: str) -> int | None:
-        """Get the run ID associated with a given token."""
-        self._cleanup_expired_tokens()
-        query = "SELECT run_id FROM token_store WHERE token = :token;"
-        data = {"token": token}
-        rows = self.query(query, data)
-        if not rows:
-            return None
-        return int64_to_uint64(rows[0]["run_id"])
-
-    def acknowledge_app_heartbeat(self, token: str) -> bool:
-        """Acknowledge an app heartbeat with the provided token."""
-        # Clean up expired tokens
-        self._cleanup_expired_tokens()
-
-        # Update the active_until field
-        current = now().timestamp()
-        active_until = current + HEARTBEAT_PATIENCE * HEARTBEAT_DEFAULT_INTERVAL
-        query = """
-            UPDATE token_store
-            SET active_until = :active_until
-            WHERE token = :token
-            RETURNING run_id;
-        """
-        data = {"active_until": active_until, "token": token}
-        rows = self.query(query, data)
-        return len(rows) > 0
-
-    def _cleanup_expired_tokens(self) -> None:
-        """Remove expired tokens and perform additional cleanup.
-
-        This method is called before token operations to ensure integrity.
-        Subclasses can override `_on_tokens_expired` to add custom cleanup logic.
-        """
-        current = now().timestamp()
-
-        with self.conn:
-            # Delete expired tokens and get their run_ids and active_until timestamps
-            query = """
-                DELETE FROM token_store
-                WHERE active_until < :current
-                RETURNING run_id, active_until;
-            """
-            rows = self.conn.execute(query, {"current": current}).fetchall()
-            expired_records = [
-                (int64_to_uint64(row["run_id"]), row["active_until"]) for row in rows
-            ]
-
-            # Hook for subclasses
-            if expired_records:
-                self._on_tokens_expired(expired_records)
-
-    def _on_tokens_expired(self, expired_records: list[tuple[int, float]]) -> None:
-        """Handle cleanup of expired tokens.
-
-        Override in subclasses to add custom cleanup logic.
-
-        Parameters
-        ----------
-        expired_records : list[tuple[int, float]]
-            List of tuples containing (run_id, active_until timestamp)
-            for expired tokens.
-        """

flwr/supercore/object_store/sqlite_object_store.py

@@ -1,253 +0,0 @@
-# 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.
-# ==============================================================================
-"""Flower SQLite ObjectStore implementation."""
-
-
-from typing import cast
-
-from flwr.common.inflatable import (
-    get_object_id,
-    is_valid_sha256_hash,
-    iterate_object_tree,
-)
-from flwr.common.inflatable_utils import validate_object_content
-from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
-from flwr.supercore.sqlite_mixin import SqliteMixin
-from flwr.supercore.utils import uint64_to_int64
-
-from .object_store import NoObjectInStoreError, ObjectStore
-
-SQL_CREATE_OBJECTS = """
-CREATE TABLE IF NOT EXISTS objects (
-    object_id               TEXT PRIMARY KEY,
-    content                 BLOB,
-    is_available            INTEGER NOT NULL CHECK (is_available IN (0,1)),
-    ref_count               INTEGER NOT NULL
-);
-"""
-SQL_CREATE_OBJECT_CHILDREN = """
-CREATE TABLE IF NOT EXISTS object_children (
-    parent_id               TEXT NOT NULL,
-    child_id                TEXT NOT NULL,
-    FOREIGN KEY (parent_id) REFERENCES objects(object_id) ON DELETE CASCADE,
-    FOREIGN KEY (child_id)  REFERENCES objects(object_id) ON DELETE CASCADE,
-    PRIMARY KEY (parent_id, child_id)
-);
-"""
-SQL_CREATE_RUN_OBJECTS = """
-CREATE TABLE IF NOT EXISTS run_objects (
-    run_id                  INTEGER NOT NULL,
-    object_id               TEXT NOT NULL,
-    FOREIGN KEY (object_id) REFERENCES objects(object_id) ON DELETE CASCADE,
-    PRIMARY KEY (run_id, object_id)
-);
-"""
-
-
-class SqliteObjectStore(ObjectStore, SqliteMixin):
-    """SQLite-based implementation of the ObjectStore interface."""
-
-    def __init__(self, database_path: str, verify: bool = True) -> None:
-        super().__init__(database_path)
-        self.verify = verify
-
-    def get_sql_statements(self) -> tuple[str, ...]:
-        """Return SQL statements for ObjectStore tables."""
-        return (
-            SQL_CREATE_OBJECTS,
-            SQL_CREATE_OBJECT_CHILDREN,
-            SQL_CREATE_RUN_OBJECTS,
-        )
-
-    def preregister(self, run_id: int, object_tree: ObjectTree) -> list[str]:
-        """Identify and preregister missing objects in the `ObjectStore`."""
-        new_objects = []
-        for tree_node in iterate_object_tree(object_tree):
-            obj_id = tree_node.object_id
-            if not is_valid_sha256_hash(obj_id):
-                raise ValueError(f"Invalid object ID format: {obj_id}")
-
-            child_ids = [child.object_id for child in tree_node.children]
-            with self.conn:
-                row = self.conn.execute(
-                    "SELECT object_id, is_available FROM objects WHERE object_id=?",
-                    (obj_id,),
-                ).fetchone()
-                if row is None:
-                    # Insert new object
-                    self.conn.execute(
-                        "INSERT INTO objects"
-                        "(object_id, content, is_available, ref_count) "
-                        "VALUES (?, ?, ?, ?)",
-                        (obj_id, b"", 0, 0),
-                    )
-                    for cid in child_ids:
-                        self.conn.execute(
-                            "INSERT INTO object_children(parent_id, child_id) "
-                            "VALUES (?, ?)",
-                            (obj_id, cid),
-                        )
-                        self.conn.execute(
-                            "UPDATE objects SET ref_count = ref_count + 1 "
-                            "WHERE object_id = ?",
-                            (cid,),
-                        )
-                    new_objects.append(obj_id)
-                else:
-                    # Add to the list of new objects if not available
-                    if not row["is_available"]:
-                        new_objects.append(obj_id)
-
-                # Ensure run mapping
-                self.conn.execute(
-                    "INSERT OR IGNORE INTO run_objects(run_id, object_id) "
-                    "VALUES (?, ?)",
-                    (uint64_to_int64(run_id), obj_id),
-                )
-        return new_objects
-
-    def get_object_tree(self, object_id: str) -> ObjectTree:
-        """Get the object tree for a given object ID."""
-        with self.conn:
-            row = self.conn.execute(
-                "SELECT object_id FROM objects WHERE object_id=?", (object_id,)
-            ).fetchone()
-            if not row:
-                raise NoObjectInStoreError(
-                    f"Object {object_id} was not pre-registered."
-                )
-            children = self.query(
-                "SELECT child_id FROM object_children WHERE parent_id=?", (object_id,)
-            )
-
-            # Build the object trees of all children
-            try:
-                child_trees = [self.get_object_tree(ch["child_id"]) for ch in children]
-            except NoObjectInStoreError as e:
-                # Raise an error if any child object is missing
-                # This indicates an integrity issue
-                raise NoObjectInStoreError(
-                    f"Object tree for object ID '{object_id}' contains missing "
-                    "children. This may indicate a corrupted object store."
-                ) from e
-
-            # Create and return the ObjectTree for the current object
-            return ObjectTree(object_id=object_id, children=child_trees)
-
-    def put(self, object_id: str, object_content: bytes) -> None:
-        """Put an object into the store."""
-        if self.verify:
-            # Verify object_id and object_content match
-            object_id_from_content = get_object_id(object_content)
-            if object_id != object_id_from_content:
-                raise ValueError(f"Object ID {object_id} does not match content hash")
-
-            # Validate object content
-            validate_object_content(content=object_content)
-
-        with self.conn:
-            # Only allow adding the object if it has been preregistered
-            row = self.conn.execute(
-                "SELECT is_available FROM objects WHERE object_id=?", (object_id,)
-            ).fetchone()
-            if row is None:
-                raise NoObjectInStoreError(
-                    f"Object with ID '{object_id}' was not pre-registered."
-                )
-
-            # Return if object is already present in the store
-            if row["is_available"]:
-                return
-
-            # Update the object entry in the store
-            self.conn.execute(
-                "UPDATE objects SET content=?, is_available=1 WHERE object_id=?",
-                (object_content, object_id),
-            )
-
-    def get(self, object_id: str) -> bytes | None:
-        """Get an object from the store."""
-        rows = self.query("SELECT content FROM objects WHERE object_id=?", (object_id,))
-        return rows[0]["content"] if rows else None
-
-    def delete(self, object_id: str) -> None:
-        """Delete an object and its unreferenced descendants from the store."""
-        with self.conn:
-            row = self.conn.execute(
-                "SELECT ref_count FROM objects WHERE object_id=?", (object_id,)
-            ).fetchone()
-
-            # If the object is not in the store, nothing to delete
-            if row is None:
-                return
-
-            # Skip deletion if there are still references
-            if row["ref_count"] > 0:
-                return
-
-            # Deleting will cascade via FK, but we need to decrement children first
-            children = self.conn.execute(
-                "SELECT child_id FROM object_children WHERE parent_id=?", (object_id,)
-            ).fetchall()
-            child_ids = [child["child_id"] for child in children]
-
-            if child_ids:
-                placeholders = ", ".join("?" for _ in child_ids)
-                query = f"""
-                    UPDATE objects SET ref_count = ref_count - 1
-                    WHERE object_id IN ({placeholders})
-                """
-                self.conn.execute(query, child_ids)
-
-            self.conn.execute("DELETE FROM objects WHERE object_id=?", (object_id,))
-
-            # Recursively clean children
-            for child_id in child_ids:
-                self.delete(child_id)
-
-    def delete_objects_in_run(self, run_id: int) -> None:
-        """Delete all objects that were registered in a specific run."""
-        run_id_sint = uint64_to_int64(run_id)
-        with self.conn:
-            objs = self.conn.execute(
-                "SELECT object_id FROM run_objects WHERE run_id=?", (run_id_sint,)
-            ).fetchall()
-            for obj in objs:
-                object_id = obj["object_id"]
-                row = self.conn.execute(
-                    "SELECT ref_count FROM objects WHERE object_id=?", (object_id,)
-                ).fetchone()
-                if row and row["ref_count"] == 0:
-                    self.delete(object_id)
-            self.conn.execute("DELETE FROM run_objects WHERE run_id=?", (run_id_sint,))
-
-    def clear(self) -> None:
-        """Clear the store."""
-        with self.conn:
-            self.conn.execute("DELETE FROM object_children;")
-            self.conn.execute("DELETE FROM run_objects;")
-            self.conn.execute("DELETE FROM objects;")
-
-    def __contains__(self, object_id: str) -> bool:
-        """Check if an object_id is in the store."""
-        row = self.conn.execute(
-            "SELECT 1 FROM objects WHERE object_id=?", (object_id,)
-        ).fetchone()
-        return row is not None
-
-    def __len__(self) -> int:
-        """Return the number of objects in the store."""
-        row = self.conn.execute("SELECT COUNT(*) AS cnt FROM objects;").fetchone()
-        return cast(int, row["cnt"])

flwr/supercore/sqlite_mixin.py

@@ -1,156 +0,0 @@
-# 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))