diff --git a/diracx-db/pyproject.toml b/diracx-db/pyproject.toml index f7a145fc5..3bdabfa2f 100644 --- a/diracx-db/pyproject.toml +++ b/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" diff --git a/diracx-db/src/diracx/db/sql/__init__.py b/diracx-db/src/diracx/db/sql/__init__.py index f98785e23..3be3af8a3 100644 --- a/diracx-db/src/diracx/db/sql/__init__.py +++ b/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 diff --git a/diracx-db/src/diracx/db/sql/pilot_agents/__init__.py b/diracx-db/src/diracx/db/sql/pilot_agents/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diracx-db/src/diracx/db/sql/pilot_agents/db.py b/diracx-db/src/diracx/db/sql/pilot_agents/db.py new file mode 100644 index 000000000..a5c9b6bcf --- /dev/null +++ b/diracx-db/src/diracx/db/sql/pilot_agents/db.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from ..utils import BaseSQLDB +from .schema import PilotAgentsDBBase + + +class PilotAgentsDB(BaseSQLDB): + metadata = PilotAgentsDBBase.metadata diff --git a/diracx-db/src/diracx/db/sql/pilot_agents/schema.py b/diracx-db/src/diracx/db/sql/pilot_agents/schema.py new file mode 100644 index 000000000..366e07745 --- /dev/null +++ b/diracx-db/src/diracx/db/sql/pilot_agents/schema.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) diff --git a/diracx-routers/src/diracx/routers/dependencies.py b/diracx-routers/src/diracx/routers/dependencies.py index 418aedfbc..9abe5f673 100644 --- a/diracx-routers/src/diracx/routers/dependencies.py +++ b/diracx-routers/src/diracx/routers/dependencies.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) ]