diracx-db 0.0.1a18__py3-none-any.whl → 0.0.1a19__py3-none-any.whl

diracx/db/__main__.py

@@ -31,6 +31,7 @@ async def init_sql():
     from diracx.db.sql.utils import BaseSQLDB
 
     for db_name, db_url in BaseSQLDB.available_urls().items():
+
         logger.info("Initialising %s", db_name)
         db = BaseSQLDB.available_implementations(db_name)[0](db_url)
         async with db.engine_context():

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/sql/job/schema.py

@@ -1,4 +1,3 @@
-import sqlalchemy.types as types
 from sqlalchemy import (
     DateTime,
     Enum,
@@ -10,37 +9,11 @@ from sqlalchemy import (
 )
 from sqlalchemy.orm import declarative_base
 
-from ..utils import Column, NullColumn
+from ..utils import Column, EnumBackedBool, NullColumn
 
 JobDBBase = declarative_base()
 
 
-class EnumBackedBool(types.TypeDecorator):
-    """Maps a ``EnumBackedBool()`` column to True/False in Python."""
-
-    impl = types.Enum
-    cache_ok: bool = True
-
-    def __init__(self) -> None:
-        super().__init__("True", "False")
-
-    def process_bind_param(self, value, dialect) -> str:
-        if value is True:
-            return "True"
-        elif value is False:
-            return "False"
-        else:
-            raise NotImplementedError(value, dialect)
-
-    def process_result_value(self, value, dialect) -> bool:
-        if value == "True":
-            return True
-        elif value == "False":
-            return False
-        else:
-            raise NotImplementedError(f"Unknown {value=}")
-
-
 class Jobs(JobDBBase):
     __tablename__ = "Jobs"
 
@@ -134,24 +107,6 @@ class AtticJobParameters(JobDBBase):
     RescheduleCycle = Column(Integer)
 
 
-class SiteMask(JobDBBase):
-    __tablename__ = "SiteMask"
-    Site = Column(String(64), primary_key=True)
-    Status = Column(String(64))
-    LastUpdateTime = Column(DateTime)
-    Author = Column(String(255))
-    Comment = Column(Text)
-
-
-class SiteMaskLogging(JobDBBase):
-    __tablename__ = "SiteMaskLogging"
-    Site = Column(String(64), primary_key=True)
-    UpdateTime = Column(DateTime, primary_key=True)
-    Status = Column(String(64))
-    Author = Column(String(255))
-    Comment = Column(Text)
-
-
 class HeartBeatLoggingInfo(JobDBBase):
     __tablename__ = "HeartBeatLoggingInfo"
     JobID = Column(

diracx/db/sql/pilot_agents/db.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+from sqlalchemy import insert
+
+from ..utils import BaseSQLDB
+from .schema import PilotAgents, PilotAgentsDBBase
+
+
+class PilotAgentsDB(BaseSQLDB):
+    """PilotAgentsDB class is a front-end to the PilotAgents Database."""
+
+    metadata = PilotAgentsDBBase.metadata
+
+    async def add_pilot_references(
+        self,
+        pilot_ref: list[str],
+        vo: str,
+        grid_type: str = "DIRAC",
+        pilot_stamps: dict | None = None,
+    ) -> None:
+
+        if pilot_stamps is None:
+            pilot_stamps = {}
+
+        now = datetime.now(tz=timezone.utc)
+
+        # Prepare the list of dictionaries for bulk insertion
+        values = [
+            {
+                "PilotJobReference": ref,
+                "VO": vo,
+                "GridType": grid_type,
+                "SubmissionTime": now,
+                "LastUpdateTime": now,
+                "Status": "Submitted",
+                "PilotStamp": pilot_stamps.get(ref, ""),
+            }
+            for ref in pilot_ref
+        ]
+
+        # Insert multiple rows in a single execute call
+        stmt = insert(PilotAgents).values(values)
+        await self.conn.execute(stmt)
+        return

diracx/db/sql/pilot_agents/schema.py

@@ -0,0 +1,58 @@
+from sqlalchemy import (
+    DateTime,
+    Double,
+    Index,
+    Integer,
+    String,
+    Text,
+)
+from sqlalchemy.orm import declarative_base
+
+from ..utils import Column, EnumBackedBool, 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)
+    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")
+    VO = Column("VO", String(128))
+    GridType = Column("GridType", String(32), default="LCG")
+    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")
+    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/sql/utils/__init__.py

@@ -13,6 +13,7 @@ from datetime import datetime, timedelta, timezone
 from functools import partial
 from typing import TYPE_CHECKING, Self, cast
 
+import sqlalchemy.types as types
 from pydantic import TypeAdapter
 from sqlalchemy import Column as RawColumn
 from sqlalchemy import DateTime, Enum, MetaData, select
@@ -128,6 +129,32 @@ def EnumColumn(enum_type, **kwargs):
     return Column(Enum(enum_type, native_enum=False, length=16), **kwargs)
 
 
+class EnumBackedBool(types.TypeDecorator):
+    """Maps a ``EnumBackedBool()`` column to True/False in Python."""
+
+    impl = types.Enum
+    cache_ok: bool = True
+
+    def __init__(self) -> None:
+        super().__init__("True", "False")
+
+    def process_bind_param(self, value, dialect) -> str:
+        if value is True:
+            return "True"
+        elif value is False:
+            return "False"
+        else:
+            raise NotImplementedError(value, dialect)
+
+    def process_result_value(self, value, dialect) -> bool:
+        if value == "True":
+            return True
+        elif value == "False":
+            return False
+        else:
+            raise NotImplementedError(f"Unknown {value=}")
+
+
 class SQLDBError(Exception):
     pass
 

{diracx_db-0.0.1a18.dist-info → diracx_db-0.0.1a19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diracx-db
-Version: 0.0.1a18
+Version: 0.0.1a19
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research

{diracx_db-0.0.1a18.dist-info → diracx_db-0.0.1a19.dist-info}/RECORD

@@ -1,11 +1,11 @@
 diracx/db/__init__.py,sha256=2oeUeVwZq53bo_ZOflEYZsBn7tcR5Tzb2AIu0TAWELM,109
-diracx/db/__main__.py,sha256=3yaUP1ig-yaPSQM4wy6CtSXXHivQg-hIz2FeBt7joBc,1714
+diracx/db/__main__.py,sha256=tU4tp3OAClYCiPMxlRj524sZGBx9oy4CoWHd8pMuEEs,1715
 diracx/db/exceptions.py,sha256=-LSkEwsvjwU7vXqx-xeLvLKInTRAhjwB7K_AKfQcIH8,41
 diracx/db/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/os/__init__.py,sha256=IZr6z6SefrRvuC8sTC4RmB3_wwOyEt1GzpDuwSMH8O4,112
 diracx/db/os/job_parameters.py,sha256=Knca19uT2G-5FI7MOFlaOAXeHn4ecPVLIH30TiwhaTw,858
 diracx/db/os/utils.py,sha256=9UyhgMqaI8loh8chW2zHW-9JAOtH5YfktC-d-uY5Wnk,11346
-diracx/db/sql/__init__.py,sha256=oJKlsWofwDMinJ5dkRy6BuP1IuVXgrmQ5HHBLpLf7no,304
+diracx/db/sql/__init__.py,sha256=JYu0b0IVhoXy3lX2m2r2dmAjsRS7IbECBUMEDvX0Te4,391
 diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/auth/db.py,sha256=Wi4oeHCL4pPiaysZEx8R0KNk9BDxncAAtOD2qnD-NnY,10206
 diracx/db/sql/auth/schema.py,sha256=W5whp1ZK_SNt-wxWVRBegmrc9IgqCR1LFY1FWwUlEBs,2828
@@ -14,20 +14,23 @@ diracx/db/sql/dummy/db.py,sha256=4Xyo7gUh_5b6Q2a_ggJG6e7fCtc9HrP_BRXfKGfqZIs,164
 diracx/db/sql/dummy/schema.py,sha256=uEkGDNVZbmJecytkHY1CO-M1MiKxe5w1_h0joJMPC9E,680
 diracx/db/sql/job/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/job/db.py,sha256=DaU1SGeXl7TqX1QxT1RldCeMIOYGnzNwkqBrwgGE90A,16248
-diracx/db/sql/job/schema.py,sha256=3FKVc8BlPg08vAALcFS2HPGGu7QExlBuPIDOqEnBNvI,5519
+diracx/db/sql/job/schema.py,sha256=w9Ht9LyVK-fB5T9-hYGsqifzneeG2YP123j1-Mx8Xio,4283
 diracx/db/sql/job_logging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/job_logging/db.py,sha256=OGVHYIDcWhGYVfHacsz9DEPSoJ7aRbKVoQOyCCj8XvU,5036
 diracx/db/sql/job_logging/schema.py,sha256=dD2arl-6bffeK8INT6tZ1HWEpJuYTx2iNiVzswVXXF8,812
+diracx/db/sql/pilot_agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+diracx/db/sql/pilot_agents/db.py,sha256=7-cuCbh_KhM0jlybsHMWV-W66bHsPHIVBpbuqwjncj0,1232
+diracx/db/sql/pilot_agents/schema.py,sha256=d4sVlhWpMVgHguS-VlDR3zqM0Yj5QVEjon3gBgrr-kk,2091
 diracx/db/sql/sandbox_metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/sandbox_metadata/db.py,sha256=ZoiKjs-Rdd4p-kOZOLoUkpP0xGIp0oGjou3D2oP_6GE,6452
 diracx/db/sql/sandbox_metadata/schema.py,sha256=rngYYkJxBhjETBHGLD1CTipDGe44mRYR0wdaFoAJwp0,1400
 diracx/db/sql/task_queue/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/task_queue/db.py,sha256=e6yauZO0nWaUVqjqQycH8iPO4wXLXaC82eaIq1K_KI8,9102
 diracx/db/sql/task_queue/schema.py,sha256=fvzQyCw_xWAOWTLW6Qrp1m-WzEKb0tlYmafoLTbCy1I,3222
-diracx/db/sql/utils/__init__.py,sha256=pS0886KOZZvGDurQU4peJ4RGSCMH9umCx95qp66OSpY,14793
+diracx/db/sql/utils/__init__.py,sha256=8YcDra_P_FL3QIDzl53HvN1tHUE396F0tGKB-TPQxeM,15541
 diracx/db/sql/utils/job_status.py,sha256=GNQTKiyguhnB348mLIB7BT-PEOEKpKljR4JzvOd_h8M,10414
-diracx_db-0.0.1a18.dist-info/METADATA,sha256=5SEgwo0STgA0X_475cghDUlzvqRVnmBipxztPG4Ya8E,688
-diracx_db-0.0.1a18.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-diracx_db-0.0.1a18.dist-info/entry_points.txt,sha256=xEFGu_zgmPgQPlUeFtdahQfQIboJ1ugFOK8eMio9gtw,271
-diracx_db-0.0.1a18.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
-diracx_db-0.0.1a18.dist-info/RECORD,,
+diracx_db-0.0.1a19.dist-info/METADATA,sha256=_e0zgiSjl92T83BUpBXkEHzMeCcLcr4L_sE9h9v_iwQ,688
+diracx_db-0.0.1a19.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+diracx_db-0.0.1a19.dist-info/entry_points.txt,sha256=YLI4f6640bri8Ud6Jt9WNq79pSTVQAkfUasb9f75fR8,315
+diracx_db-0.0.1a19.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
+diracx_db-0.0.1a19.dist-info/RECORD,,

{diracx_db-0.0.1a18.dist-info → diracx_db-0.0.1a19.dist-info}/entry_points.txt

@@ -5,5 +5,6 @@ JobParametersDB = diracx.db.os:JobParametersDB
 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