import uuid from sqlalchemy import select, delete, and_, bindparam from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session from starlette.exceptions import HTTPException from starlette.requests import Request from fastapi import APIRouter from starlette.responses import JSONResponse from pillar_tool.db import Host, Environment, EnvironmentAssignment from pillar_tool.db.models.top_data import State, StateAssignment, TopFile from pillar_tool.util.validation import validate_state_name router = APIRouter( prefix="/top", tags=["top"], ) @router.get("/query/{host}") def top_get(req: Request, host: str): db: Session = req.state.db # build the hierarchy host_stmt = select(Host).where(Host.name == host) result = db.execute(host_stmt).fetchall() if len(result) == 0: return JSONResponse(status_code=404, content={"message": "Host '{}' not found".format(host)}) elif len(result) > 1: return JSONResponse(status_code=500, content={"message": "More than one host found"}) else: target_host: Host = result[0][0] parent_stmt = select(Host).where(Host.id == bindparam("parent_id")) parents = [] current: Host = target_host while current is not None and current.parent_id is not None: parents.append(current) result = db.execute(parent_stmt, {'parent_id': current.parent_id}).fetchall() if len(result) == 0: current = None elif len(result) > 1: return JSONResponse(status_code=500, content={"message": "More than one parent host found"}) else: current = result[0][0] if current is not None: parents.append(current) env_stmt = (select(Environment) .join(EnvironmentAssignment, EnvironmentAssignment.environment_id == Environment.id) .where(EnvironmentAssignment.host_id == bindparam("host_id")) ) env: Environment | None = None for host in reversed(parents): env_res = db.execute(env_stmt, {'host_id': host.id}).fetchall() if len(env_res) == 1: env: Environment = env_res[0][0] state_stmt = (select(State) .join(TopFile, State.id == TopFile.state_id) .join(StateAssignment, State.id == StateAssignment.state_id) .where(and_(StateAssignment.environment_id == env.id, TopFile.host_id == bindparam("host_id"))) ) assigned_states = [ [ row[0] for row in db.execute(state_stmt, {'host_id': host.id}).fetchall() ] for host in parents ] all_assigned_states = set(s for states in assigned_states for s in states) env_name = env.name return JSONResponse(status_code=200, content={ env.name: list(map(lambda state: state.name, all_assigned_states)), }) @router.post("/setenv/{host}/{environment}") def top_setenv(req: Request, host: str, environment: str): db: Session = req.state.db # get the target host id host_stmt = select(Host).where(Host.name == host) host_res = db.execute(host_stmt).fetchall() if len(host_res) == 0: return JSONResponse(status_code=404, content={"error": "No host found"}) elif len(host_res) == 1: host_res = host_res[0][0] else: # Note that this should be prevented by the database return JSONResponse(status_code=404, content={"error": "Too many hosts found??? This should not happen"}) # get the environment id env_stmt = select(Environment).where(Environment.name == environment) env_res = db.execute(env_stmt).fetchall() if len(env_res) == 0: return JSONResponse(status_code=404, content={"error": "No environment found"}) elif len(env_res) == 1: env_res = env_res[0][0] else: # Note that this should be prevented by the database return JSONResponse(status_code=404, content={"error": "Too many environments found??? This should not happen"}) insert_stmt = insert(EnvironmentAssignment).values(environment_id=env_res.id, host_id=host_res.id) result = db.execute(insert_stmt) return JSONResponse(status_code=200, content={}) @router.post("/assign/{host_name}/{state_name}") def top_state_assign(req: Request, host_name: str, state_name: str): db: Session = req.state.db # get the host in question host_stmt = select(Host).where(Host.name == host_name) host_res = db.execute(host_stmt).fetchall() if len(host_res) != 1: return JSONResponse(status_code=404, content={"error": f"Host '{host_name} not found"}) host: Host = host_res[0][0] parent_stmt = select(Host).where(Host.id == bindparam("parent_id")) parents: list[Host] = [] current: Host = host while current is not None: parents.append(current) if current.parent_id is None: current = None else: parent = db.execute(parent_stmt, {'parent_id': current.parent_id}).fetchall() if len(parent) == 0: return JSONResponse(status_code=500, content={"error": f"Host Hierarchy seems broken: parent_id '{current.parent_id}' does not exist"}) # Note: more than one result is impossible, since the id is a primary key current: Host = parent[0][0] # get the hosts environment env_assign_stmt = select(EnvironmentAssignment).where(EnvironmentAssignment.host_id == bindparam("host_id")) env_assign: EnvironmentAssignment | None = None for current_host in parents: env_res = db.execute(env_assign_stmt, {'host_id': current_host.id}).fetchall() if len(env_res) == 1: env_assign: EnvironmentAssignment = env_res[0][0] break env_stmt = select(Environment).where(Environment.id == env_assign.environment_id) env_res = db.execute(env_stmt).fetchall() if len(env_res) != 1: return JSONResponse(status_code=404, content={"error": f"Host '{host_name}' has no environment assigned"}) env: Environment = env_res[0][0] # get the state in question state_stmt = (select(State).join(StateAssignment, State.id == StateAssignment.state_id) .where(and_(State.name == state_name, StateAssignment.environment_id == env.id)) ) state_res = db.execute(state_stmt).fetchall() if len(state_res) != 1: return JSONResponse(status_code=404, content={"error": f"No state '{state_name}' found in environment '{env.name}'"}) state: State = state_res[0][0] # insert the relation into the database db.execute(insert(TopFile).on_conflict_do_nothing('pillar_tool_top_file_unique_state_host').values(state_id=state.id, host_id=host.id)) return JSONResponse(status_code=200, content={}) @router.delete("/assign/{host_name}/{state_name}") def top_state_unassign(req: Request, host_name: str, state_name: str): db: Session = req.state.db # get the host in question host_stmt = select(Host).where(Host.name == host_name) host_res = db.execute(host_stmt).fetchall() if len(host_res) != 1: return JSONResponse(status_code=404, content={"error": f"Host '{host_name} not found"}) host: Host = host_res[0][0] parent_stmt = select(Host).where(Host.id == bindparam("parent_id")) parents: list[Host] = [] current: Host = host while current is not None: parents.append(current) if current.parent_id is None: current = None else: parent = db.execute(parent_stmt, {'parent_id': current.parent_id}).fetchall() if len(parent) == 0: return JSONResponse(status_code=500, content={"error": f"Host Hierarchy seems broken: parent_id '{current.parent_id}' does not exist"}) # Note: more than one result is impossible, since the id is a primary key current: Host = parent[0][0] # get the hosts environment env_assign_stmt = select(EnvironmentAssignment).where(EnvironmentAssignment.host_id == bindparam("host_id")) env_assign: EnvironmentAssignment | None = None for current_host in parents: env_res = db.execute(env_assign_stmt, {'host_id': current_host.id}).fetchall() if len(env_res) == 1: env_assign: EnvironmentAssignment = env_res[0][0] break env_stmt = select(Environment).where(Environment.id == env_assign.environment_id) env_res = db.execute(env_stmt).fetchall() if len(env_res) != 1: return JSONResponse(status_code=404, content={"error": f"Host '{host_name}' has no environment assigned"}) env: Environment = env_res[0][0] # get the state in question state_stmt = (select(State).join(StateAssignment, State.id == StateAssignment.state_id) .where(and_(State.name == state_name, StateAssignment.environment_id == env.id)) ) state_res = db.execute(state_stmt).fetchall() if len(state_res) != 1: return JSONResponse(status_code=404, content={"error": f"No state '{state_name}' found in environment '{env.name}'"}) state: State = state_res[0][0] # delete the relation from the database db.execute(delete(TopFile).where(and_(TopFile.state_id == state.id, TopFile.host_id == host.id))) return JSONResponse(status_code=200, content={})