feat: Implement PilotAgents schema

Author: martynia

diracx-db/pyproject.toml

@@ -31,6 +31,7 @@ testing = [
 AuthDB = "diracx.db.sql:AuthDB"
 JobDB = "diracx.db.sql:JobDB"
 JobLoggingDB = "diracx.db.sql:JobLoggingDB"
+PilotAgentsDB = "diracx.db.sql:PilotAgentsDB"
 SandboxMetadataDB = "diracx.db.sql:SandboxMetadataDB"
 TaskQueueDB = "diracx.db.sql:TaskQueueDB"
 

diracx-db/src/diracx/db/sql/__init__.py

@@ -1,9 +1,17 @@
 from __future__ import annotations
 
-__all__ = ("AuthDB", "JobDB", "JobLoggingDB", "SandboxMetadataDB", "TaskQueueDB")
+__all__ = (
+    "AuthDB",
+    "JobDB",
+    "JobLoggingDB",
+    "PilotAgentsDB",
+    "SandboxMetadataDB",
+    "TaskQueueDB",
+)
 
 from .auth.db import AuthDB
 from .job.db import JobDB
 from .job_logging.db import JobLoggingDB
+from .pilot_agents.db import PilotAgentsDB
 from .sandbox_metadata.db import SandboxMetadataDB
 from .task_queue.db import TaskQueueDB

diracx-db/src/diracx/db/sql/pilot_agents/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from ..utils import BaseSQLDB
+from .schema import PilotAgentsDBBase
+
+
+class PilotAgentsDB(BaseSQLDB):
+    metadata = PilotAgentsDBBase.metadata

diracx-db/src/diracx/db/sql/pilot_agents/db.py

@@ -0,0 +1,65 @@
+from sqlalchemy import (
+    DateTime,
+    Double,
+    Index,
+    Integer,
+    String,
+    Text,
+)
+from sqlalchemy.orm import declarative_base
+
+from ..job.schema import EnumBackedBool
+from ..utils import Column, NullColumn
+
+PilotAgentsDBBase = declarative_base()
+
+
+class PilotAgents(PilotAgentsDBBase):
+    __tablename__ = "PilotAgents"
+
+    PilotID = Column("PilotID", Integer, autoincrement=True, primary_key=True)
+    InitialJobID = Column("InitialJobID", Integer, default=0)
+    CurrentJobID = Column("CurrentJobID", Integer, default=0)
+    TaskQueueID = Column("TaskQueueID", Integer, default=0)
+    PilotJobReference = Column("PilotJobReference", String(255), default="Unknown")
+    PilotStamp = Column("PilotStamp", String(32), default="")
+    DestinationSite = Column("DestinationSite", String(128), default="NotAssigned")
+    Queue = Column("Queue", String(128), default="Unknown")
+    GridSite = Column("GridSite", String(128), default="Unknown")
+    Broker = Column("Broker", String(128), default="Unknown")
+    OwnerDN = Column("OwnerDN", String(255))
+    OwnerGroup = Column("OwnerGroup", String(128))
+    GridType = Column("GridType", String(32), default="LCG")
+    GridRequirements = Column("GridRequirements", Text)
+    BenchMark = Column("BenchMark", Double, default=0.0)
+    SubmissionTime = NullColumn("SubmissionTime", DateTime)
+    LastUpdateTime = NullColumn("LastUpdateTime", DateTime)
+    Status = Column("Status", String(32), default="Unknown")
+    StatusReason = Column("StatusReason", String(255), default="Unknown")
+    ParentID = Column("ParentID", Integer, default=0)
+    OutputReady = Column("OutputReady", EnumBackedBool(), default=False)
+    AccountingSent = Column("AccountingSent", EnumBackedBool(), default=False)
+
+    __table_args__ = (
+        Index("PilotJobReference", "PilotJobReference"),
+        Index("Status", "Status"),
+        Index("Statuskey", "GridSite", "DestinationSite", "Status"),
+    )
+
+
+class JobToPilotMapping(PilotAgentsDBBase):
+    __tablename__ = "JobToPilotMapping"
+
+    PilotID = Column("PilotID", Integer, primary_key=True)
+    JobID = Column("JobID", Integer, primary_key=True)
+    StartTime = Column("StartTime", DateTime)
+
+    __table_args__ = (Index("JobID", "JobID"), Index("PilotID", "PilotID"))
+
+
+class PilotOutput(PilotAgentsDBBase):
+    __tablename__ = "PilotOutput"
+
+    PilotID = Column("PilotID", Integer, primary_key=True)
+    StdOutput = Column("StdOutput", Text)
+    StdError = Column("StdError", Text)

diracx-db/src/diracx/db/sql/pilot_agents/schema.py

@@ -7,6 +7,7 @@
     "JobLoggingDB",
     "SandboxMetadataDB",
     "TaskQueueDB",
+    "PilotAgentsDB",
     "add_settings_annotation",
     "AvailableSecurityProperties",
 )
@@ -22,6 +23,7 @@
 from diracx.db.sql import AuthDB as _AuthDB
 from diracx.db.sql import JobDB as _JobDB
 from diracx.db.sql import JobLoggingDB as _JobLoggingDB
+from diracx.db.sql import PilotAgentsDB as _PilotAgentsDB
 from diracx.db.sql import SandboxMetadataDB as _SandboxMetadataDB
 from diracx.db.sql import TaskQueueDB as _TaskQueueDB
 
@@ -37,6 +39,7 @@ def add_settings_annotation(cls: T) -> T:
 AuthDB = Annotated[_AuthDB, Depends(_AuthDB.transaction)]
 JobDB = Annotated[_JobDB, Depends(_JobDB.transaction)]
 JobLoggingDB = Annotated[_JobLoggingDB, Depends(_JobLoggingDB.transaction)]
+PilotAgentsDB = Annotated[_PilotAgentsDB, Depends(_PilotAgentsDB.transaction)]
 SandboxMetadataDB = Annotated[
     _SandboxMetadataDB, Depends(_SandboxMetadataDB.transaction)
 ]