PillarTool/pillar_tool/routers/environment.py

import uuid


from sqlalchemy import select, insert, bindparam, delete
from sqlalchemy.orm import Session
from starlette.exceptions import HTTPException
from starlette.requests import Request
from fastapi import APIRouter, Query, Depends
from starlette.responses import JSONResponse

from pillar_tool.db import Host
from pillar_tool.db.models.top_data import Environment, EnvironmentAssignment
from pillar_tool.schemas import HostgroupParams, get_model_from_query
from pillar_tool.util.validation import split_and_validate_path, validate_environment_name

router = APIRouter(
    prefix="/environment",
    tags=["environment"],
)


@router.get("")
def environments_get(req: Request):
    """
    Retrieve all environments.

    Fetches and returns a list of environment names from the database.

    Returns:
        JSONResponse: A JSON response with status code 200 containing a list of environment names (strings).
    """
    db: Session = req.state.db

    result = db.execute(select(Environment)).fetchall()
    environments: list[Environment] = list(map(lambda x: x[0], result))

    return JSONResponse(status_code=200, content=[env.name for env in environments])


@router.get("/{name}")
def environment_get(req: Request, name: str):
    """
    Retrieve a specific environment by name.

    Fetches and returns details of the specified environment.
    Returns 404 if no such environment exists.

    Args:
        req (Request): The incoming request object.
        name (str): The name of the environment to retrieve.

    Returns:
        JSONResponse: A JSON response with status code 200 and the environment details on success,
                      or 404 if not found.
    """
    db: Session = req.state.db

    # Validate name before query
    if not validate_environment_name(name):
        raise HTTPException(status_code=400, detail="Invalid environment name format")

    stmt = select(Environment).where(Environment.name == name)
    result = db.execute(stmt).fetchall()

    if len(result) == 0:
        raise HTTPException(status_code=404, detail="No such environment exists")

    assert len(result) == 1

    env: Environment = result[0][0]

    # Get assigned hosts count as an example of additional info
    hosts_stmt = select(Host).join(EnvironmentAssignment, Host.id == EnvironmentAssignment.host_id)\
                              .where(EnvironmentAssignment.environment_id == env.id)
    hosts_count = db.execute(hosts_stmt).fetchall().__len__()

    return JSONResponse(status_code=200, content={
        'environment': env.name,
        'host_count': hosts_count
    })


@router.post("/{name}")
def environment_create(req: Request, name: str):
    """
    Create a new environment.

    Creates a new environment record in the database with the provided parameters.

    Args:
        req (Request): The incoming request object.
        name (str): The name of the environment (must be unique).

    Returns:
        JSONResponse: A JSON response with status code 201 on success,
                      or appropriate error codes (e.g., 409 if already exists, 400 for invalid format).
    """
    db = req.state.db

    # Validate name format
    if not validate_environment_name(name):
        raise HTTPException(status_code=400,
                            detail="Invalid environment name. Use only alphanumeric, underscore or dash characters.")

    # Check if environment already exists
    stmt_check = select(Environment).where(Environment.name == name)
    existing = db.execute(stmt_check).fetchall()

    if len(existing) > 0:
        raise HTTPException(status_code=409, detail="Environment already exists")

    new_id = uuid.uuid4()
    db.execute(insert(Environment).values(id=new_id, name=name))

    return JSONResponse(status_code=201, content={
        'id': str(new_id),
        'name': name
    })


@router.delete("/{name}")
def environment_delete(req: Request, name: str):
    """
    Delete an environment by name.

    Deletes the specified environment from the database.
    Returns 409 if hosts are still assigned to this environment.
    Returns 404 if no such environment exists.

    Args:
        req (Request): The incoming request object.
        name (str): The name of the environment to delete.

    Returns:
        JSONResponse: A JSON response with status code 204 on successful deletion,
                      or appropriate error codes for conflicts or not found.
    """
    db = req.state.db

    # Validate name format
    if not validate_environment_name(name):
        raise HTTPException(status_code=400, detail="Invalid environment name format")

    stmt = select(Environment).where(Environment.name == name)
    result = db.execute(stmt).fetchall()

    if len(result) == 0:
        raise HTTPException(status_code=404, detail="No such environment exists")

    assert len(result) == 1

    env: Environment = result[0][0]

    # Check for assigned hosts before deleting
    assignments_stmt = select(EnvironmentAssignment).where(
        EnvironmentAssignment.environment_id == env.id
    )
    assignments = db.execute(assignments_stmt).fetchall()

    if len(assignments) > 0:
        host_ids_stmt = select(EnvironmentAssignment.host_id).where(
            EnvironmentAssignment.environment_id == env.id
        )
        host_ids = [row[0] for row in db.execute(host_ids_stmt).fetchall()]

        # Get host names (could optimize this)
        hosts_stmt = select(Host).where(Host.id.in_(host_ids))
        hosts: list[Host] = list(map(lambda x: x[0], db.execute(hosts_stmt).fetchall()))

        return JSONResponse(status_code=409, content={
            'message': "Cannot delete an environment that still has hosts assigned",
            'assigned_hosts': [h.name for h in hosts]
        })

    # Delete the environment
    db.execute(delete(Environment).where(Environment.id == env.id))

    return JSONResponse(status_code=204, content={})