feat: Add load_usage_logs parameter to node retrieval functions and update related calls

Author: M03ED

app/db/crud/admin.py

@@ -530,9 +530,7 @@ async def get_usage_percentage_reached_admins(
     return list((await db.execute(stmt)).scalars().all())
 
 
-async def bulk_create_admin_notification_reminders(
-    db: AsyncSession, reminder_data: list[dict]
-) -> list[dict]:
+async def bulk_create_admin_notification_reminders(db: AsyncSession, reminder_data: list[dict]) -> list[dict]:
     """Bulk-insert admin reminder rows after successful sends."""
     if not reminder_data:
         return []
@@ -550,11 +548,7 @@ async def bulk_create_admin_notification_reminders(
     types = {d["type"] for d in unique_reminder_data}
 
     # Lock the Admin rows to serialize reminder checks/creation for these admins
-    await db.execute(
-        select(Admin.id)
-        .where(Admin.id.in_(list(admin_ids)))
-        .with_for_update()
-    )
+    await db.execute(select(Admin.id).where(Admin.id.in_(list(admin_ids))).with_for_update())
 
     # Fetch existing reminders that match these criteria
     stmt = select(AdminNotificationReminder).where(

app/db/crud/node.py

@@ -46,14 +46,16 @@ def _build_node_simple_sort_clause(sort_option: NodeSimpleSortOption):
     return column.desc() if sort_option.value.startswith("-") else column.asc()
 
 
-async def load_node_attrs(node: Node):
+async def load_node_attrs(node: Node, *, load_usage_logs: bool = True):
+    if not load_usage_logs:
+        return
     try:
         await node.awaitable_attrs.usage_logs
     except AttributeError:
         pass
 
 
-async def get_node(db: AsyncSession, name: str) -> Node | None:
+async def get_node(db: AsyncSession, name: str, *, load_usage_logs: bool = True) -> Node | None:
     """
     Retrieves a node by its name.
 
@@ -66,11 +68,11 @@ async def get_node(db: AsyncSession, name: str) -> Node | None:
     """
     node = (await db.execute(select(Node).where(Node.name == name))).unique().scalar_one_or_none()
     if node:
-        await load_node_attrs(node)
+        await load_node_attrs(node, load_usage_logs=load_usage_logs)
     return node
 
 
-async def get_node_by_id(db: AsyncSession, node_id: int) -> Node | None:
+async def get_node_by_id(db: AsyncSession, node_id: int, *, load_usage_logs: bool = True) -> Node | None:
     """
     Retrieves a node by its ID.
 
@@ -83,13 +85,15 @@ async def get_node_by_id(db: AsyncSession, node_id: int) -> Node | None:
     """
     node = (await db.execute(select(Node).where(Node.id == node_id))).unique().scalar_one_or_none()
     if node:
-        await load_node_attrs(node)
+        await load_node_attrs(node, load_usage_logs=load_usage_logs)
     return node
 
 
 async def get_nodes(
     db: AsyncSession,
     query: NodeListQuery,
+    *,
+    load_usage_logs: bool = True,
 ) -> tuple[list[Node], int]:
     """
     Retrieves nodes based on optional status, enabled, id, and search filters.
@@ -143,8 +147,9 @@ async def get_nodes(
     # Order by created_at and id for consistent results
     stmt = stmt.order_by(Node.created_at.asc(), Node.id.asc())
 
-    # Eagerly load usage_logs to avoid N+1 queries (one extra SELECT per node)
-    stmt = stmt.options(selectinload(Node.usage_logs))
+    # Eagerly load usage_logs for API lifetime_* fields (skip for jobs/connect)
+    if load_usage_logs:
+        stmt = stmt.options(selectinload(Node.usage_logs))
 
     db_nodes = (await db.execute(stmt)).unique().scalars().all()
 
@@ -208,7 +213,7 @@ async def get_limited_nodes(db: AsyncSession) -> list[Node]:
     Returns:
         list[Node]: Nodes that should be limited
     """
-    query = select(Node).options(selectinload(Node.usage_logs)).where(
+    query = select(Node).where(
         and_(
             Node.status.in_([NodeStatus.error, NodeStatus.connected, NodeStatus.connecting]),
             Node.is_limited,
@@ -503,7 +508,6 @@ async def update_node_status(
         # If the instance was detached (e.g., used across sessions), re-fetch it
         db_node = (await db.execute(select(Node).where(Node.id == db_node.id))).scalar_one()
 
-    await load_node_attrs(db_node)
     return db_node
 
 
@@ -799,10 +803,11 @@ async def bulk_reset_node_usage(db: AsyncSession, nodes: list[Node]) -> list[Nod
     # Re-fetch all nodes in a single query instead of N individual refreshes
     node_ids = [node.id for node in nodes]
     refreshed = (
-        await db.execute(
-            select(Node).options(selectinload(Node.usage_logs)).where(Node.id.in_(node_ids))
-        )
-    ).unique().scalars().all()
+        (await db.execute(select(Node).options(selectinload(Node.usage_logs)).where(Node.id.in_(node_ids))))
+        .unique()
+        .scalars()
+        .all()
+    )
     # Preserve input order
     refreshed_by_id = {n.id: n for n in refreshed}
     return [refreshed_by_id[nid] for nid in node_ids if nid in refreshed_by_id]

app/db/models.py

@@ -751,7 +751,9 @@ class AdminNotificationReminder(Base, CreatedAtUTCMixin):
     __tablename__ = "admin_notification_reminders"
     __table_args__ = (
         Index("ix_admin_notification_reminders_admin_id_type", "admin_id", "type"),
-        UniqueConstraint("admin_id", "type", "threshold", name="uq_admin_notification_reminders_admin_id_type_threshold"),
+        UniqueConstraint(
+            "admin_id", "type", "threshold", name="uq_admin_notification_reminders_admin_id_type_threshold"
+        ),
     )
     admin_id: Mapped[int] = fk_id_column("admins.id", ondelete="CASCADE")
     admin: Mapped[Admin] = relationship(back_populates="notification_reminders", init=False)

app/jobs/node_checker.py

@@ -209,7 +209,7 @@ async def node_health_check():
     if not runtime_settings.role.runs_node:
         return
     async with GetDB() as db:
-        db_nodes, _ = await get_nodes(db=db, query=NodeListQuery(status=ACTIVE_NODE_STATUSES))
+        db_nodes, _ = await get_nodes(db=db, query=NodeListQuery(status=ACTIVE_NODE_STATUSES), load_usage_logs=False)
 
     dict_nodes = await node_manager.get_nodes()
     check_tasks = [process_node_health_check(db_node, dict_nodes.get(db_node.id)) for db_node in db_nodes]
@@ -224,7 +224,7 @@ async def initialize_nodes():
     logger.info("Starting nodes' cores...")
 
     async with GetDB() as db:
-        db_nodes, _ = await get_nodes(db=db, query=NodeListQuery(status=ACTIVE_NODE_STATUSES))
+        db_nodes, _ = await get_nodes(db=db, query=NodeListQuery(status=ACTIVE_NODE_STATUSES), load_usage_logs=False)
 
         if not db_nodes:
             logger.warning("Attention: You have no node, you need to have at least one node")

app/jobs/review_admins.py

@@ -79,16 +79,11 @@ async def _send_usage_limit_warning_notifications(db):
         candidate_ids = [admin.id for admin in candidate_admins]
 
         # Lock the Admin rows to serialize reminder checks/creation for these admins
-        await db.execute(
-            select(Admin.id)
-            .where(Admin.id.in_(candidate_ids))
-            .with_for_update()
-        )
+        await db.execute(select(Admin.id).where(Admin.id.in_(candidate_ids)).with_for_update())
 
         # Fetch existing reminders for this threshold to avoid duplicate notifications
         result = await db.execute(
-            select(AdminNotificationReminder.admin_id)
-            .where(
+            select(AdminNotificationReminder.admin_id).where(
                 AdminNotificationReminder.admin_id.in_(candidate_ids),
                 AdminNotificationReminder.type == ReminderType.data_usage,
                 AdminNotificationReminder.threshold == threshold,

app/node/worker.py

@@ -154,7 +154,7 @@ async def _update_node(self, data: dict):
         if not node_id:
             return
         async with GetDB() as db:
-            db_node = await get_node_by_id(db, node_id)
+            db_node = await get_node_by_id(db, node_id, load_usage_logs=False)
         if db_node:
             await node_manager.update_node(db_node)
 
@@ -180,14 +180,15 @@ async def _connect_nodes_bulk(self, data: dict):
         await core_manager._reload_from_cache()
         async with GetDB() as db:
             if node_ids:
-                nodes, _ = await get_nodes(db, query=NodeListQuery(ids=node_ids))
+                nodes, _ = await get_nodes(db, query=NodeListQuery(ids=node_ids), load_usage_logs=False)
             else:
                 nodes, _ = await get_nodes(
                     db,
                     query=NodeListQuery(
                         core_id=core_id,
                         status=[NodeStatus.connected, NodeStatus.connecting, NodeStatus.error],
                     ),
+                    load_usage_logs=False,
                 )
             await self._node_operator.connect_nodes_bulk(db, nodes)
 

app/operation/__init__.py

@@ -254,9 +254,9 @@ async def get_validated_user_template(self, db: AsyncSession, template_id: int)
             await self.raise_error("User Template not found", 404)
         return dbuser_template
 
-    async def get_validated_node(self, db: AsyncSession, node_id) -> Node:
+    async def get_validated_node(self, db: AsyncSession, node_id, *, load_usage_logs: bool = True) -> Node:
         """Dependency: Fetch node or return not found error."""
-        db_node = await get_node_by_id(db, node_id)
+        db_node = await get_node_by_id(db, node_id, load_usage_logs=load_usage_logs)
         if not db_node:
             await self.raise_error(message="Node not found", code=404)
         return db_node

app/operation/node.py

@@ -154,7 +154,7 @@ async def _update_single_node_status(
             node_version (str): Node version.
             send_notification (bool): Whether to send notification.
         """
-        db_node = await get_node_by_id(db, node_id)
+        db_node = await get_node_by_id(db, node_id, load_usage_logs=False)
         if not db_node:
             return
 
@@ -540,15 +540,15 @@ async def clear_usage_data(self, db: AsyncSession, table: UsageTable, query: Nod
             await self.raise_error(code=400, message=f"Deletion failed due to server error: {e!s}")
 
     async def update_node(self, db: AsyncSession, node_id: int) -> dict:
-        await self.get_validated_node(db, node_id)
+        await self.get_validated_node(db, node_id, load_usage_logs=False)
         return await self._update_node_api_impl(node_id)
 
     async def update_core(self, db: AsyncSession, node_id: int, node_core_update: NodeCoreUpdate) -> dict:
-        await self.get_validated_node(db, node_id)
+        await self.get_validated_node(db, node_id, load_usage_logs=False)
         return await self._update_core_impl(node_id, node_core_update)
 
     async def update_geofiles(self, db: AsyncSession, node_id: int, node_geofiles_update: NodeGeoFilesUpdate) -> dict:
-        await self.get_validated_node(db, node_id)
+        await self.get_validated_node(db, node_id, load_usage_logs=False)
         return await self._update_geofiles_impl(node_id, node_geofiles_update)
 
     async def _update_node_local(self, db_node: Node) -> None:
@@ -635,7 +635,7 @@ async def _connect_nodes_bulk_remote(self, db: AsyncSession, nodes: list[Node])
         await node_nats_client.publish("connect_nodes_bulk", {"node_ids": [node.id for node in nodes]})
 
     async def _connect_single_node_local(self, db: AsyncSession, node_id: int) -> None:
-        db_node = await get_node_by_id(db, node_id)
+        db_node = await get_node_by_id(db, node_id, load_usage_logs=False)
         if db_node is None or db_node.status in (NodeStatus.disabled, NodeStatus.limited):
             return
 
@@ -714,6 +714,7 @@ async def _restart_all_nodes_local(self, db: AsyncSession, admin: AdminDetails,
                 core_id=core_id,
                 status=[NodeStatus.connected, NodeStatus.connecting, NodeStatus.error],
             ),
+            load_usage_logs=False,
         )
         await self.connect_nodes_bulk(db, nodes)
 
@@ -1011,7 +1012,7 @@ async def _get_validated_nodes(self, db: AsyncSession, node_ids: list[int] | set
             return []
 
         ids_list = list(node_ids)
-        db_nodes, _ = await get_nodes(db, NodeListQuery(ids=ids_list, limit=len(ids_list)))
+        db_nodes, _ = await get_nodes(db, NodeListQuery(ids=ids_list, limit=len(ids_list)), load_usage_logs=False)
 
         found_ids = {n.id for n in db_nodes}
         missing = set(ids_list) - found_ids

tests/api/test_bulk.py

@@ -127,9 +127,7 @@ def test_bulk_add_wireguard_group_allocates_existing_user_peer_ips(access_token)
             type="wg",
             fallbacks=[],
         )
-        wg_group = create_group(
-            access_token, name=unique_name("wg_bulk_add_group"), inbound_tags=[interface_name]
-        )
+        wg_group = create_group(access_token, name=unique_name("wg_bulk_add_group"), inbound_tags=[interface_name])
 
         response = client.post(
             "/api/groups/bulk/add",

tests/api/test_host.py

@@ -371,4 +371,3 @@ def test_host_finalmask_new_types(access_token):
     finally:
         client.delete(f"/api/host/{host_id}", headers={"Authorization": f"Bearer {access_token}"})
         delete_core(access_token, core["id"])
-