Update backend/app/repositories/experiment_repo.py

backend/app/repositories/experiment_repo.py

@@ -1,16 +1,46 @@
-from sqlalchemy import Column, Integer, String, Text
-from backend.app.db import Base
-
-
-class Experiment(Base):
-    __tablename__ = "experiments"
-
-    id = Column(String, primary_key=True, index=True)
-    dataset_id = Column(String, nullable=False, index=True)
-    algorithm = Column(String, nullable=False)
-    status = Column(String, nullable=False)
-    config_json = Column(Text, nullable=True)
-    metrics_json = Column(Text, nullable=True)
-    summary_json = Column(Text, nullable=True)
-    runtime_ms = Column(Integer, nullable=True)
-    error_message = Column(Text, nullable=True)
+from sqlalchemy.orm import Session
+
+from backend.app.models.experiment import Experiment
+
+
+def create_experiment(
+    db: Session,
+    id: str,
+    dataset_id: str,
+    algorithm: str,
+    status: str,
+    config_json: str | None = None,
+    metrics_json: str | None = None,
+    summary_json: str | None = None,
+    runtime_ms: int | None = None,
+    error_message: str | None = None,
+) -> Experiment:
+    experiment = Experiment(
+        id=id,
+        dataset_id=dataset_id,
+        algorithm=algorithm,
+        status=status,
+        config_json=config_json,
+        metrics_json=metrics_json,
+        summary_json=summary_json,
+        runtime_ms=runtime_ms,
+        error_message=error_message,
+    )
+    db.add(experiment)
+    db.commit()
+    db.refresh(experiment)
+    return experiment
+
+
+def get_experiment(db: Session, experiment_id: str) -> Experiment | None:
+    return db.query(Experiment).filter(Experiment.id == experiment_id).first()
+
+
+def list_experiments(db: Session) -> list[Experiment]:
+    return db.query(Experiment).order_by(Experiment.id.desc()).all()
+
+
+def get_experiments_by_ids(db: Session, experiment_ids: list[str]) -> list[Experiment]:
+    if not experiment_ids:
+        return []
+    return db.query(Experiment).filter(Experiment.id.in_(experiment_ids)).all()