diracx-db 0.0.1a18__tar.gz → 0.0.1a20__tar.gz

{diracx_db-0.0.1a18 → diracx_db-0.0.1a20}/PKG-INFO

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

{diracx_db-0.0.1a18 → diracx_db-0.0.1a20}/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"
 
@@ -47,6 +48,7 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools_scm]
 root = ".."
 
+
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = [

{diracx_db-0.0.1a18 → diracx_db-0.0.1a20}/src/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-0.0.1a18 → diracx_db-0.0.1a20}/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-0.0.1a18 → diracx_db-0.0.1a20}/src/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-0.0.1a20/src/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-0.0.1a20/src/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-0.0.1a18 → diracx_db-0.0.1a20}/src/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 → diracx_db-0.0.1a20}/src/diracx_db.egg-info/PKG-INFO

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

{diracx_db-0.0.1a18 → diracx_db-0.0.1a20}/src/diracx_db.egg-info/SOURCES.txt

@@ -20,6 +20,9 @@ src/diracx/db/sql/job/schema.py
 src/diracx/db/sql/job_logging/__init__.py
 src/diracx/db/sql/job_logging/db.py
 src/diracx/db/sql/job_logging/schema.py
+src/diracx/db/sql/pilot_agents/__init__.py
+src/diracx/db/sql/pilot_agents/db.py
+src/diracx/db/sql/pilot_agents/schema.py
 src/diracx/db/sql/sandbox_metadata/__init__.py
 src/diracx/db/sql/sandbox_metadata/db.py
 src/diracx/db/sql/sandbox_metadata/schema.py
@@ -43,4 +46,6 @@ tests/jobs/test_jobLoggingDB.py
 tests/jobs/test_sandbox_metadata.py
 tests/opensearch/test_connection.py
 tests/opensearch/test_index_template.py
-tests/opensearch/test_search.py
+tests/opensearch/test_search.py
+tests/pilot_agents/__init__.py
+tests/pilot_agents/test_pilotAgentsDB.py

{diracx_db-0.0.1a18 → diracx_db-0.0.1a20}/src/diracx_db.egg-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

diracx_db-0.0.1a20/tests/pilot_agents/test_pilotAgentsDB.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import pytest
+
+from diracx.db.sql.pilot_agents.db import PilotAgentsDB
+
+
+@pytest.fixture
+async def pilot_agents_db(tmp_path) -> PilotAgentsDB:
+    agents_db = PilotAgentsDB("sqlite+aiosqlite:///:memory:")
+    async with agents_db.engine_context():
+        async with agents_db.engine.begin() as conn:
+            await conn.run_sync(agents_db.metadata.create_all)
+        yield agents_db
+
+
+async def test_insert_and_select(pilot_agents_db: PilotAgentsDB):
+
+    async with pilot_agents_db as pilot_agents_db:
+        # Add a pilot reference
+        refs = [f"ref_{i}" for i in range(10)]
+        stamps = [f"stamp_{i}" for i in range(10)]
+        stamp_dict = dict(zip(refs, stamps))
+
+        await pilot_agents_db.add_pilot_references(
+            refs, "test_vo", grid_type="DIRAC", pilot_stamps=stamp_dict
+        )
+
+        await pilot_agents_db.add_pilot_references(
+            refs, "test_vo", grid_type="DIRAC", pilot_stamps=None
+        )