refactor: restructure to follow n-layer architecture

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/tests/test_rss.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import pytest
+
+from diracx.core.models.rss import (
+    AllowedStatus,
+    BannedStatus,
+    map_status,
+)
+
+
+@pytest.mark.parametrize("status", ["Active", "Degraded"])
+async def test_map_status_allowed(status):
+    assert bool(map_status(status, "")) is True
+    assert isinstance(map_status(status, ""), AllowedStatus)
+
+
+@pytest.mark.parametrize("status", ["Banned", "Probing", "Error", "Unknown"])
+async def test_map_status_banned(status):
+    result = map_status(status, "CE banned")
+    assert bool(result) is False
+    assert isinstance(result, BannedStatus)
+    assert result.reason == "CE banned"
+
+
+async def test_map_status_unknown_banned():
+    result = map_status("WeirdValue", "")
+    assert bool(result) is False
+    assert isinstance(result, BannedStatus)
+    assert result.reason == "Unknown status: WeirdValue"

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/db.py

@@ -1,17 +1,9 @@
 from __future__ import annotations
 
 from sqlalchemy import select
+from sqlalchemy.engine import Row
 
-from diracx.core.models.rss import (
-    ComputeElementStatus,
-    FTSStatus,
-    ResourceType,
-    StorageElementStatus,
-    map_status,
-)
-from diracx.core.models.rss import (
-    SiteStatus as SiteStatusModel,
-)
+from diracx.core.exceptions import ResourceNotFoundError
 
 from ..utils import BaseSQLDB
 from .schema import (
@@ -26,67 +18,35 @@ class ResourceStatusDB(BaseSQLDB):
 
     metadata = RSSBase.metadata
 
-    async def get_site_status(self, name: str, vo: str = "all") -> SiteStatusModel:
+    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.statustype == "all",
             SiteStatus.vo == vo,
         )
         result = await self.conn.execute(stmt)
-        row = result.first()
+        row = result.one_or_none()
         if not row:
-            raise ValueError(f"Site {name} not found")
+            raise ResourceNotFoundError(name)
 
-        return SiteStatusModel(all=map_status(row.Status, row.Reason))
+        return row.Status, row.Reason
 
     async def get_resource_status(
         self,
         name: str,
+        statustypes: list[str] = ["all"],
         vo: str = "all",
-    ) -> ComputeElementStatus | FTSStatus:
+    ) -> dict[str, Row]:
         stmt = select(
-            ResourceStatus.status, ResourceStatus.reason, ResourceStatus.elementtype
+            ResourceStatus.status, ResourceStatus.reason, ResourceStatus.statustype
         ).where(
             ResourceStatus.name == name,
-            ResourceStatus.statustype == "all",
+            ResourceStatus.statustype.in_(statustypes),
             ResourceStatus.vo == vo,
         )
         result = await self.conn.execute(stmt)
-        row = result.first()
-
-        if not row:
-            raise ValueError(f"Resource {name} not found")
-
-        element_type = ResourceType(row.ElementType)
-
-        if element_type == ResourceType.Compute:
-            return ComputeElementStatus(all=map_status(row.Status, row.Reason))
-        if element_type == ResourceType.FTS:
-            return FTSStatus(all=map_status(row.Status, row.Reason))
+        rows = result.all()
 
-        raise ValueError(f"Unexpected resource type {element_type}")
-
-    async def get_storage_status(
-        self, name: str, vo: str = "all"
-    ) -> StorageElementStatus:
-        async def get_status(statustype: str):
-            stmt = select(ResourceStatus.status, ResourceStatus.reason).where(
-                ResourceStatus.name == name,
-                ResourceStatus.statustype == statustype,
-                ResourceStatus.vo == vo,
-            )
-
-            result = await self.conn.execute(stmt)
-            row = result.first()
-
-            if not row:
-                raise ValueError(f"StorageElement {name} not found")
-
-            return map_status(row.Status, row.Reason)
-
-        return StorageElementStatus(
-            read=await get_status("ReadAccess"),
-            write=await get_status("WriteAccess"),
-            check=await get_status("CheckAccess"),
-            remove=await get_status("RemoveAccess"),
-        )
+        if not rows:
+            raise ResourceNotFoundError(name)
+        return {row.StatusType: row for row in rows}

diracx-db/tests/rss/test_rss_db.py

@@ -1,20 +1,11 @@
-# diracx-db/tests/sql/rss/test_rss.py
 from __future__ import annotations
 
 from datetime import datetime, timezone
 
 import pytest
 from sqlalchemy import insert
 
-from diracx.core.models.rss import (
-    AllowedStatus,
-    BannedStatus,
-    ComputeElementStatus,
-    FTSStatus,
-    StorageElementStatus,
-    map_status,
-)
-from diracx.core.models.rss import SiteStatus as SiteStatusModel
+from diracx.core.exceptions import ResourceNotFoundError
 from diracx.db.sql.rss.db import ResourceStatusDB
 
 _NOW = datetime(2024, 1, 1, tzinfo=timezone.utc)
@@ -30,27 +21,6 @@ async def rss_db(tmp_path):
         yield rss_db
 
 
-@pytest.mark.parametrize("status", ["Active", "Degraded"])
-async def test_map_status_allowed(status):
-    assert bool(map_status(status, "")) is True
-    assert isinstance(map_status(status, ""), AllowedStatus)
-
-
-@pytest.mark.parametrize("status", ["Banned", "Probing", "Error", "Unknown"])
-async def test_map_status_banned(status):
-    result = map_status(status, "CE banned")
-    assert bool(result) is False
-    assert isinstance(result, BannedStatus)
-    assert result.reason == "CE banned"
-
-
-async def test_map_status_unknown_banned():
-    result = map_status("WeirdValue", "")
-    assert bool(result) is False
-    assert isinstance(result, BannedStatus)
-    assert result.reason == "Unknown status: WeirdValue"
-
-
 async def test_site_status(rss_db: ResourceStatusDB):
     # Insert a test Site
     async with rss_db.engine.begin() as conn:
@@ -71,13 +41,12 @@ async def test_site_status(rss_db: ResourceStatusDB):
 
     # Test with the test Site (should be found)
     async with rss_db as rss_db:
-        result = await rss_db.get_site_status("TestSite")
-    assert isinstance(result, SiteStatusModel)
-    assert isinstance(result.all, AllowedStatus)
-    assert bool(result.all) is True
+        status, reason = await rss_db.get_site_status("TestSite")
+    assert status == "Active"
+    assert reason == "All good"
 
     # Test with an unknow Site (should not be found)
-    with pytest.raises(ValueError, match="Site Unknown not found"):
+    with pytest.raises(ResourceNotFoundError):
         async with rss_db as rss_db:
             await rss_db.get_site_status("Unknown")
 
@@ -114,50 +83,7 @@ async def test_resource_status(rss_db: ResourceStatusDB):
                 TokenOwner="test",
             )
         )
-        # Insert a wrong test
-        await conn.execute(
-            insert(rss_db.metadata.tables["ResourceStatus"]).values(
-                Name="WrongTest",
-                StatusType="all",
-                VO="all",
-                Status="Active",
-                Reason="All good",
-                DateEffective=_NOW,
-                TokenExpiration=_FAR,
-                LastCheckTime=_NOW,
-                ElementType="WrongTest",
-                TokenOwner="WrongTest",
-            )
-        )
-
-    # Test with the test Compute Element (should be found)
-    async with rss_db as rss_db:
-        result = await rss_db.get_resource_status("TestCompute")
-    assert isinstance(result, ComputeElementStatus)
-    assert isinstance(result.all, AllowedStatus)
-    assert bool(result.all) is True
-
-    # Test with the test FTS (should be found)
-    async with rss_db as rss_db:
-        result = await rss_db.get_resource_status("TestFTS")
-    assert isinstance(result, FTSStatus)
-    assert isinstance(result.all, AllowedStatus)
-    assert bool(result.all) is True
-
-    # Test with a wrong Resource type
-    with pytest.raises(ValueError, match="not a valid ResourceType"):
-        async with rss_db as rss_db:
-            await rss_db.get_resource_status("WrongTest")
-
-    # Test with an unknow Resource (should not be found)
-    with pytest.raises(ValueError, match="Resource Unknown not found"):
-        async with rss_db as rss_db:
-            await rss_db.get_resource_status("Unknown")
-
-
-async def test_storage_status(rss_db: ResourceStatusDB):
-    # Insert a test Storage Element with all StatusType
-    async with rss_db.engine.begin() as conn:
+        # 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(
@@ -174,20 +100,36 @@ async def test_storage_status(rss_db: ResourceStatusDB):
                 )
             )
 
+    # Test with the test Compute Element (should be found)
+    async with rss_db as rss_db:
+        result = await rss_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 rss_db:
+        result = await rss_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 rss_db:
-        result = await rss_db.get_storage_status("TestStorage")
-    assert isinstance(result, StorageElementStatus)
-    assert isinstance(result.read, AllowedStatus)
-    assert isinstance(result.write, AllowedStatus)
-    assert isinstance(result.check, AllowedStatus)
-    assert isinstance(result.remove, AllowedStatus)
-    assert bool(result.read) is True
-    assert bool(result.write) is True
-    assert bool(result.check) is True
-    assert bool(result.remove) is True
-
-    # Test with an unknow Storage Element (should not be found)
-    with pytest.raises(ValueError, match="StorageElement Unknown not found"):
+        result = await rss_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 rss_db:
-            await rss_db.get_storage_status("Unknown")
+            await rss_db.get_resource_status("Unknown")

diracx-logic/src/diracx/logic/rss/__init__.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from diracx.core.models.rss import (
+    ComputeElementStatus,
+    FTSStatus,
+    StorageElementStatus,
+    map_status,
+)
+from diracx.core.models.rss import (
+    SiteStatus as SiteStatusModel,
+)
+from diracx.db.sql import ResourceStatusDB
+
+
+async def get_site_status(
+    resource_status_db: ResourceStatusDB, name: str, vo: str
+) -> SiteStatusModel:
+    status, reason = await resource_status_db.get_site_status(name, vo)
+    return SiteStatusModel(all=map_status(status, reason))
+
+
+async def get_compute_status(
+    resource_status_db: ResourceStatusDB, name: str, vo: str
+) -> ComputeElementStatus:
+    rows = await resource_status_db.get_resource_status(name, ["all"], vo)
+    return ComputeElementStatus(all=map_status(rows["all"].Status, rows["all"].Reason))
+
+
+async def get_fts_status(
+    resource_status_db: ResourceStatusDB, name: str, vo: str
+) -> FTSStatus:
+    rows = await resource_status_db.get_resource_status(name, ["all"], vo)
+    return FTSStatus(all=map_status(rows["all"].Status, rows["all"].Reason))
+
+
+async def get_storage_status(
+    resource_status_db: ResourceStatusDB, name: str, vo: str
+) -> StorageElementStatus:
+    rows = await resource_status_db.get_resource_status(
+        name, ["ReadAccess", "WriteAccess", "CheckAccess", "RemoveAccess"], vo
+    )
+    return StorageElementStatus(
+        read=map_status(rows["ReadAccess"].Status, rows["ReadAccess"].Reason),
+        write=map_status(rows["WriteAccess"].Status, rows["WriteAccess"].Reason),
+        check=map_status(rows["CheckAccess"].Status, rows["CheckAccess"].Reason),
+        remove=map_status(rows["RemoveAccess"].Status, rows["RemoveAccess"].Reason),
+    )