orm get jobs by experiment

mozilla-ai · Jan 15, 2025 · 9283c23 · 9283c23
1 parent 8c838c9
commit 9283c23
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/lumigator/python/mzai/backend/backend/repositories/jobs.py b/lumigator/python/mzai/backend/backend/repositories/jobs.py
@@ -1,5 +1,6 @@
 from uuid import UUID
 
+from sqlalchemy import desc
 from sqlalchemy.orm import Session
 
 from backend.records.jobs import JobRecord, JobResultRecord
@@ -16,8 +17,7 @@ def __init__(self, session: Session):
         super().__init__(JobResultRecord, session)
 
     def get_by_job_id(self, job_id: UUID) -> JobResultRecord | None:
-        return (
-            self.session.query(JobResultRecord)
-            .where(JobResultRecord.job_id == job_id)
-            .first()
-        )
+        return self.session.query(JobResultRecord).where(JobResultRecord.job_id == job_id).first()
+
+    def get_jobs_by_experiment_id(self, experiment_id: UUID) -> list[JobRecord]:
+        return self.session.query(JobRecord).order_by(desc(JobRecord.created_at)).limit(2).all()