flwr 1.21.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -18,18 +18,16 @@
 # pylint: disable=too-many-lines
 
 import json
-import re
 import secrets
 import sqlite3
-import time
 from collections.abc import Sequence
-from logging import DEBUG, ERROR, WARNING
+from logging import ERROR, WARNING
 from typing import Any, Optional, Union, cast
 
 from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
     FLWR_APP_TOKEN_LENGTH,
-    HEARTBEAT_MAX_INTERVAL,
+    HEARTBEAT_INTERVAL_INF,
     HEARTBEAT_PATIENCE,
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
@@ -47,10 +45,14 @@ from flwr.common.typing import Run, RunStatus, UserConfig
 
 # pylint: disable=E0611
 from flwr.proto.error_pb2 import Error as ProtoError
+from flwr.proto.node_pb2 import NodeInfo
 from flwr.proto.recorddict_pb2 import RecordDict as ProtoRecordDict
 
 # pylint: enable=E0611
 from flwr.server.utils.validator import validate_message
+from flwr.supercore.constant import NodeStatus
+from flwr.supercore.sqlite_mixin import SqliteMixin
+from flwr.supercore.utils import int64_to_uint64, uint64_to_int64
 
 from .linkstate import LinkState
 from .utils import (
@@ -59,9 +61,7 @@ from .utils import (
     configrecord_to_bytes,
     context_from_bytes,
     context_to_bytes,
-    convert_sint64_to_uint64,
     convert_sint64_values_in_dict_to_uint64,
-    convert_uint64_to_sint64,
     convert_uint64_values_in_dict_to_sint64,
     generate_rand_int_from_bytes,
     has_valid_sub_status,
@@ -72,10 +72,16 @@ from .utils import (
 
 SQL_CREATE_TABLE_NODE = """
 CREATE TABLE IF NOT EXISTS node(
-    node_id         INTEGER UNIQUE,
-    online_until    REAL,
-    heartbeat_interval   REAL,
-    public_key      BLOB
+    node_id                 INTEGER UNIQUE,
+    owner_aid               TEXT,
+    status                  TEXT,
+    registered_at           TEXT,
+    last_activated_at       TEXT NULL,
+    last_deactivated_at     TEXT NULL,
+    unregistered_at         TEXT NULL,
+    online_until            TIMESTAMP NULL,
+    heartbeat_interval      REAL,
+    public_key              BLOB UNIQUE
 );
 """
 
@@ -89,6 +95,14 @@ SQL_CREATE_INDEX_ONLINE_UNTIL = """
 CREATE INDEX IF NOT EXISTS idx_online_until ON node (online_until);
 """
 
+SQL_CREATE_INDEX_OWNER_AID = """
+CREATE INDEX IF NOT EXISTS idx_node_owner_aid ON node(owner_aid);
+"""
+
+SQL_CREATE_INDEX_NODE_STATUS = """
+CREATE INDEX IF NOT EXISTS idx_node_status ON node(status);
+"""
+
 SQL_CREATE_TABLE_RUN = """
 CREATE TABLE IF NOT EXISTS run(
     run_id                INTEGER UNIQUE,
@@ -172,94 +186,26 @@ CREATE TABLE IF NOT EXISTS token_store (
 );
 """
 
-DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
 
-
-class SqliteLinkState(LinkState):  # pylint: disable=R0904
+class SqliteLinkState(LinkState, SqliteMixin):  # pylint: disable=R0904
     """SQLite-based LinkState implementation."""
 
-    def __init__(
-        self,
-        database_path: str,
-    ) -> None:
-        """Initialize an SqliteLinkState.
-
-        Parameters
-        ----------
-        database : (path-like object)
-            The path to the database file to be opened. Pass ":memory:" to open
-            a connection to a database that is in RAM, instead of on disk.
-        """
-        self.database_path = database_path
-        self.conn: Optional[sqlite3.Connection] = None
-
     def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
-        """Create tables if they don't exist yet.
-
-        Parameters
-        ----------
-        log_queries : bool
-            Log each query which is executed.
-
-        Returns
-        -------
-        list[tuple[str]]
-            The list of all tables in the DB.
-        """
-        self.conn = sqlite3.connect(self.database_path)
-        self.conn.execute("PRAGMA foreign_keys = ON;")
-        self.conn.row_factory = dict_factory
-        if log_queries:
-            self.conn.set_trace_callback(lambda query: log(DEBUG, query))
-        cur = self.conn.cursor()
-
-        # Create each table if not exists queries
-        cur.execute(SQL_CREATE_TABLE_RUN)
-        cur.execute(SQL_CREATE_TABLE_LOGS)
-        cur.execute(SQL_CREATE_TABLE_CONTEXT)
-        cur.execute(SQL_CREATE_TABLE_MESSAGE_INS)
-        cur.execute(SQL_CREATE_TABLE_MESSAGE_RES)
-        cur.execute(SQL_CREATE_TABLE_NODE)
-        cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
-        cur.execute(SQL_CREATE_TABLE_TOKEN_STORE)
-        cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
-        res = cur.execute("SELECT name FROM sqlite_schema;")
-        return res.fetchall()
-
-    def query(
-        self,
-        query: str,
-        data: Optional[Union[Sequence[DictOrTuple], DictOrTuple]] = None,
-    ) -> list[dict[str, Any]]:
-        """Execute a SQL query."""
-        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
+        """Connect to the DB, enable FK support, and create tables if needed."""
+        return self._ensure_initialized(
+            SQL_CREATE_TABLE_RUN,
+            SQL_CREATE_TABLE_LOGS,
+            SQL_CREATE_TABLE_CONTEXT,
+            SQL_CREATE_TABLE_MESSAGE_INS,
+            SQL_CREATE_TABLE_MESSAGE_RES,
+            SQL_CREATE_TABLE_NODE,
+            SQL_CREATE_TABLE_PUBLIC_KEY,
+            SQL_CREATE_TABLE_TOKEN_STORE,
+            SQL_CREATE_INDEX_ONLINE_UNTIL,
+            SQL_CREATE_INDEX_OWNER_AID,
+            SQL_CREATE_INDEX_NODE_STATUS,
+            log_queries=log_queries,
+        )
 
     def store_message_ins(self, message: Message) -> Optional[str]:
         """Store one Message."""
@@ -293,8 +239,10 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             return None
 
         # Validate destination node ID
-        query = "SELECT node_id FROM node WHERE node_id = ?;"
-        if not self.query(query, (data[0]["dst_node_id"],)):
+        query = "SELECT node_id FROM node WHERE node_id = ? AND status IN (?, ?);"
+        if not self.query(
+            query, (data[0]["dst_node_id"], NodeStatus.ONLINE, NodeStatus.OFFLINE)
+        ):
             log(
                 ERROR,
                 "Invalid destination node ID for Message: %s",
@@ -323,7 +271,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         data: dict[str, Union[str, int]] = {}
 
         # Convert the uint64 value to sint64 for SQLite
-        data["node_id"] = convert_uint64_to_sint64(node_id)
+        data["node_id"] = uint64_to_int64(node_id)
 
         # Retrieve all Messages for node_id
         query = """
@@ -398,8 +346,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         if (
             msg_ins
             and message
-            and convert_sint64_to_uint64(msg_ins["dst_node_id"])
-            != res_metadata.src_node_id
+            and int64_to_uint64(msg_ins["dst_node_id"]) != res_metadata.src_node_id
         ):
             return None
 
@@ -451,7 +398,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         ret: dict[str, Message] = {}
 
         # Verify Message IDs
-        current = time.time()
+        current = now().timestamp()
         query = f"""
             SELECT *
             FROM message_ins
@@ -475,20 +422,20 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         dst_node_ids: set[int] = set()
         for message_id in message_ids:
             in_message = found_message_ins_dict[message_id]
-            sint_node_id = convert_uint64_to_sint64(in_message.metadata.dst_node_id)
+            sint_node_id = uint64_to_int64(in_message.metadata.dst_node_id)
             dst_node_ids.add(sint_node_id)
         query = f"""
-                    SELECT node_id, online_until
-                    FROM node
-                    WHERE node_id IN ({",".join(["?"] * len(dst_node_ids))});
-                """
-        rows = self.query(query, tuple(dst_node_ids))
+            SELECT node_id, online_until
+            FROM node
+            WHERE node_id IN ({",".join(["?"] * len(dst_node_ids))})
+            AND status != ?
+        """
+        rows = self.query(query, tuple(dst_node_ids) + (NodeStatus.UNREGISTERED,))
         tmp_ret_dict = check_node_availability_for_in_message(
             inquired_in_message_ids=message_ids,
             found_in_message_dict=found_message_ins_dict,
             node_id_to_online_until={
-                convert_sint64_to_uint64(row["node_id"]): row["online_until"]
-                for row in rows
+                int64_to_uint64(row["node_id"]): row["online_until"] for row in rows
             },
             current_time=current,
         )
@@ -589,7 +536,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             WHERE run_id = :run_id;
         """
 
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         data = {"run_id": sint64_run_id}
 
         with self.conn:
@@ -597,7 +544,9 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return {row["message_id"] for row in rows}
 
-    def create_node(self, heartbeat_interval: float) -> int:
+    def create_node(
+        self, owner_aid: str, public_key: bytes, heartbeat_interval: float
+    ) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random uint64 as node_id
         uint64_node_id = generate_rand_int_from_bytes(
@@ -605,50 +554,126 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         )
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_node_id = convert_uint64_to_sint64(uint64_node_id)
+        sint64_node_id = uint64_to_int64(uint64_node_id)
 
-        query = (
-            "INSERT INTO node "
-            "(node_id, online_until, heartbeat_interval, public_key) "
-            "VALUES (?, ?, ?, ?)"
-        )
+        query = """
+            INSERT INTO node
+            (node_id, owner_aid, status, registered_at, last_activated_at,
+            last_deactivated_at, unregistered_at, online_until, heartbeat_interval,
+            public_key)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+        """
 
-        # Mark the node online util time.time() + heartbeat_interval
+        # Mark the node online until now().timestamp() + heartbeat_interval
         try:
             self.query(
                 query,
                 (
-                    sint64_node_id,
-                    time.time() + heartbeat_interval,
-                    heartbeat_interval,
-                    b"",  # Initialize with an empty public key
+                    sint64_node_id,  # node_id
+                    owner_aid,  # owner_aid
+                    NodeStatus.REGISTERED,  # status
+                    now().isoformat(),  # registered_at
+                    None,  # last_activated_at
+                    None,  # last_deactivated_at
+                    None,  # unregistered_at
+                    None,  # online_until, initialized with offline status
+                    heartbeat_interval,  # heartbeat_interval
+                    public_key,  # public_key
                 ),
             )
-        except sqlite3.IntegrityError:
+        except sqlite3.IntegrityError as e:
+            if "UNIQUE constraint failed: node.public_key" in str(e):
+                raise ValueError("Public key already in use.") from None
+            # Must be node ID conflict, almost impossible unless system is compromised
             log(ERROR, "Unexpected node registration failure.")
             return 0
 
         # Note: we need to return the uint64 value of the node_id
         return uint64_node_id
 
-    def delete_node(self, node_id: int) -> None:
+    def delete_node(self, owner_aid: str, node_id: int) -> None:
         """Delete a node."""
-        # Convert the uint64 value to sint64 for SQLite
-        sint64_node_id = convert_uint64_to_sint64(node_id)
+        sint64_node_id = uint64_to_int64(node_id)
 
-        query = "DELETE FROM node WHERE node_id = ?"
-        params = (sint64_node_id,)
+        query = """
+            UPDATE node
+            SET status = ?, unregistered_at = ?,
+            online_until = IIF(online_until > ?, ?, online_until)
+            WHERE node_id = ? AND status != ? AND owner_aid = ?
+            RETURNING node_id
+        """
+        current = now()
+        params = (
+            NodeStatus.UNREGISTERED,
+            current.isoformat(),
+            current.timestamp(),
+            current.timestamp(),
+            sint64_node_id,
+            NodeStatus.UNREGISTERED,
+            owner_aid,
+        )
 
-        if self.conn is None:
-            raise AttributeError("LinkState is not initialized.")
+        rows = self.query(query, params)
+        if not rows:
+            raise ValueError(
+                f"Node {node_id} already deleted, not found or unauthorized "
+                "deletion attempt."
+            )
 
-        try:
-            with self.conn:
-                rows = self.conn.execute(query, params)
-                if rows.rowcount < 1:
-                    raise ValueError(f"Node {node_id} not found")
-        except KeyError as exc:
-            log(ERROR, {"query": query, "data": params, "exception": exc})
+    def activate_node(self, node_id: int, heartbeat_interval: float) -> bool:
+        """Activate the node with the specified `node_id`."""
+        with self.conn:
+            self._check_and_tag_offline_nodes([node_id])
+
+            # Only activate if the node is currently registered or offline
+            current_dt = now()
+            query = """
+                UPDATE node
+                SET status = ?,
+                    last_activated_at = ?,
+                    online_until = ?,
+                    heartbeat_interval = ?
+                WHERE node_id = ? AND status in (?, ?)
+                RETURNING node_id
+            """
+            params = (
+                NodeStatus.ONLINE,
+                current_dt.isoformat(),
+                current_dt.timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval,
+                heartbeat_interval,
+                uint64_to_int64(node_id),
+                NodeStatus.REGISTERED,
+                NodeStatus.OFFLINE,
+            )
+
+            row = self.conn.execute(query, params).fetchone()
+            return row is not None
+
+    def deactivate_node(self, node_id: int) -> bool:
+        """Deactivate the node with the specified `node_id`."""
+        with self.conn:
+            self._check_and_tag_offline_nodes([node_id])
+
+            # Only deactivate if the node is currently online
+            current_dt = now()
+            query = """
+                UPDATE node
+                SET status = ?,
+                    last_deactivated_at = ?,
+                    online_until = ?
+                WHERE node_id = ? AND status = ?
+                RETURNING node_id
+            """
+            params = (
+                NodeStatus.OFFLINE,
+                current_dt.isoformat(),
+                current_dt.timestamp(),
+                uint64_to_int64(node_id),
+                NodeStatus.ONLINE,
+            )
+
+            row = self.conn.execute(query, params).fetchone()
+            return row is not None
 
     def get_nodes(self, run_id: int) -> set[int]:
         """Retrieve all currently stored node IDs as a set.
@@ -658,69 +683,117 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         If the provided `run_id` does not exist or has no matching nodes,
         an empty `Set` MUST be returned.
         """
+        if self.conn is None:
+            raise AttributeError("LinkState not initialized")
+
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
 
         # Validate run ID
-        query = "SELECT COUNT(*) FROM run WHERE run_id = ?;"
-        if self.query(query, (sint64_run_id,))[0]["COUNT(*)"] == 0:
+        query = "SELECT COUNT(*) FROM run WHERE run_id = ?"
+        rows = self.query(query, (sint64_run_id,))
+        if rows[0]["COUNT(*)"] == 0:
             return set()
 
-        # Get nodes
-        query = "SELECT node_id FROM node WHERE online_until > ?;"
-        rows = self.query(query, (time.time(),))
-
-        # Convert sint64 node_ids to uint64
-        result: set[int] = {convert_sint64_to_uint64(row["node_id"]) for row in rows}
-        return result
-
-    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
-        """Set `public_key` for the specified `node_id`."""
-        # Convert the uint64 value to sint64 for SQLite
-        sint64_node_id = convert_uint64_to_sint64(node_id)
-
-        # Check if the node exists in the `node` table
-        query = "SELECT 1 FROM node WHERE node_id = ?"
-        if not self.query(query, (sint64_node_id,)):
-            raise ValueError(f"Node {node_id} not found")
-
-        # Check if the public key is already in use in the `node` table
-        query = "SELECT 1 FROM node WHERE public_key = ?"
-        if self.query(query, (public_key,)):
-            raise ValueError("Public key already in use")
+        # Retrieve all online nodes
+        return {
+            node.node_id for node in self.get_node_info(statuses=[NodeStatus.ONLINE])
+        }
 
-        # Update the `node` table to set the public key for the given node ID
-        query = "UPDATE node SET public_key = ? WHERE node_id = ?"
-        self.query(query, (public_key, sint64_node_id))
+    def _check_and_tag_offline_nodes(
+        self, node_ids: Optional[list[int]] = None
+    ) -> None:
+        """Check and tag offline nodes."""
+        # strftime will convert POSIX timestamp to ISO format
+        query = """
+            UPDATE node SET status = ?,
+            last_deactivated_at =
+            strftime("%Y-%m-%dT%H:%M:%f+00:00", online_until, "unixepoch")
+            WHERE online_until <= ? AND status == ?
+        """
+        params = [
+            NodeStatus.OFFLINE,
+            now().timestamp(),
+            NodeStatus.ONLINE,
+        ]
+        if node_ids is not None:
+            placeholders = ",".join(["?"] * len(node_ids))
+            query += f" AND node_id IN ({placeholders})"
+            params.extend(uint64_to_int64(node_id) for node_id in node_ids)
+        self.conn.execute(query, params)
 
-    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
+    def get_node_info(
+        self,
+        *,
+        node_ids: Optional[Sequence[int]] = None,
+        owner_aids: Optional[Sequence[str]] = None,
+        statuses: Optional[Sequence[str]] = None,
+    ) -> Sequence[NodeInfo]:
+        """Retrieve information about nodes based on the specified filters."""
+        with self.conn:
+            self._check_and_tag_offline_nodes()
+
+            # Build the WHERE clause based on provided filters
+            conditions = []
+            params: list[Any] = []
+            if node_ids is not None:
+                sint64_node_ids = [uint64_to_int64(node_id) for node_id in node_ids]
+                placeholders = ",".join(["?"] * len(sint64_node_ids))
+                conditions.append(f"node_id IN ({placeholders})")
+                params.extend(sint64_node_ids)
+            if owner_aids is not None:
+                placeholders = ",".join(["?"] * len(owner_aids))
+                conditions.append(f"owner_aid IN ({placeholders})")
+                params.extend(owner_aids)
+            if statuses is not None:
+                placeholders = ",".join(["?"] * len(statuses))
+                conditions.append(f"status IN ({placeholders})")
+                params.extend(statuses)
+
+            # Construct the final query
+            query = "SELECT * FROM node"
+            if conditions:
+                query += " WHERE " + " AND ".join(conditions)
+
+            rows = self.conn.execute(query, params).fetchall()
+
+            result: list[NodeInfo] = []
+            for row in rows:
+                # Convert sint64 node_id to uint64
+                row["node_id"] = int64_to_uint64(row["node_id"])
+                result.append(NodeInfo(**row))
+
+            return result
+
+    def get_node_public_key(self, node_id: int) -> bytes:
         """Get `public_key` for the specified `node_id`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_node_id = convert_uint64_to_sint64(node_id)
+        sint64_node_id = uint64_to_int64(node_id)
 
         # Query the public key for the given node_id
-        query = "SELECT public_key FROM node WHERE node_id = ?"
-        rows = self.query(query, (sint64_node_id,))
+        query = "SELECT public_key FROM node WHERE node_id = ? AND status != ?;"
+        rows = self.query(query, (sint64_node_id, NodeStatus.UNREGISTERED))
 
         # If no result is found, return None
         if not rows:
-            raise ValueError(f"Node {node_id} not found")
+            raise ValueError(f"Node ID {node_id} not found")
 
-        # Return the public key if it is not empty, otherwise return None
-        return rows[0]["public_key"] or None
+        # Return the public key
+        return cast(bytes, rows[0]["public_key"])
 
-    def get_node_id(self, node_public_key: bytes) -> Optional[int]:
-        """Retrieve stored `node_id` filtered by `node_public_keys`."""
-        query = "SELECT node_id FROM node WHERE public_key = :public_key;"
-        row = self.query(query, {"public_key": node_public_key})
-        if len(row) > 0:
-            node_id: int = row[0]["node_id"]
+    def get_node_id_by_public_key(self, public_key: bytes) -> Optional[int]:
+        """Get `node_id` for the specified `public_key` if it exists and is not
+        deleted."""
+        query = "SELECT node_id FROM node WHERE public_key = ? AND status != ?;"
+        rows = self.query(query, (public_key, NodeStatus.UNREGISTERED))
 
-            # Convert the sint64 value to uint64 after reading from SQLite
-            uint64_node_id = convert_sint64_to_uint64(node_id)
+        # If no result is found, return None
+        if not rows:
+            return None
 
-            return uint64_node_id
-        return None
+        # Convert sint64 node_id to uint64
+        node_id = int64_to_uint64(rows[0]["node_id"])
+        return node_id
 
     # pylint: disable=too-many-arguments,too-many-positional-arguments
     def create_run(
@@ -737,7 +810,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         uint64_run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(uint64_run_id)
+        sint64_run_id = uint64_to_int64(uint64_run_id)
 
         # Check conflicts
         query = "SELECT COUNT(*) FROM run WHERE run_id = ?;"
@@ -773,28 +846,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
-    def clear_supernode_auth_keys(self) -> None:
-        """Clear stored `node_public_keys` in the link state if any."""
-        self.query("DELETE FROM public_key;")
-
-    def store_node_public_keys(self, public_keys: set[bytes]) -> None:
-        """Store a set of `node_public_keys` in the link state."""
-        query = "INSERT INTO public_key (public_key) VALUES (?)"
-        data = [(key,) for key in public_keys]
-        self.query(query, data)
-
-    def store_node_public_key(self, public_key: bytes) -> None:
-        """Store a `node_public_key` in the link state."""
-        query = "INSERT INTO public_key (public_key) VALUES (:public_key)"
-        self.query(query, {"public_key": public_key})
-
-    def get_node_public_keys(self) -> set[bytes]:
-        """Retrieve all currently stored `node_public_keys` as a set."""
-        query = "SELECT public_key FROM public_key"
-        rows = self.query(query)
-        result: set[bytes] = {row["public_key"] for row in rows}
-        return result
-
     def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
         """Retrieve all run IDs if `flwr_aid` is not specified.
 
@@ -807,7 +858,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             )
         else:
             rows = self.query("SELECT run_id FROM run;", ())
-        return {convert_sint64_to_uint64(row["run_id"]) for row in rows}
+        return {int64_to_uint64(row["run_id"]) for row in rows}
 
     def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:
         """Check if any runs are no longer active.
@@ -815,7 +866,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         Marks runs with status 'starting' or 'running' as failed
         if they have not sent a heartbeat before `active_until`.
         """
-        sint_run_ids = [convert_uint64_to_sint64(run_id) for run_id in run_ids]
+        sint_run_ids = [uint64_to_int64(run_id) for run_id in run_ids]
         query = "UPDATE run SET finished_at = ?, sub_status = ?, details = ? "
         query += "WHERE starting_at != '' AND finished_at = '' AND active_until < ?"
         query += f" AND run_id IN ({','.join(['?'] * len(run_ids))});"
@@ -837,13 +888,13 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids={run_id})
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint64_run_id,))
         if rows:
             row = rows[0]
             return Run(
-                run_id=convert_sint64_to_uint64(row["run_id"]),
+                run_id=int64_to_uint64(row["run_id"]),
                 fab_id=row["fab_id"],
                 fab_version=row["fab_version"],
                 fab_hash=row["fab_hash"],
@@ -868,13 +919,13 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids=run_ids)
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_ids = (convert_uint64_to_sint64(run_id) for run_id in set(run_ids))
+        sint64_run_ids = (uint64_to_int64(run_id) for run_id in set(run_ids))
         query = f"SELECT * FROM run WHERE run_id IN ({','.join(['?'] * len(run_ids))});"
         rows = self.query(query, tuple(sint64_run_ids))
 
         return {
             # Restore uint64 run IDs
-            convert_sint64_to_uint64(row["run_id"]): RunStatus(
+            int64_to_uint64(row["run_id"]): RunStatus(
                 status=determine_run_status(row),
                 sub_status=row["sub_status"],
                 details=row["details"],
@@ -888,7 +939,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids={run_id})
 
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint64_run_id,))
 
@@ -933,7 +984,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         # when switching to starting or running
         current = now()
         if new_status.status in (Status.STARTING, Status.RUNNING):
-            heartbeat_interval = HEARTBEAT_MAX_INTERVAL
+            heartbeat_interval = HEARTBEAT_INTERVAL_INF
             active_until = current.timestamp() + heartbeat_interval
         else:
             heartbeat_interval = 0
@@ -954,7 +1005,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             new_status.details,
             active_until,
             heartbeat_interval,
-            convert_uint64_to_sint64(run_id),
+            uint64_to_int64(run_id),
         )
         self.query(query % timestamp_fld, data)
         return True
@@ -967,14 +1018,14 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         query = "SELECT * FROM run WHERE starting_at = '' LIMIT 1;"
         rows = self.query(query)
         if rows:
-            pending_run_id = convert_sint64_to_uint64(rows[0]["run_id"])
+            pending_run_id = int64_to_uint64(rows[0]["run_id"])
 
         return pending_run_id
 
     def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
         """Retrieve the federation options for the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
         query = "SELECT federation_options FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint64_run_id,))
 
@@ -996,26 +1047,38 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         HEARTBEAT_PATIENCE = N allows for N-1 missed heartbeat before
         the node is marked as offline.
         """
-        sint64_node_id = convert_uint64_to_sint64(node_id)
+        if self.conn is None:
+            raise AttributeError("LinkState not initialized")
 
-        # Check if the node exists in the `node` table
-        query = "SELECT 1 FROM node WHERE node_id = ?"
-        if not self.query(query, (sint64_node_id,)):
-            return False
+        sint64_node_id = uint64_to_int64(node_id)
 
-        # Update `online_until` and `heartbeat_interval` for the given `node_id`
-        query = (
-            "UPDATE node SET online_until = ?, heartbeat_interval = ? WHERE node_id = ?"
-        )
-        self.query(
-            query,
-            (
-                time.time() + HEARTBEAT_PATIENCE * heartbeat_interval,
+        with self.conn:
+            # Check if node exists and not deleted
+            query = "SELECT status FROM node WHERE node_id = ? AND status != ?"
+            row = self.conn.execute(
+                query, (sint64_node_id, NodeStatus.UNREGISTERED)
+            ).fetchone()
+            if row is None:
+                return False
+
+            # Construct query and params
+            current_dt = now()
+            query = "UPDATE node SET online_until = ?, heartbeat_interval = ?"
+            params: list[Any] = [
+                current_dt.timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval,
                 heartbeat_interval,
-                sint64_node_id,
-            ),
-        )
-        return True
+            ]
+
+            # Set timestamp if the status changes
+            if row["status"] != NodeStatus.ONLINE:
+                query += ", status = ?, last_activated_at = ?"
+                params += [NodeStatus.ONLINE, current_dt.isoformat()]
+
+            # Execute the query, refreshing `online_until` and `heartbeat_interval`
+            query += " WHERE node_id = ?"
+            params += [sint64_node_id]
+            self.conn.execute(query, params)
+            return True
 
     def acknowledge_app_heartbeat(self, run_id: int, heartbeat_interval: float) -> bool:
         """Acknowledge a heartbeat received from a ServerApp for a given run.
@@ -1029,7 +1092,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self._check_and_tag_inactive_run(run_ids={run_id})
 
         # Search for the run
-        sint_run_id = convert_uint64_to_sint64(run_id)
+        sint_run_id = uint64_to_int64(run_id)
         query = "SELECT * FROM run WHERE run_id = ?;"
         rows = self.query(query, (sint_run_id,))
 
@@ -1059,7 +1122,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """Get the context for the specified `run_id`."""
         # Retrieve context if any
         query = "SELECT context FROM context WHERE run_id = ?;"
-        rows = self.query(query, (convert_uint64_to_sint64(run_id),))
+        rows = self.query(query, (uint64_to_int64(run_id),))
         context = context_from_bytes(rows[0]["context"]) if rows else None
         return context
 
@@ -1067,7 +1130,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """Set the context for the specified `run_id`."""
         # Convert context to bytes
         context_bytes = context_to_bytes(context)
-        sint_run_id = convert_uint64_to_sint64(run_id)
+        sint_run_id = uint64_to_int64(run_id)
 
         # Check if any existing Context assigned to the run_id
         query = "SELECT COUNT(*) FROM context WHERE run_id = ?;"
@@ -1086,7 +1149,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     def add_serverapp_log(self, run_id: int, log_message: str) -> None:
         """Add a log entry to the ServerApp logs for the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
 
         # Store log
         try:
@@ -1102,7 +1165,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     ) -> tuple[str, float]:
         """Get the ServerApp logs for the specified `run_id`."""
         # Convert the uint64 value to sint64 for SQLite
-        sint64_run_id = convert_uint64_to_sint64(run_id)
+        sint64_run_id = uint64_to_int64(run_id)
 
         # Check if the run_id exists
         query = "SELECT run_id FROM run WHERE run_id = ?;"
@@ -1140,7 +1203,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         message_ins = rows[0]
         created_at = message_ins["created_at"]
         ttl = message_ins["ttl"]
-        current_time = time.time()
+        current_time = now().timestamp()
 
         # Check if Message is expired
         if ttl is not None and created_at + ttl <= current_time:
@@ -1152,7 +1215,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """Create a token for the given run ID."""
         token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
         query = "INSERT INTO token_store (run_id, token) VALUES (:run_id, :token);"
-        data = {"run_id": convert_uint64_to_sint64(run_id), "token": token}
+        data = {"run_id": uint64_to_int64(run_id), "token": token}
         try:
             self.query(query, data)
         except sqlite3.IntegrityError:
@@ -1162,7 +1225,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     def verify_token(self, run_id: int, token: str) -> bool:
         """Verify a token for the given run ID."""
         query = "SELECT token FROM token_store WHERE run_id = :run_id;"
-        data = {"run_id": convert_uint64_to_sint64(run_id)}
+        data = {"run_id": uint64_to_int64(run_id)}
         rows = self.query(query, data)
         if not rows:
             return False
@@ -1171,7 +1234,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
     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": convert_uint64_to_sint64(run_id)}
+        data = {"run_id": uint64_to_int64(run_id)}
         self.query(query, data)
 
     def get_run_id_by_token(self, token: str) -> Optional[int]:
@@ -1181,19 +1244,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         rows = self.query(query, data)
         if not rows:
             return None
-        return convert_sint64_to_uint64(rows[0]["run_id"])
-
-
-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))
+        return int64_to_uint64(rows[0]["run_id"])
 
 
 def message_to_dict(message: Message) -> dict[str, Any]:
@@ -1248,5 +1299,5 @@ def determine_run_status(row: dict[str, Any]) -> str:
                 return Status.RUNNING
             return Status.STARTING
         return Status.PENDING
-    run_id = convert_sint64_to_uint64(row["run_id"])
+    run_id = int64_to_uint64(row["run_id"])
     raise sqlite3.IntegrityError(f"The run {run_id} does not have a valid status.")