querying the top file from salt works

pillar_tool/ptcli/cli/top.py

@@ -36,3 +36,18 @@ def top_setenv(host: str, environment: str):
         click.echo("Assigned environment")
     except requests.exceptions.HTTPError as e:
         raise click.ClickException(f"Failed to assign environment:\n{e}")
+
+
+@top.command("assign")
+@click.argument("host")
+@click.argument("state")
+def top_assign(host: str, state: str):
+    click.echo("Assigning state to host...")
+
+    try:
+        response = requests.post(f'{base_url()}/top/assign/{host}/{state}', headers=auth_header())
+        response.raise_for_status()
+
+        click.echo("Assigned state")
+    except requests.exceptions.HTTPError as e:
+        raise click.ClickException(f"Failed to assign state:\n{e}")

pillar_tool/routers/environment.py

@@ -239,7 +239,7 @@ def environment_patch(req: Request, name: str) -> JSONResponse:
         # ensure that all the state assignments exist
         state_ids = list(map(lambda x: x.id, states.values()))
         state_assignments_known = [
-            x[0].id
+            x[0].state_id
             for x in db.execute(select(StateAssignment).where(StateAssignment.environment_id == env.id)).fetchall()
         ]
         state_assignments_new = [

pillar_tool/routers/top.py

@@ -1,6 +1,7 @@
 import uuid
 
-from sqlalchemy import select, insert, delete, and_, bindparam
+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
@@ -8,7 +9,7 @@ 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
+from pillar_tool.db.models.top_data import State, StateAssignment, TopFile
 from pillar_tool.util.validation import validate_state_name
 
 router = APIRouter(
@@ -29,34 +30,49 @@ def top_get(req: Request, host: str):
     elif len(result) > 1:
         return JSONResponse(status_code=500, content={"message": "More than one host found"})
     else:
-        target_host = result[0][0]
+        target_host: Host = result[0][0]
 
     parent_stmt = select(Host).where(Host.id == bindparam("parent_id"))
     parents = []
-    current = target_host
-    while current is not None:
-        result = db.execute(parent_stmt, {'parend_id': current.id}).fetchall()
+    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:
-            parents.append(result[0][0])
             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
+    ]
 
-    # TODO: states should be hierarchical, same as pillars are
-    select_stmt = (select(Host, EnvironmentAssignment, Environment)
-                   .where(and_(Host.name == host, Host.is_hostgroup == False))
-                   .join(EnvironmentAssignment, EnvironmentAssignment.host_id == Host.id)
-                   .join(Environment, EnvironmentAssignment.environment_id == Environment.id)
-                   )
+    all_assigned_states = set(s for states in assigned_states for s in states)
+    env_name = env.name
 
-    result = db.execute(select_stmt).fetchall()
-    print(result[0])
-
-    return JSONResponse(status_code=200, content={})
+    return JSONResponse(status_code=200, content={
+        env.name: list(map(lambda state: state.name, all_assigned_states)),
+    })
 
 
 @router.post("/setenv/{host}/{environment}")
@@ -64,7 +80,7 @@ def top_setenv(req: Request, host: str, environment: str):
     db: Session = req.state.db
 
     # get the target host id
-    host_stmt = select(Host).where(and_(Host.name == host, Host.is_hostgroup == False))
+    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"})
@@ -91,5 +107,68 @@ def top_setenv(req: Request, host: str, environment: str):
     return JSONResponse(status_code=200, content={})
 
 
-def top_state_assign(req: Request):
-    pass
+@router.post("/assign/{host_name}/{state_name}")
+def top_state_assign(req: Request, host_name: str, state_name: str):
+    db: Session = req.state.db
+
+    print(f"Assigning {state_name} to {host_name}")
+
+    # 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]
+    print(f"Found host: {host.id}")
+
+    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()
+        print(f"Looking at host: {current_host.name}")
+        if len(env_res) == 1:
+            env_assign: EnvironmentAssignment = env_res[0][0]
+            print(f"Found host with assigned environment: {current_host.name} with environment: {env_assign.environment_id}")
+            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]
+    print(f"Environment found: {env.name if env else 'None'}")
+
+    # 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))
+    )
+    print(f"Check 1: {state_stmt}")
+    state_res = db.execute(state_stmt).fetchall()
+    print("Check 2")
+    if len(state_res) != 1:
+        print(f"Check 3: State '{state_name}' not found in environment '{env.name}'")
+        return JSONResponse(status_code=404, content={"error": f"No state '{state_name}' found in environment '{env.name}'"})
+    state: State = state_res[0][0]
+    print("Check 4")
+
+    # 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))
+    print("Check 5")
+
+    return JSONResponse(status_code=200, content={})