fix: test and logic

Author: TheauW

diracx-core/src/diracx/core/models.py

@@ -73,6 +73,71 @@ class JobSearchParams(BaseModel):
     # TODO: Add more validation
 
 
+class JobParameters(BaseModel):
+    """All the parameters that can be set for a job."""
+
+    timestamp: datetime | None = None
+    CPUNormalizationFactor: int | None = None
+    NormCPUTime_s: int | None = Field(None, alias="NormCPUTime(s)")
+    TotalCPUTime_s: int | None = Field(None, alias="TotalCPUTime(s)")
+    HostName: str | None = None
+    GridCE: str | None = None
+    ModelName: str | None = None
+    PilotAgent: str | None = None
+    Pilot_Reference: str | None = None
+    Memory_MB: int | None = Field(None, alias="Memory(MB)")
+    LocalAccount: str | None = None
+    PayloadPID: int | None = None
+    CEQueue: str | None = None
+    BatchSystem: str | None = None
+    JobType: str | None = None
+    JobStatus: str | None = None
+
+    class Config:
+        """Configuration for the JobParameters model."""
+
+        extra = "forbid"  # Disallow additional fields
+
+
+class JobAttributes(BaseModel):
+    """All the attributes that can be set for a job."""
+
+    JobType: str | None = None
+    JobGroup: str | None = None
+    Site: str | None = None
+    JobName: str | None = None
+    Owner: str | None = None
+    OwnerGroup: str | None = None
+    VO: str | None = None
+    SubmissionTime: datetime | None = None
+    RescheduleTime: datetime | None = None
+    LastUpdateTime: datetime | None = None
+    StartExecTime: datetime | None = None
+    HeartBeatTime: datetime | None = None
+    EndExecTime: datetime | None = None
+    Status: str | None = None
+    MinorStatus: str | None = None
+    ApplicationStatus: str | None = None
+    UserPriority: int | None = None
+    RescheduleCounter: int | None = None
+    VerifiedFlag: bool | None = None
+    AccountedFlag: bool | str | None = None
+
+    class Config:
+        """Configuration for the JobAttributes model."""
+
+        extra = "forbid"  # Disallow additional fields
+
+
+class JobMetaData(JobAttributes, JobParameters):
+    """A model that combines both JobAttributes and JobParameters."""
+
+    class Config:
+        """Configuration for the JobMetaData model."""
+
+        extra = "forbid"  # Disallow additional fields
+
+
 class JobStatus(StrEnum):
     SUBMITTING = "Submitting"
     RECEIVED = "Received"

diracx-logic/src/diracx/logic/jobs/status.py

@@ -31,18 +31,20 @@
 from diracx.core.config.schema import Config
 from diracx.core.models import (
     HeartbeatData,
+    JobAttributes,
     JobCommand,
     JobLoggingRecord,
+    JobMetaData,
     JobMinorStatus,
+    JobParameters,
     JobStatus,
     JobStatusUpdate,
     SetJobStatusReturn,
     VectorSearchOperator,
     VectorSearchSpec,
 )
 from diracx.db.os.job_parameters import JobParametersDB
-from diracx.db.sql.job.db import JobDB, _get_columns
-from diracx.db.sql.job.schema import Jobs
+from diracx.db.sql.job.db import JobDB
 from diracx.db.sql.job_logging.db import JobLoggingDB
 from diracx.db.sql.sandbox_metadata.db import SandboxMetadataDB
 from diracx.db.sql.task_queue.db import TaskQueueDB
@@ -52,6 +54,10 @@
 
 logger = logging.getLogger(__name__)
 
+# Precalculate valid field sets for performance
+VALID_JOB_ATTRIBUTES = set(JobAttributes.model_fields.keys())
+VALID_JOB_PARAMETERS = set(JobParameters.model_fields.keys())
+
 
 async def remove_jobs(
     job_ids: list[int],
@@ -502,39 +508,27 @@ async def remove_jobs_from_task_queue(
 
 
 async def set_job_parameters_or_attributes(
-    updates: dict[int, dict[str, Any]],
+    updates: dict[int, JobMetaData],
     job_db: JobDB,
     job_parameters_db: JobParametersDB,
 ):
     """Set job parameters or attributes for a list of jobs."""
-    attribute_columns: list[str] = [
-        col.name for col in _get_columns(Jobs.__table__, None)
-    ]
-    attribute_columns_lower: list[str] = [col.lower() for col in attribute_columns]
-
+    # Those dicts create a mapping of job_id -> {attribute_name: value}
     attr_updates: dict[int, dict[str, Any]] = {}
     param_updates: dict[int, dict[str, Any]] = {}
 
     for job_id, metadata in updates.items():
         attr_updates[job_id] = {}
         param_updates[job_id] = {}
-        for pname, pvalue in metadata.items():
+        for pname, pvalue in metadata.model_dump(exclude_none=True).items():
             # An argument can be a job attribute and/or a job parameter
 
-            # If the attribute exactly matches one of the allowed columns, treat it as an attribute.
-            if pname in attribute_columns:
+            # Check if the argument is a valid job attribute
+            if pname in VALID_JOB_ATTRIBUTES:
                 attr_updates[job_id][pname] = pvalue
-            # Otherwise, if the lower-case version is valid, the user likely mis-cased the key.
-            elif pname.lower() in attribute_columns_lower:
-                correct_name = attribute_columns[
-                    attribute_columns_lower.index(pname.lower())
-                ]
-                raise ValueError(
-                    f"Attribute column '{pname}' is mis-cased. Did you mean '{correct_name}'?"
-                )
 
             # Check if the argument is a valid job parameter
-            if pname in job_parameters_db.fields:
+            if pname in VALID_JOB_PARAMETERS:
                 param_updates[job_id][pname] = pvalue
 
     # Bulk set job attributes if required

diracx-routers/src/diracx/routers/jobs/sandboxes.py

@@ -166,11 +166,11 @@ async def unassign_job_sandboxes(
 
 @router.delete("/sandbox")
 async def unassign_bulk_jobs_sandboxes(
-    jobs_ids: Annotated[list[int], Query()],
+    job_ids: Annotated[list[int], Query()],
     sandbox_metadata_db: SandboxMetadataDB,
     job_db: JobDB,
     check_permissions: CheckWMSPolicyCallable,
 ):
     """Delete bulk jobs sandbox mapping."""
-    await check_permissions(action=ActionType.MANAGE, job_db=job_db, job_ids=jobs_ids)
-    await unassign_jobs_sandboxes_bl(jobs_ids, sandbox_metadata_db)
+    await check_permissions(action=ActionType.MANAGE, job_db=job_db, job_ids=job_ids)
+    await unassign_jobs_sandboxes_bl(job_ids, sandbox_metadata_db)

diracx-routers/src/diracx/routers/jobs/status.py

@@ -9,6 +9,7 @@
 from diracx.core.models import (
     HeartbeatData,
     JobCommand,
+    JobMetaData,
     JobStatusUpdate,
     SetJobStatusReturn,
 )
@@ -291,15 +292,13 @@ async def reschedule_jobs(
 
 @router.patch("/metadata", status_code=HTTPStatus.NO_CONTENT)
 async def patch_metadata(
-    updates: Annotated[
-        dict[int, dict[str, Any]], Body(openapi_examples=EXAMPLE_METADATA)
-    ],
+    updates: Annotated[dict[int, JobMetaData], Body(openapi_examples=EXAMPLE_METADATA)],
     job_db: JobDB,
     job_parameters_db: JobParametersDB,
     check_permissions: CheckWMSPolicyCallable,
 ):
     """Update job metadata such as UserPriority, HeartBeatTime, JobType, etc.
-    The parameters are all the attributes of a job (except the ID).
+    The argument  are all the attributes/parameters of a job (except the ID).
     """
     await check_permissions(action=ActionType.MANAGE, job_db=job_db, job_ids=updates)
     try:

diracx-routers/tests/jobs/test_sandboxes.py

@@ -163,7 +163,7 @@ def test_assign_then_unassign_sandboxes_to_jobs(normal_user_client: TestClient):
 
     # Unassign sb to job:
     job_ids = [job_id]
-    r = normal_user_client.delete("/api/jobs/sandbox", params={"jobs_ids": job_ids})
+    r = normal_user_client.delete("/api/jobs/sandbox", params={"job_ids": job_ids})
     assert r.status_code == 200
 
     # Get the sb again, it should'nt be there anymore:

diracx-routers/tests/jobs/test_status.py

@@ -445,6 +445,7 @@ def test_reschedule_job_attr_update(normal_user_client: TestClient):
             assert successful_results[str(jid)]["RescheduleCounter"] == i + 1
         for jid in fail_resched_ids:
             assert str(jid) in failed_results, result
+            # FIXME
             # assert successful_results[jid]["Status"] == JobStatus.RECEIVED
             # assert successful_results[jid]["MinorStatus"] == "Job Rescheduled"
             # assert successful_results[jid]["RescheduleCounter"] == i + 1
@@ -851,7 +852,7 @@ def test_patch_metadata(normal_user_client: TestClient, valid_job_id: int):
         assert j["ApplicationStatus"] == "Unknown"
 
     # Act
-    hbt = str(datetime.now(timezone.utc))
+    hbt = datetime.now(timezone.utc).isoformat()
     r = normal_user_client.patch(
         "/api/jobs/metadata",
         json={
@@ -883,11 +884,14 @@ def test_patch_metadata(normal_user_client: TestClient, valid_job_id: int):
     )
     assert r.status_code == 200, r.json()
 
+    # TODO: This should be timezone aware
+    hbt1 = datetime.fromisoformat(r.json()[0]["HeartBeatTime"])
+    assert hbt1.tzinfo is None
+    hbt1 = hbt1.replace(tzinfo=timezone.utc)
+
     assert r.json()[0]["JobID"] == valid_job_id
     assert r.json()[0]["JobType"] == "VerySpecialIndeed"
-    assert datetime.fromisoformat(
-        r.json()[0]["HeartBeatTime"]
-    ) == datetime.fromisoformat(hbt)
+    assert hbt1 == datetime.fromisoformat(hbt)
     assert r.json()[0]["UserPriority"] == 2
 
 
@@ -929,9 +933,7 @@ def test_bad_patch_metadata(normal_user_client: TestClient, valid_job_id: int):
     )
 
     # Assert
-    assert r.status_code == 400, (
-        "PATCH metadata should 400 Bad Request if an attribute column's case is incorrect"
-    )
+    assert r.status_code == 422, r.text
 
 
 def test_diracx_476(normal_user_client: TestClient, valid_job_id: int):