from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.entities import UploadSession


class UploadRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get_by_id(self, upload_id: str) -> UploadSession | None:
        return await self.session.get(UploadSession, upload_id)

    async def create(self, upload_session: UploadSession) -> UploadSession:
        self.session.add(upload_session)
        await self.session.flush()
        await self.session.refresh(upload_session)
        return upload_session

    async def save(self, upload_session: UploadSession) -> UploadSession:
        await self.session.flush()
        await self.session.refresh(upload_session)
        return upload_session