Xiaoxia AI
56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
from sqlalchemy.orm import Session
|
|
|
|
from packages.adapters.sqlalchemy_impl.models import ClassificationJobModel
|
|
from packages.domain import ClassificationJob, ClassificationJobStatus
|
|
|
|
|
|
class SQLAlchemyClassificationJobRepository:
|
|
def __init__(self, session: Session):
|
|
self.session = session
|
|
|
|
def create(self, job: ClassificationJob) -> ClassificationJob:
|
|
model = ClassificationJobModel(
|
|
id=job.id,
|
|
workspace_id=job.workspace_id,
|
|
project_id=job.project_id,
|
|
asset_id=job.asset_id,
|
|
status=job.status.value,
|
|
classification=job.classification,
|
|
confidence=job.confidence,
|
|
error_message=job.error_message,
|
|
created_at=job.created_at,
|
|
updated_at=job.updated_at,
|
|
)
|
|
self.session.add(model)
|
|
self.session.commit()
|
|
return job
|
|
|
|
def get(self, job_id: str) -> ClassificationJob | None:
|
|
model = self.session.query(ClassificationJobModel).filter(ClassificationJobModel.id == job_id).first()
|
|
if model is None:
|
|
return None
|
|
return ClassificationJob(
|
|
id=model.id,
|
|
workspace_id=model.workspace_id,
|
|
project_id=model.project_id,
|
|
asset_id=model.asset_id,
|
|
status=ClassificationJobStatus(model.status),
|
|
classification=model.classification,
|
|
confidence=model.confidence,
|
|
error_message=model.error_message,
|
|
created_at=model.created_at,
|
|
updated_at=model.updated_at,
|
|
)
|
|
|
|
def update(self, job: ClassificationJob) -> ClassificationJob:
|
|
model = self.session.query(ClassificationJobModel).filter(ClassificationJobModel.id == job.id).first()
|
|
if model is None:
|
|
raise ValueError(f"ClassificationJob {job.id} not found")
|
|
model.status = job.status.value
|
|
model.classification = job.classification
|
|
model.confidence = job.confidence
|
|
model.error_message = job.error_message
|
|
model.updated_at = job.updated_at
|
|
self.session.commit()
|
|
return job
|