projects.py

"""Definitions of the endpoints related the projects."""
import logging
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import create_engine, text, update
from sqlalchemy.exc import DBAPIError
from sqlalchemy.future import select
from sqlalchemy.orm import Session

from ...db import DBProject, DBProjectServer
from ...helpers import fix_sqlalchemy2_stubs_non_nullable_column
from ...schemas import Project
from ..depends import get_session, get_session_as_you_go

logger = logging.getLogger(__name__)
router = APIRouter()


@router.get('', summary='Lists the projects.', response_model=list[Project])
def list_projects(*, session: Session = Depends(get_session)) -> list[DBProject]:
    """Lists the projects visible to a user."""
    stmt = select(DBProject)
    return session.execute(stmt).scalars().all()


@router.post('', response_model=Project)
def create_project(
    *, session: Session = Depends(get_session_as_you_go), project: Project
) -> DBProject:
    """Creates a new project."""
    stmt: Any = select(DBProject).filter_by(name=project.name)
    db_project = session.execute(stmt).scalars().first()
    if db_project is not None:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=f"The project '{project.name} already exists.",
        )

    server = _find_best_project_server(session, project)

    # preempt the required storage
    _update_server_available_size(session, server, -project.max_size)
    session.commit()

    try:
        with create_engine(
            fix_sqlalchemy2_stubs_non_nullable_column(server.uri),
            pool_pre_ping=True,
            future=True,
        ).connect() as conn:
            stmt = text(f'CREATE DATABASE "{project.name}"')
            conn.execution_options(isolation_level='AUTOCOMMIT').execute(stmt)
    except DBAPIError as exc:
        logger.error(str(exc))
        # release the preempted storage
        _update_server_available_size(session, server, +project.max_size)
        session.commit()
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail='The project database could not be created: '
            f'{type(exc).__name__}: {exc}',
        )

    uri = f'{server.uri}/{project.name}'
    db_project = DBProject(
        project_server_id=server.id,
        name=project.name,
        description=project.description,
        uri=uri,
        max_size=project.max_size,
    )
    session.add(db_project)
    session.commit()
    return db_project


def _find_best_project_server(session: Session, project: Project) -> DBProjectServer:
    stmt = select(DBProjectServer).where(
        DBProjectServer.available_size >= project.max_size
    )
    server = session.execute(stmt).scalars().first()
    if server is None:
        raise HTTPException(
            status_code=status.HTTP_507_INSUFFICIENT_STORAGE,
            detail='No database has enough storage for this project.',
        )
    return server


def _update_server_available_size(
    session: Session, server: DBProjectServer, size: int
) -> None:
    available_size = (
        fix_sqlalchemy2_stubs_non_nullable_column(server.available_size) + size
    )
    stmt = (
        update(DBProjectServer)
        .where(DBProjectServer.id == server.id)
        .values(available_size=available_size)
    )
    session.execute(stmt)


@router.get('/{project}', summary='Gets a project.', response_model=Project)
def get_project(project: str, *, session: Session = Depends(get_session)) -> DBProject:
    """Gets a project visible to a user."""
    stmt = select(DBProject).where(DBProject.name == project)
    db_project = session.execute(stmt).scalars().first()
    if db_project is None:
        msg = f"The project '{project}' is not known."
        raise HTTPException(status.HTTP_404_NOT_FOUND, msg)
    return db_project