Nexus Agent
202 lines
7.4 KiB
Python
202 lines
7.4 KiB
Python
"""Nexus — Server Repository (Async SQLAlchemy implementation)
|
|
Concrete implementation of ServerRepository protocol.
|
|
"""
|
|
|
|
from datetime import datetime, timezone
|
|
from typing import Optional, List, Tuple
|
|
from sqlalchemy import select, update, func, asc, desc, or_
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from server.domain.models import Server
|
|
|
|
UNCATEGORIZED_CATEGORY = "uncategorized"
|
|
|
|
|
|
def _category_filter(category: str):
|
|
"""Map API category filter to SQLAlchemy criterion."""
|
|
if category == UNCATEGORIZED_CATEGORY:
|
|
return or_(Server.category.is_(None), Server.category == "")
|
|
return Server.category == category
|
|
|
|
|
|
# Whitelist of sortable columns (prevents SQL injection via sort_by)
|
|
_SORT_COLUMNS = {
|
|
"id": Server.id,
|
|
"name": Server.name,
|
|
"domain": Server.domain,
|
|
"status": Server.is_online,
|
|
"heartbeat": Server.last_heartbeat,
|
|
"agent_version": Server.agent_version,
|
|
"category": Server.category,
|
|
"target_path": Server.target_path,
|
|
}
|
|
|
|
|
|
class ServerRepositoryImpl:
|
|
"""Async SQLAlchemy implementation of Server data access"""
|
|
|
|
def __init__(self, session: AsyncSession):
|
|
self.session = session
|
|
|
|
async def get_by_id(self, id: int) -> Optional[Server]:
|
|
result = await self.session.execute(select(Server).where(Server.id == id))
|
|
return result.scalar_one_or_none()
|
|
|
|
async def get_by_ids(self, ids: List[int]) -> List[Server]:
|
|
"""Fetch multiple servers in one query; order matches *ids* (skips missing)."""
|
|
if not ids:
|
|
return []
|
|
result = await self.session.execute(select(Server).where(Server.id.in_(ids)))
|
|
by_id = {s.id: s for s in result.scalars().all()}
|
|
return [by_id[i] for i in ids if i in by_id]
|
|
|
|
async def get_all(self) -> List[Server]:
|
|
result = await self.session.execute(select(Server).order_by(Server.id))
|
|
return list(result.scalars().all())
|
|
|
|
async def get_paginated(
|
|
self,
|
|
category: Optional[str] = None,
|
|
platform_id: Optional[int] = None,
|
|
node_id: Optional[int] = None,
|
|
offset: int = 0,
|
|
limit: int = 50,
|
|
search: Optional[str] = None,
|
|
sort_by: Optional[str] = None,
|
|
sort_order: Optional[str] = "asc",
|
|
is_online: Optional[bool] = None,
|
|
) -> Tuple[List[Server], int]:
|
|
"""Get paginated server list with total count.
|
|
|
|
Supports text search (name/domain), sorting, and online status filter.
|
|
Returns (servers, total_count) — DB-level pagination for 2000+ servers.
|
|
When is_online filter is active, we over-fetch (limit + 10) to compensate
|
|
for Redis overlay post-filter corrections in the API layer.
|
|
"""
|
|
filters = []
|
|
if category:
|
|
filters.append(_category_filter(category))
|
|
if platform_id:
|
|
filters.append(Server.platform_id == platform_id)
|
|
if node_id:
|
|
filters.append(Server.node_id == node_id)
|
|
if search:
|
|
term = f"%{search}%"
|
|
filters.append(Server.name.ilike(term) | Server.domain.ilike(term))
|
|
if is_online is not None:
|
|
filters.append(Server.is_online == is_online)
|
|
|
|
# Count query
|
|
count_query = select(func.count(Server.id))
|
|
for f in filters:
|
|
count_query = count_query.where(f)
|
|
count_result = await self.session.execute(count_query)
|
|
total = count_result.scalar_one() or 0
|
|
|
|
# Over-fetch when is_online filter is active (Redis overlay may change status)
|
|
fetch_limit = limit + 10 if is_online is not None else limit
|
|
|
|
# Paginated data query with sort
|
|
data_query = select(Server)
|
|
for f in filters:
|
|
data_query = data_query.where(f)
|
|
|
|
# Sorting (whitelist-validated)
|
|
sort_col = _SORT_COLUMNS.get(sort_by, Server.id)
|
|
order_func = desc if sort_order == "desc" else asc
|
|
data_query = data_query.order_by(order_func(sort_col), Server.id)
|
|
data_query = data_query.offset(offset).limit(fetch_limit)
|
|
result = await self.session.execute(data_query)
|
|
servers = list(result.scalars().all())
|
|
|
|
return servers, total
|
|
|
|
async def get_filtered(
|
|
self,
|
|
category: Optional[str] = None,
|
|
platform_id: Optional[int] = None,
|
|
node_id: Optional[int] = None,
|
|
search: Optional[str] = None,
|
|
is_online: Optional[bool] = None,
|
|
max_rows: int = 5000,
|
|
) -> Tuple[List[Server], int]:
|
|
"""All servers matching filters (no sort/pagination), capped at *max_rows*."""
|
|
filters = []
|
|
if category:
|
|
filters.append(_category_filter(category))
|
|
if platform_id:
|
|
filters.append(Server.platform_id == platform_id)
|
|
if node_id:
|
|
filters.append(Server.node_id == node_id)
|
|
if search:
|
|
term = f"%{search}%"
|
|
filters.append(Server.name.ilike(term) | Server.domain.ilike(term))
|
|
if is_online is not None:
|
|
filters.append(Server.is_online == is_online)
|
|
|
|
count_query = select(func.count(Server.id))
|
|
for f in filters:
|
|
count_query = count_query.where(f)
|
|
count_result = await self.session.execute(count_query)
|
|
total = count_result.scalar_one() or 0
|
|
|
|
data_query = select(Server)
|
|
for f in filters:
|
|
data_query = data_query.where(f)
|
|
data_query = data_query.order_by(Server.id).limit(max_rows)
|
|
result = await self.session.execute(data_query)
|
|
servers = list(result.scalars().all())
|
|
return servers, total
|
|
|
|
async def get_by_category(self, category: str) -> List[Server]:
|
|
result = await self.session.execute(
|
|
select(Server).where(Server.category == category).order_by(Server.id)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
async def get_online(self) -> List[Server]:
|
|
result = await self.session.execute(
|
|
select(Server).where(Server.is_online == True).order_by(Server.id)
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
async def create(self, server: Server) -> Server:
|
|
self.session.add(server)
|
|
await self.session.commit()
|
|
return server
|
|
|
|
async def update(self, server: Server) -> Server:
|
|
await self.session.commit()
|
|
return server
|
|
|
|
async def update_category_batch(self, ids: List[int], category: Optional[str]) -> int:
|
|
"""Set category for multiple servers in one UPDATE."""
|
|
if not ids:
|
|
return 0
|
|
result = await self.session.execute(
|
|
update(Server).where(Server.id.in_(ids)).values(category=category)
|
|
)
|
|
await self.session.commit()
|
|
return int(result.rowcount or 0)
|
|
|
|
async def delete(self, id: int) -> bool:
|
|
server = await self.get_by_id(id)
|
|
if server:
|
|
await self.session.delete(server)
|
|
await self.session.commit()
|
|
return True
|
|
return False
|
|
|
|
async def update_heartbeat(self, id: int, is_online: bool, system_info: str, agent_version: str) -> None:
|
|
"""Update server online status from Agent heartbeat"""
|
|
await self.session.execute(
|
|
update(Server)
|
|
.where(Server.id == id)
|
|
.values(
|
|
is_online=is_online,
|
|
last_heartbeat=datetime.now(timezone.utc),
|
|
system_info=system_info,
|
|
agent_version=agent_version,
|
|
)
|
|
)
|
|
await self.session.commit() |