your-workflow/backend/app/nodes/status_node.py

79 lines
2.3 KiB
Python

import asyncio
from fastapi import FastAPI
from core.pocket_flow import AsyncNode
from core.logger import logger
from .base import async_run_cmd
class StatusNode(AsyncNode):
def set_info(self, id, callback):
self.id = id
self.callback = callback
async def prep_async(self, shared):
if self.callback:
await self.callback(self.id, "exec")
status_url = self.params[self.id].get("status_url", "")
cmd_str = "\n".join(
[
f'status_code=$(curl -s -o /dev/null -w "%{{http_code}}" {status_url})',
'if [ "$status_code" -eq 200 ]; then',
' echo "true"',
" exit 0",
"else",
' echo "Error: status code $status_code" >&2',
" exit 1",
"fi",
]
)
logger.info(cmd_str)
return {
"host": self.params[self.id].get("host"),
"user": self.params[self.id].get("user"),
"password": self.params[self.id].get("password"),
"port": self.params[self.id].get("port"),
"sudo": self.params[self.id].get("sudo", False),
"wait": self.params[self.id].get("wait", 3),
"cmd": cmd_str,
}
async def exec_async(self, prep_res):
if self.callback:
await self.callback(self.id, "exec")
err = []
out = []
async for stream_type, line in async_run_cmd(**prep_res):
if stream_type == "stderr":
err.append(str(line))
logger.debug(f"err: {line}")
else:
out.append(str(line))
logger.debug(f"out: {line}")
err_str = "\n".join(err)
out_str = "\n".join(out)
# 补充一下密码失效问题
if "Sorry, try again." in out_str and not err_str:
err_str = "SSH connection failed no password was provided"
return {"err": err_str, "out": out_str}
async def post_async(self, shared, prep_res, exec_res):
shared[self.id] = exec_res
result = "true"
if exec_res.get("err"):
result = "false"
if self.callback:
await self.callback(self.id, "post")
return result