add server api handlers. Fix bug with cmd output

2026-06-24 13:49:12 +03:00
parent fe2db1b6e9
commit 7ebfea097d
12 changed files with 612 additions and 25 deletions
@@ -246,22 +246,25 @@ def main(argv: list[str] | None = None) -> int:
    try:
        config = load_client_config()
        result = run_once(config)
+        if result is None:
+            logger.info("heartbeat accepted; no command")
+            return 0
+
+        while result is not None:
+            logger.info(
+                "command result: %s",
+                json.dumps(
+                    result.model_dump(mode="json"),
+                    ensure_ascii=False,
+                    separators=(",", ":"),
+                ),
+            )
+            result = run_once(config, last_command_result=result)
+            logger.info("command result reported to server")
    except (ClientConfigError, HeartbeatError, ValueError) as error:
        logger.error("nexus-sync client error: %s", error)
        return 1

-    if result is None:
-        logger.info("heartbeat accepted; no command")
-        return 0
-
-    logger.info(
-        "command result: %s",
-        json.dumps(
-            result.model_dump(mode="json"),
-            ensure_ascii=False,
-            separators=(",", ":"),
-        ),
-    )
    return 0


@@ -1,11 +1,13 @@
 from nexus_sync.server.app import create_app
 from nexus_sync.server.config import DEFAULT_COMMAND_POLL_SECONDS, DEFAULT_IDLE_POLL_SECONDS
+from nexus_sync.server.sqlalchemy_store import SQLAlchemyStore
 from nexus_sync.server.store import InMemoryStore, Store

 __all__ = [
    "DEFAULT_COMMAND_POLL_SECONDS",
    "DEFAULT_IDLE_POLL_SECONDS",
    "InMemoryStore",
+    "SQLAlchemyStore",
    "Store",
    "create_app",
 ]
@@ -1,16 +1,27 @@
 from datetime import UTC, datetime
+from typing import Any
+from uuid import uuid4

 from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
+from pydantic import BaseModel, Field

-from nexus_sync.common import HeartbeatRequest, HeartbeatResponse
+from nexus_sync.common import Command, CommandKind, HeartbeatRequest, HeartbeatResponse
 from nexus_sync.server.config import (
    DEFAULT_COMMAND_POLL_SECONDS,
    DEFAULT_IDLE_POLL_SECONDS,
    load_client_tokens,
+    load_database_url,
 )
-from nexus_sync.server.store import InMemoryStore, Store
+from nexus_sync.server.sqlalchemy_store import SQLAlchemyStore
+from nexus_sync.server.store import Store
+
+
+class CommandCreateRequest(BaseModel):
+    name: str
+    args: dict[str, Any] = Field(default_factory=dict)
+    timeout_seconds: int = Field(default=30, gt=0)


 def create_app(
@@ -19,7 +30,7 @@ def create_app(
    client_tokens: dict[str, str] | None = None,
 ) -> FastAPI:
    app = FastAPI(title="nexus-sync", version="0.1.0")
-    app.state.store = store or InMemoryStore()
+    app.state.store = store or SQLAlchemyStore(load_database_url())
    app.state.client_tokens = client_tokens if client_tokens is not None else load_client_tokens()

    @app.exception_handler(RequestValidationError)
@@ -71,6 +82,43 @@ def create_app(
            command=command,
        )

+    @app.get("/api/v1/server/clients")
+    def list_clients() -> dict[str, list[dict[str, Any]]]:
+        return {
+            "clients": [client.model_dump(mode="json") for client in app.state.store.list_clients()]
+        }
+
+    @app.get("/api/v1/server/clients/{client_id}")
+    def get_client(client_id: str) -> dict[str, Any]:
+        client = app.state.store.get_client(client_id)
+        if client is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="client not found")
+        return client.model_dump(mode="json")
+
+    @app.post("/api/v1/server/clients/{client_id}/commands", status_code=status.HTTP_201_CREATED)
+    def create_command(client_id: str, payload: CommandCreateRequest) -> dict[str, Any]:
+        if app.state.store.get_client(client_id) is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="client not found")
+
+        return _enqueue_command(
+            store=app.state.store,
+            client_id=client_id,
+            name=payload.name,
+            args=payload.args,
+            timeout_seconds=payload.timeout_seconds,
+        )
+
+    @app.get("/api/v1/server/commands/{command_id}")
+    def get_command(command_id: str) -> dict[str, Any]:
+        command = app.state.store.get_command(command_id)
+        if command is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="command not found")
+
+        body = command.model_dump(mode="json")
+        result = app.state.store.get_command_result(command_id)
+        body["result"] = result.model_dump(mode="json") if result else None
+        return body
+
    return app


@@ -81,4 +129,24 @@ def _client_id_for_token(client_tokens: dict[str, str], token: str) -> str | Non
    return None


+def _enqueue_command(
+    *,
+    store: Store,
+    client_id: str,
+    name: str,
+    args: dict[str, Any],
+    timeout_seconds: int,
+) -> dict[str, Any]:
+    now = datetime.now(UTC)
+    command = Command(
+        id=f"cmd_{uuid4().hex}",
+        kind=CommandKind.EXEC,
+        name=name,
+        args=args,
+        timeout_seconds=timeout_seconds,
+    )
+    record = store.enqueue_command(command, client_id, now)
+    return record.model_dump(mode="json")
+
+
 app = create_app()
@@ -2,6 +2,11 @@ import os

 DEFAULT_IDLE_POLL_SECONDS = 60
 DEFAULT_COMMAND_POLL_SECONDS = 10
+DEFAULT_DATABASE_URL = "sqlite:///nexus-sync.db"
+
+
+def load_database_url() -> str:
+    return os.environ.get("NEXUS_SYNC_DATABASE_URL", DEFAULT_DATABASE_URL)


 def load_client_tokens() -> dict[str, str]:
@@ -0,0 +1,205 @@
+from datetime import datetime
+
+from fastapi import HTTPException, status
+from sqlalchemy import Column, MetaData, String, Table, create_engine, select
+from sqlalchemy.engine import Engine
+
+from nexus_sync.common import (
+    ClientRecord,
+    Command,
+    CommandRecord,
+    CommandResultRecord,
+    CommandStatus,
+    HeartbeatRequest,
+)
+from nexus_sync.server.store import TERMINAL_COMMAND_STATUSES
+
+
+class SQLAlchemyStore:
+    def __init__(self, database_url: str = "sqlite:///nexus-sync.db") -> None:
+        self.engine = create_engine(database_url)
+        self.metadata = MetaData()
+        self.clients = Table(
+            "clients",
+            self.metadata,
+            Column("id", String, primary_key=True),
+            Column("json", String, nullable=False),
+        )
+        self.commands = Table(
+            "commands",
+            self.metadata,
+            Column("id", String, primary_key=True),
+            Column("client_id", String, nullable=False, index=True),
+            Column("status", String, nullable=False, index=True),
+            Column("created_at", String, nullable=False, index=True),
+            Column("json", String, nullable=False),
+        )
+        self.results = Table(
+            "command_results",
+            self.metadata,
+            Column("command_id", String, primary_key=True),
+            Column("json", String, nullable=False),
+        )
+        self.metadata.create_all(self.engine)
+
+    def list_clients(self) -> list[ClientRecord]:
+        with self.engine.begin() as connection:
+            rows = connection.execute(select(self.clients.c.json).order_by(self.clients.c.id)).all()
+        return [ClientRecord.model_validate_json(row.json) for row in rows]
+
+    def get_client(self, client_id: str) -> ClientRecord | None:
+        with self.engine.begin() as connection:
+            row = connection.execute(
+                select(self.clients.c.json).where(self.clients.c.id == client_id)
+            ).first()
+        return ClientRecord.model_validate_json(row.json) if row else None
+
+    def upsert_client_from_payload(self, payload: dict, now: datetime) -> ClientRecord:
+        return self.upsert_client(HeartbeatRequest.model_validate(payload), now)
+
+    def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord:
+        existing = self.get_client(heartbeat.client_id)
+        created_at = existing.created_at if existing else now
+        token_hash = existing.token_hash if existing else None
+        record = ClientRecord(
+            id=heartbeat.client_id,
+            hostname=heartbeat.client.hostname,
+            platform=heartbeat.client.platform,
+            version=heartbeat.client.version,
+            created_at=created_at,
+            last_seen_at=now,
+            is_active=True,
+            token_hash=token_hash,
+            available_commands=heartbeat.available_commands,
+        )
+        with self.engine.begin() as connection:
+            connection.execute(self.clients.delete().where(self.clients.c.id == record.id))
+            connection.execute(
+                self.clients.insert().values(id=record.id, json=record.model_dump_json())
+            )
+        return record
+
+    def enqueue_command(self, command: Command, client_id: str, now: datetime) -> CommandRecord:
+        record = CommandRecord(
+            id=command.id,
+            client_id=client_id,
+            kind=command.kind,
+            name=command.name,
+            args=command.args,
+            status=CommandStatus.PENDING,
+            timeout_seconds=command.timeout_seconds,
+            created_at=now,
+        )
+        self._write_command(record)
+        return record
+
+    def record_command_result_from_payload(self, payload: dict, now: datetime) -> None:
+        self.record_command_result(HeartbeatRequest.model_validate(payload), now)
+
+    def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None:
+        result = heartbeat.last_command_result
+        if result is None:
+            return
+
+        command = self.get_command(result.command_id)
+        if command is None:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail="command result references unknown command",
+            )
+        if command.client_id != heartbeat.client_id:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail="command result references command for another client",
+            )
+        if command.status in TERMINAL_COMMAND_STATUSES:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail="command result references terminal command",
+            )
+
+        finished_at = result.finished_at or now
+        self._write_command(
+            command.model_copy(
+                update={
+                    "status": CommandStatus(result.status.value),
+                    "finished_at": finished_at,
+                }
+            )
+        )
+        record = CommandResultRecord(
+            command_id=result.command_id,
+            client_id=heartbeat.client_id,
+            status=result.status,
+            started_at=result.started_at,
+            finished_at=finished_at,
+            return_code=result.return_code,
+            stdout=result.stdout,
+            stderr=result.stderr,
+            received_at=now,
+        )
+        with self.engine.begin() as connection:
+            connection.execute(
+                self.results.delete().where(self.results.c.command_id == record.command_id)
+            )
+            connection.execute(
+                self.results.insert().values(
+                    command_id=record.command_id,
+                    json=record.model_dump_json(),
+                )
+            )
+
+    def take_next_command(self, client_id: str, now: datetime) -> Command | None:
+        with self.engine.begin() as connection:
+            rows = connection.execute(
+                select(self.commands.c.json)
+                .where(self.commands.c.client_id == client_id)
+                .where(self.commands.c.status == CommandStatus.PENDING.value)
+                .order_by(self.commands.c.created_at)
+            ).all()
+        if not rows:
+            return None
+
+        record = CommandRecord.model_validate_json(rows[0].json)
+        delivered = record.model_copy(
+            update={
+                "status": CommandStatus.DELIVERED,
+                "attempts": record.attempts + 1,
+                "delivered_at": now,
+            }
+        )
+        self._write_command(delivered)
+        return Command(
+            id=delivered.id,
+            kind=delivered.kind,
+            name=delivered.name,
+            args=delivered.args,
+            timeout_seconds=delivered.timeout_seconds,
+        )
+
+    def get_command(self, command_id: str) -> CommandRecord | None:
+        with self.engine.begin() as connection:
+            row = connection.execute(
+                select(self.commands.c.json).where(self.commands.c.id == command_id)
+            ).first()
+        return CommandRecord.model_validate_json(row.json) if row else None
+
+    def get_command_result(self, command_id: str) -> CommandResultRecord | None:
+        with self.engine.begin() as connection:
+            row = connection.execute(
+                select(self.results.c.json).where(self.results.c.command_id == command_id)
+            ).first()
+        return CommandResultRecord.model_validate_json(row.json) if row else None
+
+    def _write_command(self, record: CommandRecord) -> None:
+        with self.engine.begin() as connection:
+            connection.execute(self.commands.delete().where(self.commands.c.id == record.id))
+            connection.execute(
+                self.commands.insert().values(
+                    id=record.id,
+                    client_id=record.client_id,
+                    status=record.status.value,
+                    created_at=record.created_at.isoformat(),
+                    json=record.model_dump_json(),
+                )
+            )
@@ -15,15 +15,30 @@ from nexus_sync.common import (


 class Store(Protocol):
+    def list_clients(self) -> list[ClientRecord]:
+        pass
+
+    def get_client(self, client_id: str) -> ClientRecord | None:
+        pass
+
    def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord:
        pass

+    def enqueue_command(self, command: Command, client_id: str, now: datetime) -> CommandRecord:
+        pass
+
    def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None:
        pass

    def take_next_command(self, client_id: str, now: datetime) -> Command | None:
        pass

+    def get_command(self, command_id: str) -> CommandRecord | None:
+        pass
+
+    def get_command_result(self, command_id: str) -> CommandResultRecord | None:
+        pass
+

@dataclass
 class InMemoryStore:
@@ -31,6 +46,12 @@ class InMemoryStore:
    commands: dict[str, CommandRecord] = field(default_factory=dict)
    results: list[CommandResultRecord] = field(default_factory=list)

+    def list_clients(self) -> list[ClientRecord]:
+        return sorted(self.clients.values(), key=lambda client: client.id)
+
+    def get_client(self, client_id: str) -> ClientRecord | None:
+        return self.clients.get(client_id)
+
    def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord:
        existing = self.clients.get(heartbeat.client_id)
        created_at = existing.created_at if existing else now
@@ -135,6 +156,15 @@ class InMemoryStore:
            timeout_seconds=record.timeout_seconds,
        )

+    def get_command(self, command_id: str) -> CommandRecord | None:
+        return self.commands.get(command_id)
+
+    def get_command_result(self, command_id: str) -> CommandResultRecord | None:
+        for result in reversed(self.results):
+            if result.command_id == command_id:
+                return result
+        return None
+

 TERMINAL_COMMAND_STATUSES = {
    CommandStatus.SUCCEEDED,