Add endpoint for jobs and add job type field

alembic/versions/2024_01_25_7a614c201092_add_job_type.py

@@ -0,0 +1,34 @@
+"""Add upload jobs
+
+Revision ID: 486e8c875e0b
+Revises: 6c3da78d1af0
+Create Date: 2024-01-25 10:15:59.984529
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "486e8c875e0b"
+down_revision = "6c3da78d1af0"
+branch_labels = None
+depends_on = None
+
+from alembic import op
+import sqlalchemy as sa
+
+from trailblazer.constants import JobType
+
+
+def upgrade():
+    op.add_column(
+        table_name="job",
+        column=sa.Column(
+            name="job_type",
+            type_=sa.Enum(*JobType.types()),
+            server_default=JobType.ANALYSIS,
+            nullable=False,
+        ),
+    )
+
+
+def downgrade():
+    op.drop_column(table_name="job", column_name="job_type")

tests/integration/endpoints/test_add_job.py

@@ -0,0 +1,26 @@
+from flask.testing import FlaskClient
+from http import HTTPStatus
+from trailblazer.dto.create_job_request import CreateJobRequest
+
+from trailblazer.store.models import Analysis
+
+TYPE_JSON = "application/json"
+
+
+def test_add_job_to_analysis(client: FlaskClient, analysis: Analysis):
+    # GIVEN an analysis
+
+    # GIVEN a valid request to add a job to the analysis
+    create_job_request = CreateJobRequest(name="job", slurm_id="12345")
+    data: str = create_job_request.model_dump_json()
+
+    # WHEN sending the request
+    response = client.post(
+        f"/api/v1/analyses/{analysis.id}/jobs", data=data, content_type=TYPE_JSON
+    )
+
+    # THEN it gives a success response
+    assert response.status_code == HTTPStatus.CREATED
+
+    # THEN it should return the job
+    assert response.json["name"] == create_job_request.name

trailblazer/constants.py

@@ -170,3 +170,12 @@ class TrailblazerStatusColor(StrEnum):
 }
 
 PIPELINES = [pipeline.lower() for pipeline in Pipeline]
+
+
+class JobType(StrEnum):
+    UPLOAD: str = "upload"
+    ANALYSIS: str = "analysis"
+
+    @classmethod
+    def types(cls) -> tuple:
+        return tuple(cls)

trailblazer/dto/create_job_request.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+from trailblazer.constants import JobType
+
+
+class CreateJobRequest(BaseModel):
+    name: str
+    job_type: JobType | None = JobType.ANALYSIS
+    slurm_id: int

trailblazer/dto/job_response.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+from trailblazer.constants import JobType
+
+
+class JobResponse(BaseModel):
+    name: str

trailblazer/server/api.py

@@ -24,12 +24,14 @@
     FailedJobsRequest,
     FailedJobsResponse,
 )
+from trailblazer.dto.create_job_request import CreateJobRequest
 from trailblazer.exc import MissingAnalysis
 from trailblazer.server.ext import store
 from trailblazer.server.utils import (
     parse_analyses_request,
     parse_analysis_update_request,
     parse_get_failed_jobs_request,
+    parse_job_create_request,
     stringify_timestamps,
 )
 from trailblazer.services import AnalysisService, JobService
@@ -80,6 +82,21 @@ def get_analysis(analysis_id):
         return jsonify(error=str(error)), HTTPStatus.NOT_FOUND
 
 
+@blueprint.route("/analyses/<int:analysis_id>/jobs", methods=["POST"])
+def add_job(analysis_id):
+    job_service: JobService = current_app.extensions.get("job_service")
+    try:
+        data: CreateJobRequest = parse_job_create_request(request)
+        response: AnalysisResponse = job_service.add_job(analysis_id=analysis_id, data=data)
+        return jsonify(response.model_dump()), HTTPStatus.CREATED
+    except MissingAnalysis as error:
+        return jsonify(error=str(error)), HTTPStatus.NOT_FOUND
+    except ValidationError as error:
+        return jsonify(error=str(error)), HTTPStatus.BAD_REQUEST
+    except Exception as error:
+        return jsonify(error=str(error)), HTTPStatus.CONFLICT
+
+
 @blueprint.route("/analyses/<int:analysis_id>", methods=["PUT"])
 def update_analysis(analysis_id):
     analysis_service: AnalysisService = current_app.extensions.get("analysis_service")

trailblazer/server/utils.py

@@ -3,6 +3,7 @@
 from flask import Request
 
 from trailblazer.dto import AnalysesRequest, AnalysisUpdateRequest, FailedJobsRequest
+from trailblazer.dto.create_job_request import CreateJobRequest
 
 
 def parse_analyses_request(request: Request) -> AnalysesRequest:
@@ -34,3 +35,7 @@ def stringify_timestamps(data: dict) -> dict[str, str]:
         if isinstance(val, datetime.datetime):
             data[key] = str(val)
     return data
+
+
+def parse_job_create_request(request: Request) -> CreateJobRequest:
+    return CreateJobRequest.model_validate(request.json)

trailblazer/services/job_service.py

@@ -2,7 +2,10 @@
 
 from trailblazer.constants import TrailblazerStatus
 from trailblazer.dto import FailedJobsRequest, FailedJobsResponse
-from trailblazer.services.utils import create_jobs_response
+from trailblazer.dto.create_job_request import CreateJobRequest
+from trailblazer.dto.job_response import JobResponse
+from trailblazer.services.utils import create_job_response, create_jobs_response
+from trailblazer.store.models import Job
 from trailblazer.store.store import Store
 from trailblazer.utils.datetime import get_date_number_of_days_ago
 
@@ -17,3 +20,7 @@ def get_failed_jobs(self, request: FailedJobsRequest) -> FailedJobsResponse:
             status=TrailblazerStatus.FAILED, since_when=time_window
         )
         return create_jobs_response(failed_jobs)
+
+    def add_job(self, analysis_id: int, data: CreateJobRequest) -> JobResponse:
+        job: Job = self.store.add_job(analysis_id=analysis_id, data=data)
+        return create_job_response(job)

trailblazer/services/utils.py

@@ -1,5 +1,6 @@
 from trailblazer.dto import AnalysisResponse, FailedJobsResponse
-from trailblazer.store.models import Analysis
+from trailblazer.dto.job_response import JobResponse
+from trailblazer.store.models import Analysis, Job
 
 
 def create_jobs_response(failed_job_statistics: list[dict]) -> FailedJobsResponse:
@@ -11,3 +12,7 @@ def create_analysis_response(analysis: Analysis) -> AnalysisResponse:
     analysis_data["jobs"] = [job.to_dict() for job in analysis.jobs]
     analysis_data["user"] = analysis.user.to_dict() if analysis.user else None
     return AnalysisResponse.model_validate(analysis_data)
+
+
+def create_job_response(job: Job) -> JobResponse:
+    return JobResponse(name=job.name)

trailblazer/store/crud/create.py

@@ -2,10 +2,11 @@
 
 from sqlalchemy.orm import Session
 
-from trailblazer.constants import TrailblazerStatus
+from trailblazer.constants import JobType, SlurmJobStatus, TrailblazerStatus
+from trailblazer.dto.create_job_request import CreateJobRequest
 from trailblazer.store.base import BaseHandler
 from trailblazer.store.database import get_session
-from trailblazer.store.models import Analysis, User
+from trailblazer.store.models import Analysis, Job, User
 
 
 class CreateHandler(BaseHandler):
@@ -49,3 +50,17 @@ def add_user(self, name: str, email: str) -> User:
         session.add(new_user)
         session.commit()
         return new_user
+
+    def add_job(self, analysis_id: int, data: CreateJobRequest) -> Job:
+        """Add a new job to the database."""
+        analysis: Analysis = self.get_analysis_with_id(analysis_id)
+        job = Job(
+            name=data.name,
+            status=SlurmJobStatus.PENDING,
+            slurm_id=data.slurm_id,
+            job_type=data.job_type,
+        )
+        analysis.jobs.append(job)
+        session: Session = get_session()
+        session.commit()
+        return job

trailblazer/store/models.py

@@ -5,6 +5,7 @@
 from trailblazer.constants import (
     PRIORITY_OPTIONS,
     TYPES,
+    JobType,
     SlurmJobStatus,
     TrailblazerStatus,
     WorkflowManager,
@@ -136,6 +137,7 @@ class Job(Model):
     started_at = Column(types.DateTime)
     elapsed = Column(types.Integer)
     status = Column(types.Enum(*SlurmJobStatus.statuses()))
+    job_type = Column(types.Enum(*JobType.types()), default=JobType.ANALYSIS, nullable=False)
 
     def to_dict(self) -> dict:
         """Return a dictionary representation of the object."""