221 lines
9.0 KiB
Python

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={})