be-optimumoh/src/overhaul_activity/service.py

import asyncio
from uuid import UUID
from sqlalchemy import Select, Delete, func, select, update as sqlUpdate
from sqlalchemy.orm import joinedload
from typing import List, Optional
from sqlalchemy.dialects.postgresql import insert

from src.overhaul_activity.utils import get_material_cost, get_service_cost
from src.overhaul_scope.model import OverhaulScope

from .model import OverhaulActivity
from .schema import OverhaulActivityCreate, OverhaulActivityUpdate, OverhaulActivityRead

from src.database.core import DbSession
from src.database.service import CommonParameters, search_filter_sort_paginate
from src.auth.service import CurrentUser
from src.overhaul_scope.service import get as get_session


async def get(*, db_session: DbSession, assetnum: str, overhaul_session_id: Optional[UUID] = None) -> Optional[OverhaulActivityRead]:
    """Returns a document based on the given document id."""
    query = Select(OverhaulActivity).where(
        OverhaulActivity.assetnum == assetnum).options(joinedload(OverhaulActivity.equipment))

    if overhaul_session_id:
        query = query.filter(
            OverhaulActivity.overhaul_scope_id == overhaul_session_id)

    result = await db_session.execute(query)
    return result.scalar()


async def get_all(*, common: CommonParameters, overhaul_session_id: UUID, assetnum: Optional[str] = None, scope_name: Optional[str] = None):
    query = Select(OverhaulActivity).where(
        OverhaulActivity.overhaul_scope_id == overhaul_session_id).options(joinedload(OverhaulActivity.equipment))

    if assetnum:
        query = query.filter(OverhaulActivity.assetnum == assetnum).options(
            joinedload(OverhaulActivity.overhaul_scope))

    if scope_name:
        query = query.filter(OverhaulActivity.scope_name == scope_name).options(
            joinedload(OverhaulActivity.overhaul_scope))

    results = await search_filter_sort_paginate(model=query, **common)

    return results


# async def create(*, db_session: DbSession, overhaul_activty_in: OverhaulActivityCreate, overhaul_session_id: UUID):
#     # Check if the combination of assetnum and activity_id already exists
#     existing_equipment_query = (
#         Select(OverhaulActivity)
#         .where(
#             OverhaulActivity.assetnum == overhaul_activty_in.assetnum,
#             OverhaulActivity.overhaul_scope_id == overhaul_session_id
#         )
#     )
#     result = await db_session.execute(existing_equipment_query)
#     existing_activity = result.scalar_one_or_none()

#     # If the combination exists, raise an exception or return the existing activity
#     if existing_activity:
#         raise ValueError("This assetnum already exist.")

#     activity = OverhaulActivity(
#         **overhaul_activty_in.model_dump(),
#         overhaul_scope_id=overhaul_session_id)
#     db_session.add(activity)
#     await db_session.commit()

#     # Refresh and load relationships using joinedload
#     query = (
#         Select(OverhaulActivity)
#         .options(joinedload(OverhaulActivity.equipment))
#         .where(OverhaulActivity.id == activity.id)
#     )
#     result = await db_session.execute(query)
#     activity_with_relationship = result.scalar_one()

#     return activity_with_relationship

async def create(*, db_session: DbSession, overhaul_activty_in: OverhaulActivityCreate, overhaul_session_id: UUID):
    """Creates a new document."""
    assetnums = overhaul_activty_in.assetnums
    if not assetnums:
        return []

    # Get session and count in parallel
    session = await get_session(db_session=db_session, overhaul_session_id=overhaul_session_id)
    equipment_count = await db_session.scalar(
        select(func.count())
        .select_from(OverhaulActivity)
        .where(OverhaulActivity.overhaul_scope_id == overhaul_session_id)
    )

    # Calculate costs for all records
    total_equipment = equipment_count + len(assetnums)
    material_cost = get_material_cost(
        scope=session.type, total_equipment=total_equipment)
    service_cost = get_service_cost(
        scope=session.type, total_equipment=total_equipment)

    # Create the insert statement
    stmt = insert(OverhaulActivity).values([
        {
            'assetnum': assetnum,
            'overhaul_scope_id': overhaul_session_id,
            'material_cost': material_cost,
            'service_cost': service_cost
        }
        for assetnum in assetnums
    ])

    # Add the ON CONFLICT DO NOTHING clause
    stmt = stmt.on_conflict_do_nothing(
            index_elements=["assetnum", "overhaul_scope_id"]
        )

    # Execute the statement
    db_session.execute(stmt)
    await db_session.execute(
        sqlUpdate(OverhaulActivity)
        .where(OverhaulActivity.overhaul_scope_id == overhaul_session_id)
        .values(material_cost=material_cost, service_cost=service_cost)
    )

    await db_session.commit()
    return assetnums


async def update(*, db_session: DbSession, activity: OverhaulActivity, overhaul_activity_in: OverhaulActivityUpdate):
    """Updates a document."""
    data = overhaul_activity_in.model_dump()

    update_data = overhaul_activity_in.model_dump(exclude_defaults=True)

    for field in data:
        if field in update_data:
            setattr(activity, field, update_data[field])

    await db_session.commit()

    return activity


async def delete(*, db_session: DbSession, overhaul_activity_id: str):
    """Deletes a document."""
    activity = await db_session.get(OverhaulActivity, overhaul_activity_id)
    await db_session.delete(activity)
    await db_session.commit()