Add node messages to node get (#836)

This exposes the node commands that have yet to be processed by the node.  Example use case:  The SDK can now ask "has this node installed my SSH key"
This commit is contained in:
bmc-msft
2021-04-26 16:14:58 -04:00
committed by GitHub
parent 541e745199
commit ced21b2ea3
3 changed files with 29 additions and 25 deletions

View File

@ -12,7 +12,7 @@ from onefuzztypes.responses import BoolResult
from ..onefuzzlib.endpoint_authorization import call_if_user from ..onefuzzlib.endpoint_authorization import call_if_user
from ..onefuzzlib.events import get_events from ..onefuzzlib.events import get_events
from ..onefuzzlib.request import not_ok, ok, parse_request from ..onefuzzlib.request import not_ok, ok, parse_request
from ..onefuzzlib.workers.nodes import Node, NodeTasks from ..onefuzzlib.workers.nodes import Node, NodeMessage, NodeTasks
def get(req: func.HttpRequest) -> func.HttpResponse: def get(req: func.HttpRequest) -> func.HttpResponse:
@ -33,6 +33,9 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
node_tasks = NodeTasks.get_by_machine_id(request.machine_id) node_tasks = NodeTasks.get_by_machine_id(request.machine_id)
node.tasks = [(t.task_id, t.state) for t in node_tasks] node.tasks = [(t.task_id, t.state) for t in node_tasks]
node.messages = [
x.message for x in NodeMessage.get_messages(request.machine_id)
]
return ok(node) return ok(node)

View File

@ -157,7 +157,7 @@ class Node(BASE_NODE, ORMMixin):
return ("pool_name", "machine_id") return ("pool_name", "machine_id")
def save_exclude(self) -> Optional[MappingIntStrAny]: def save_exclude(self) -> Optional[MappingIntStrAny]:
return {"tasks": ...} return {"tasks": ..., "messages": ...}
def telemetry_include(self) -> Optional[MappingIntStrAny]: def telemetry_include(self) -> Optional[MappingIntStrAny]:
return { return {

View File

@ -557,6 +557,29 @@ class NodeHeartbeatEntry(BaseModel):
data: List[Dict[str, HeartbeatType]] data: List[Dict[str, HeartbeatType]]
class StopNodeCommand(BaseModel):
pass
class StopTaskNodeCommand(BaseModel):
task_id: UUID
class NodeCommandAddSshKey(BaseModel):
public_key: str
class NodeCommand(EnumModel):
stop: Optional[StopNodeCommand]
stop_task: Optional[StopTaskNodeCommand]
add_ssh_key: Optional[NodeCommandAddSshKey]
class NodeCommandEnvelope(BaseModel):
command: NodeCommand
message_id: str
class Node(BaseModel): class Node(BaseModel):
timestamp: Optional[datetime] = Field(alias="Timestamp") timestamp: Optional[datetime] = Field(alias="Timestamp")
pool_name: PoolName pool_name: PoolName
@ -564,6 +587,7 @@ class Node(BaseModel):
state: NodeState = Field(default=NodeState.init) state: NodeState = Field(default=NodeState.init)
scaleset_id: Optional[UUID] = None scaleset_id: Optional[UUID] = None
tasks: Optional[List[Tuple[UUID, NodeTaskState]]] = None tasks: Optional[List[Tuple[UUID, NodeTaskState]]] = None
messages: Optional[List[NodeCommand]] = None
heartbeat: Optional[datetime] heartbeat: Optional[datetime]
version: str = Field(default="1.0.0") version: str = Field(default="1.0.0")
reimage_requested: bool = Field(default=False) reimage_requested: bool = Field(default=False)
@ -776,29 +800,6 @@ class NodeEventEnvelope(BaseModel):
event: NodeEventShim event: NodeEventShim
class StopNodeCommand(BaseModel):
pass
class StopTaskNodeCommand(BaseModel):
task_id: UUID
class NodeCommandAddSshKey(BaseModel):
public_key: str
class NodeCommand(EnumModel):
stop: Optional[StopNodeCommand]
stop_task: Optional[StopTaskNodeCommand]
add_ssh_key: Optional[NodeCommandAddSshKey]
class NodeCommandEnvelope(BaseModel):
command: NodeCommand
message_id: str
class TaskEvent(BaseModel): class TaskEvent(BaseModel):
timestamp: Optional[datetime] = Field(alias="Timestamp") timestamp: Optional[datetime] = Field(alias="Timestamp")
task_id: UUID task_id: UUID