feat: add RSS data models (SQLAlchemy + Pydantic) (#857)

* feat: add RSS data models (SQLAlchemy + Pydantic)

* fix: remove deprecated tables and add FTS ResourceType

* refactor: restructure to follow n-layer architecture

* chore: remove unnecessary comments

* fix: change map_status and clean up schema

Author: HeloiseJoffe

diracx-core/src/diracx/core/exceptions.py

@@ -70,6 +70,12 @@ def __init__(self, pfn: str, se_name: str, detail: str | None = None):
         )
 
 
+class ResourceNotFoundError(DiracError):
+    def __init__(self, name: str, detail: str | None = None):
+        self.name: str = name
+        super().__init__(f"{name} not found" + (f" ({detail})" if detail else ""))
+
+
 class SandboxAlreadyAssignedError(DiracError):
     def __init__(self, pfn: str, se_name: str, detail: str | None = None):
         self.pfn: str = pfn

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

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, Field
+
+
+class AllowedStatus(BaseModel):
+    allowed: Literal[True]
+    warnings: str | None = None
+
+    def __bool__(self) -> bool:
+        return True
+
+
+class BannedStatus(BaseModel):
+    allowed: Literal[False]
+    reason: str = "Unknown"
+
+    def __bool__(self) -> bool:
+        return False
+
+
+ResourceStatus = Annotated[
+    Union[AllowedStatus, BannedStatus],
+    Field(discriminator="allowed"),
+]
+
+
+class ResourceType(StrEnum):
+    Compute = "ComputeElement"
+    Storage = "StorageElement"
+    FTS = "FTS"
+
+
+class StorageElementStatus(BaseModel):
+    read: ResourceStatus
+    write: ResourceStatus
+    check: ResourceStatus
+    remove: ResourceStatus
+
+
+class ComputeElementStatus(BaseModel):
+    all: ResourceStatus
+
+
+class FTSStatus(BaseModel):
+    all: ResourceStatus
+
+
+class SiteStatus(BaseModel):
+    all: ResourceStatus
+
+
+ALLOWED = {"Active", "Degraded"}
+BANNED = {"Banned", "Probing", "Error", "Unknown"}

diracx-db/pyproject.toml

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

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

@@ -5,6 +5,7 @@
     "JobDB",
     "JobLoggingDB",
     "PilotAgentsDB",
+    "ResourceStatusDB",
     "SandboxMetadataDB",
     "TaskQueueDB",
 )
@@ -13,5 +14,6 @@
 from .job.db import JobDB
 from .job_logging.db import JobLoggingDB
 from .pilot_agents.db import PilotAgentsDB
+from .rss.db import ResourceStatusDB
 from .sandbox_metadata.db import SandboxMetadataDB
 from .task_queue.db import TaskQueueDB

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

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from sqlalchemy import select
+from sqlalchemy.engine import Row
+
+from diracx.core.exceptions import ResourceNotFoundError
+
+from ..utils import BaseSQLDB
+from .schema import (
+    ResourceStatus,
+    RSSBase,
+    SiteStatus,
+)
+
+
+class ResourceStatusDB(BaseSQLDB):
+    """Class that defines the interactions with the tables of the ResourceStatusDB."""
+
+    metadata = RSSBase.metadata
+
+    async def get_site_status(self, name: str, vo: str = "all") -> tuple[str, str]:
+        stmt = select(SiteStatus.status, SiteStatus.reason).where(
+            SiteStatus.name == name,
+            SiteStatus.status_type == "all",
+            SiteStatus.vo == vo,
+        )
+        result = await self.conn.execute(stmt)
+        row = result.one_or_none()
+        if not row:
+            raise ResourceNotFoundError(name)
+
+        return row.Status, row.Reason
+
+    async def get_resource_status(
+        self,
+        name: str,
+        status_types: list[str] | None = None,
+        vo: str = "all",
+    ) -> dict[str, Row]:
+        if not status_types:
+            status_types = ["all"]
+        stmt = select(
+            ResourceStatus.status, ResourceStatus.reason, ResourceStatus.status_type
+        ).where(
+            ResourceStatus.name == name,
+            ResourceStatus.status_type.in_(status_types),
+            ResourceStatus.vo == vo,
+        )
+        result = await self.conn.execute(stmt)
+        rows = result.all()
+
+        if not rows:
+            raise ResourceNotFoundError(name)
+        return {row.StatusType: row for row in rows}

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

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import BigInteger, String
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+from diracx.db.sql.utils.types import SmarterDateTime
+
+from ..utils import str32, str64, str128, str512
+
+# Defining the tables
+
+
+class RSSBase(DeclarativeBase):
+    type_annotation_map = {
+        str32: String(32),
+        str64: String(64),
+        str128: String(128),
+        str512: String(512),
+    }
+
+
+class ElementStatusBase:
+    name: Mapped[str64] = mapped_column("Name", primary_key=True)
+    status_type: Mapped[str128] = mapped_column(
+        "StatusType", server_default="all", primary_key=True
+    )
+    vo: Mapped[str64] = mapped_column("VO", primary_key=True, server_default="all")
+    status: Mapped[str] = mapped_column("Status", String(8), server_default="")
+    reason: Mapped[str512] = mapped_column("Reason", server_default="Unspecified")
+    date_effective: Mapped[datetime] = mapped_column("DateEffective", SmarterDateTime())
+    token_expiration: Mapped[datetime] = mapped_column(
+        "TokenExpiration", SmarterDateTime(), server_default="9999-12-31 23:59:59"
+    )
+    element_type: Mapped[str32] = mapped_column("ElementType", server_default="")
+    last_check_time: Mapped[datetime] = mapped_column(
+        "LastCheckTime", SmarterDateTime(), server_default="1000-01-01 00:00:00"
+    )
+    token_owner: Mapped[str] = mapped_column(
+        "TokenOwner", String(16), server_default="rs_svc"
+    )
+
+
+class ElementStatusBaseWithID(ElementStatusBase):
+    """Almost the same as ElementStatusBase.
+
+    Differences:
+    - there's an autoincrement ID column which is also the primary key
+    - the name and statusType components are not part of the primary key
+    """
+
+    id: Mapped[int] = mapped_column(
+        "ID", BigInteger, autoincrement=True, primary_key=True
+    )
+    name: Mapped[str64] = mapped_column("Name")
+    status_type: Mapped[str128] = mapped_column("StatusType", server_default="all")
+    vo: Mapped[str64] = mapped_column("VO", server_default="all")
+    status: Mapped[str] = mapped_column("Status", String(8), server_default="")
+    reason: Mapped[str512] = mapped_column("Reason", server_default="Unspecified")
+    date_effective: Mapped[datetime] = mapped_column("DateEffective", SmarterDateTime())
+    token_expiration: Mapped[datetime] = mapped_column(
+        "TokenExpiration", SmarterDateTime(), server_default="9999-12-31 23:59:59"
+    )
+    element_type: Mapped[str32] = mapped_column("ElementType", server_default="")
+    last_check_time: Mapped[datetime] = mapped_column(
+        "LastCheckTime", SmarterDateTime(), server_default="1000-01-01 00:00:00"
+    )
+    token_owner: Mapped[str] = mapped_column(
+        "TokenOwner", String(16), server_default="rs_svc"
+    )
+
+
+# tables with schema defined in ElementStatusBase
+
+
+class SiteStatus(ElementStatusBase, RSSBase):
+    __tablename__ = "SiteStatus"
+
+
+class ResourceStatus(ElementStatusBase, RSSBase):
+    __tablename__ = "ResourceStatus"
+
+
+# tables with schema defined in ElementStatusBaseWithID
+
+
+class SiteLog(ElementStatusBaseWithID, RSSBase):
+    __tablename__ = "SiteLog"
+
+
+class SiteHistory(ElementStatusBaseWithID, RSSBase):
+    __tablename__ = "SiteHistory"
+
+
+class ResourceLog(ElementStatusBaseWithID, RSSBase):
+    __tablename__ = "ResourceLog"
+
+
+class ResourceHistory(ElementStatusBaseWithID, RSSBase):
+    __tablename__ = "ResourceHistory"

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

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+import pytest
+from sqlalchemy import insert
+
+from diracx.core.exceptions import ResourceNotFoundError
+from diracx.db.sql.rss.db import ResourceStatusDB
+
+_NOW = datetime(2024, 1, 1, tzinfo=timezone.utc)
+_FAR = datetime(9999, 12, 31, 23, 59, 59, tzinfo=timezone.utc)
+
+
+@pytest.fixture
+async def rss_db(tmp_path):
+    rss_db = ResourceStatusDB("sqlite+aiosqlite:///:memory:")
+    async with rss_db.engine_context():
+        async with rss_db.engine.begin() as conn:
+            await conn.run_sync(rss_db.metadata.create_all)
+        yield rss_db
+
+
+async def test_site_status(rss_db: ResourceStatusDB):
+    # Insert a test Site
+    async with rss_db.engine.begin() as conn:
+        await conn.execute(
+            insert(rss_db.metadata.tables["SiteStatus"]).values(
+                Name="TestSite",
+                StatusType="all",
+                VO="all",
+                Status="Active",
+                Reason="All good",
+                DateEffective=_NOW,
+                TokenExpiration=_FAR,
+                LastCheckTime=_NOW,
+                ElementType="Site",
+                TokenOwner="test",
+            )
+        )
+
+    # Test with the test Site (should be found)
+    async with rss_db as db:
+        status, reason = await db.get_site_status("TestSite")
+    assert status == "Active"
+    assert reason == "All good"
+
+    # Test with an unknow Site (should not be found)
+    with pytest.raises(ResourceNotFoundError):
+        async with rss_db as db:
+            await db.get_site_status("Unknown")
+
+
+async def test_resource_status(rss_db: ResourceStatusDB):
+    async with rss_db.engine.begin() as conn:
+        # Insert a test Compute Element
+        await conn.execute(
+            insert(rss_db.metadata.tables["ResourceStatus"]).values(
+                Name="TestCompute",
+                StatusType="all",
+                VO="all",
+                Status="Active",
+                Reason="All good",
+                DateEffective=_NOW,
+                TokenExpiration=_FAR,
+                LastCheckTime=_NOW,
+                ElementType="ComputeElement",
+                TokenOwner="test",
+            )
+        )
+        # Insert a test FTS
+        await conn.execute(
+            insert(rss_db.metadata.tables["ResourceStatus"]).values(
+                Name="TestFTS",
+                StatusType="all",
+                VO="all",
+                Status="Active",
+                Reason="All good",
+                DateEffective=_NOW,
+                TokenExpiration=_FAR,
+                LastCheckTime=_NOW,
+                ElementType="FTS",
+                TokenOwner="test",
+            )
+        )
+        # Insert a test Storage Element with all StatusType
+        for statustype in ["ReadAccess", "WriteAccess", "CheckAccess", "RemoveAccess"]:
+            await conn.execute(
+                insert(rss_db.metadata.tables["ResourceStatus"]).values(
+                    Name="TestStorage",
+                    StatusType=statustype,
+                    VO="all",
+                    Status="Active",
+                    Reason="All good",
+                    DateEffective=_NOW,
+                    TokenExpiration=_FAR,
+                    LastCheckTime=_NOW,
+                    ElementType="StorageElement",
+                    TokenOwner="test",
+                )
+            )
+
+    # Test with the test Compute Element (should be found)
+    async with rss_db as db:
+        result = await db.get_resource_status("TestCompute")
+    assert "all" in result
+    assert result["all"].Status == "Active"
+    assert result["all"].Reason == "All good"
+
+    # Test with the test FTS (should be found)
+    async with rss_db as db:
+        result = await db.get_resource_status("TestFTS")
+    assert "all" in result
+    assert result["all"].Status == "Active"
+    assert result["all"].Reason == "All good"
+
+    # Test with the test Storage Element (should be found)
+    async with rss_db as db:
+        result = await db.get_resource_status(
+            "TestStorage", ["ReadAccess", "WriteAccess", "CheckAccess", "RemoveAccess"]
+        )
+    assert set(result.keys()) == {
+        "ReadAccess",
+        "WriteAccess",
+        "CheckAccess",
+        "RemoveAccess",
+    }
+    for row in result.values():
+        assert row.Status == "Active"
+        assert row.Reason == "All good"
+
+    # Test with an unknow Resource (should not be found)
+    with pytest.raises(ResourceNotFoundError):
+        async with rss_db as db:
+            await db.get_resource_status("Unknown")