import os import uuid import pygit2 from pygit2.enums import BranchType from sqlalchemy import select, delete, bindparam from sqlalchemy.orm import Session from sqlalchemy.dialects.postgresql import insert from sqlalchemy.sql.operators import in_op, and_ 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, StateAssignment from pillar_tool.db.models.top_data import Environment, EnvironmentAssignment, State from pillar_tool.util.files import recursive_list_dir from pillar_tool.util.validation import validate_environment_name from pillar_tool.util import config, Config from pillar_tool.git.repository import checkout_remote_branch 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 print("[DEBUG] environments_get: retrieving all environments") result = db.execute(select(Environment)).fetchall() environments: list[Environment] = list(map(lambda x: x[0], result)) print("[DEBUG] environments_get: found {} environment(s) in database".format(len(environments))) 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 print("[DEBUG] environment_get: retrieving environment '{}'".format(name)) # Validate name before query if not validate_environment_name(name): print("[DEBUG] environment_get: ERROR - Invalid environment name '{}'. Environment names must use only alphanumeric, underscore or dash characters.".format(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: print("[DEBUG] environment_get: ERROR - No environment found with name '{}'".format(name)) raise HTTPException(status_code=404, detail="No such environment exists") assert len(result) == 1 env: Environment = result[0][0] print("[DEBUG] environment_get: resolved environment '{}' with id={}".format(env.name, env.id)) # 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__() print("[DEBUG] environment_get: environment '{}' has {} host(s) assigned".format(env.name, hosts_count)) 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 print("[DEBUG] environment_create: creating new environment '{}'".format(name)) # Validate name format if not validate_environment_name(name): print("[DEBUG] environment_create: ERROR - Invalid environment name '{}'. Environment names must use only alphanumeric, underscore or dash characters.".format(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: print("[DEBUG] environment_create: ERROR - Environment '{}' already exists in database. Duplicate environment creation rejected.".format(name)) raise HTTPException(status_code=409, detail="Environment already exists") new_id = uuid.uuid4() db.execute(insert(Environment).values(id=new_id, name=name)) print("[DEBUG] environment_create: created environment '{}' with id={}".format(name, new_id)) 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 print("[DEBUG] environment_delete: deleting environment '{}'".format(name)) # Validate name format if not validate_environment_name(name): print("[DEBUG] environment_delete: ERROR - Invalid environment name '{}'. Environment names must use only alphanumeric, underscore or dash characters.".format(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: print("[DEBUG] environment_delete: ERROR - No environment found with name '{}'".format(name)) raise HTTPException(status_code=404, detail="No such environment exists") assert len(result) == 1 env: Environment = result[0][0] print("[DEBUG] environment_delete: resolved environment '{}' with id={}".format(env.name, env.id)) # 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())) print("[DEBUG] environment_delete: ERROR - Cannot delete environment '{}' because it has {} host(s) assigned: {}. All hosts must be unassigned before deletion.".format(name, len(assignments), [h.name for h in hosts])) 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)) print("[DEBUG] environment_delete: successfully deleted environment '{}' (id={})".format(env.name, env.id)) return JSONResponse(status_code=204, content={}) @router.patch("/{name}") def environment_patch(req: Request, name: str) -> JSONResponse: db: Session = req.state.db cfg: Config = config() print("[DEBUG] environment_patch: importing/syncing environment '{}' from git".format(name)) # Attempt to check the requested branch out try: checkout_remote_branch(cfg, name) print("[DEBUG] environment_patch: successfully checked out git branch '{}'".format(name)) # create the environment if it did not exist already select_env_res = db.execute(select(Environment).where(Environment.name == name)).fetchall() if len(select_env_res) == 0: insert_env_res = db.execute(insert(Environment).values(id=uuid.uuid4(), name=name).returning(Environment)).fetchall() print("[DEBUG] environment_patch: environment '{}' does not exist in database, creating new record".format(name)) if len(insert_env_res) == 0: print("[DEBUG] environment_patch: ERROR - Failed to create environment '{}' in database after git checkout. Database insert returned no results.".format(name)) raise HTTPException(status_code=404, detail=f"Failed to create non-existent environment '{name}'") else: env: Environment = insert_env_res[0][0] else: env: Environment = select_env_res[0][0] print("[DEBUG] environment_patch: environment '{}' already exists in database (id={})".format(name, env.id)) # Branch has been checked out all_files = recursive_list_dir(cfg.git.state_repo_path) sls_files = filter(lambda f: f.endswith(".sls"), all_files) state_file_paths = map(lambda x: x.replace("/init.sls", "").replace(".sls", ""), sls_files) state_names = list(map(lambda x: x.replace(f"{cfg.git.state_repo_path}/", "").replace("/", "."), state_file_paths)) print("[DEBUG] environment_patch: found {} .sls file(s) in git repo '{}', resolving to {} state name(s)".format(len(all_files), cfg.git.state_repo_path, len(state_names))) # get all the existing states and the to be created ones select_res = db.execute(select(State).where(State.name.in_(state_names))).fetchall() states_known = {} for row in select_res: state: State = row[0] states_known[state.name] = state print("[DEBUG] environment_patch: {} of {} states already exist in database".format(len(states_known), len(state_names))) states_new = { state_name: State(name=state_name, id=uuid.uuid4()) for state_name in state_names if state_name not in states_known } print("[DEBUG] environment_patch: {} new state(s) to be created".format(len(states_new))) states = {} states.update(states_known) states.update(states_new) # insert all states that don't exist already insert_stmt = insert(State) insert_assignment_stmt = insert(StateAssignment) for state in states_new.values(): db.execute(insert_stmt.values(id=state.id, name=state.name)) # ensure that all the state assignments exist state_ids = list(map(lambda x: x.id, states.values())) state_assignments_known = [ x[0].state_id for x in db.execute(select(StateAssignment).where(StateAssignment.environment_id == env.id)).fetchall() ] print("[DEBUG] environment_patch: {} existing state assignment(s) found for environment '{}'".format(len(state_assignments_known), name)) state_assignments_new = [ state_id for state_id in map(lambda x: x.id, states.values()) if state_id not in state_assignments_known ] print("[DEBUG] environment_patch: {} new state assignment(s) to be created".format(len(state_assignments_new))) state_assignments_to_delete = [ sid for sid in filter(lambda x: x not in state_ids, state_assignments_known) ] if len(state_assignments_to_delete) > 0: print("[DEBUG] environment_patch: {} stale state assignment(s) to be removed (states no longer exist in git)".format(len(state_assignments_to_delete))) for sid in state_assignments_to_delete: db.execute(delete(StateAssignment).where(and_(StateAssignment.state_id == sid, StateAssignment.environment_id == env.id))) # type: ignore insert_stmt = insert(StateAssignment) for sid in state_assignments_new: db.execute(insert_stmt.values(state_id=sid, environment_id=env.id)) return JSONResponse(status_code=200, content={'message': "done"}) except Exception as exc: print("[DEBUG] environment_patch: ERROR - Failed to import environment '{}': {}".format(name, str(exc))) raise HTTPException(status_code=404, detail=str(exc))