Model/backend/app/db/functions/bulk_address_uploads_functions.py

from uuid import UUID
from datetime import datetime, timezone

from sqlmodel import select

from backend.app.db.connection import get_db_session
from backend.app.db.models.bulk_address_uploads import BulkAddressUpload


def set_combining_status(task_id: UUID) -> None:
    now = datetime.now(timezone.utc)
    with get_db_session() as session:
        row = session.exec(
            select(BulkAddressUpload).where(BulkAddressUpload.task_id == task_id)
        ).first()
        if not row:
            raise ValueError(f"No bulk_address_uploads row for task_id {task_id}")
        row.status = "combining"
        row.updated_at = now
        session.add(row)
        session.commit()


def set_combined_output_s3_uri(task_id: UUID, s3_uri: str) -> None:
    now = datetime.now(timezone.utc)
    with get_db_session() as session:
        row = session.exec(
            select(BulkAddressUpload).where(BulkAddressUpload.task_id == task_id)
        ).first()
        if not row:
            raise ValueError(f"No bulk_address_uploads row for task_id {task_id}")
        row.combined_output_s3_uri = s3_uri
        row.status = "awaiting_review"
        row.updated_at = now
        session.add(row)
        session.commit()