6068599a47
- Adjusted function signatures in various test files to enhance clarity by aligning parameters. - Updated patching syntax for better readability across test cases. - Improved formatting and spacing in test assertions and mock setups. - Ensured consistent use of async/await patterns in async test functions. - Enhanced comments for better understanding of test intentions.
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
"""Extraction repository for database operations."""
|
|
|
|
from sqlalchemy import desc
|
|
from sqlmodel import select
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
from app.models.extraction import Extraction
|
|
from app.repositories.base import BaseRepository
|
|
|
|
|
|
class ExtractionRepository(BaseRepository[Extraction]):
|
|
"""Repository for extraction database operations."""
|
|
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
"""Initialize the extraction repository."""
|
|
super().__init__(Extraction, session)
|
|
|
|
async def get_by_service_and_id(
|
|
self,
|
|
service: str,
|
|
service_id: str,
|
|
) -> Extraction | None:
|
|
"""Get an extraction by service and service_id."""
|
|
result = await self.session.exec(
|
|
select(Extraction).where(
|
|
Extraction.service == service,
|
|
Extraction.service_id == service_id,
|
|
),
|
|
)
|
|
return result.first()
|
|
|
|
async def get_by_user(self, user_id: int) -> list[Extraction]:
|
|
"""Get all extractions for a user."""
|
|
result = await self.session.exec(
|
|
select(Extraction)
|
|
.where(Extraction.user_id == user_id)
|
|
.order_by(desc(Extraction.created_at)),
|
|
)
|
|
return list(result.all())
|
|
|
|
async def get_pending_extractions(self) -> list[Extraction]:
|
|
"""Get all pending extractions."""
|
|
result = await self.session.exec(
|
|
select(Extraction)
|
|
.where(Extraction.status == "pending")
|
|
.order_by(Extraction.created_at),
|
|
)
|
|
return list(result.all())
|
|
|
|
async def get_extractions_by_status(self, status: str) -> list[Extraction]:
|
|
"""Get extractions by status."""
|
|
result = await self.session.exec(
|
|
select(Extraction)
|
|
.where(Extraction.status == status)
|
|
.order_by(desc(Extraction.created_at)),
|
|
)
|
|
return list(result.all())
|