diracx-db 0.0.1a21__py3-none-any.whl → 0.0.6__py3-none-any.whl

diracx/db/sql/sandbox_metadata/db.py

@@ -1,32 +1,73 @@
 from __future__ import annotations
 
-from typing import Any
+import logging
+from contextlib import asynccontextmanager
+from functools import partial
+from typing import Any, AsyncGenerator
 
-import sqlalchemy
+from sqlalchemy import (
+    BigInteger,
+    Column,
+    Executable,
+    MetaData,
+    Table,
+    and_,
+    delete,
+    exists,
+    insert,
+    literal,
+    or_,
+    select,
+    update,
+)
+from sqlalchemy.exc import IntegrityError, NoResultFound
 
+from diracx.core.exceptions import (
+    SandboxAlreadyAssignedError,
+    SandboxAlreadyInsertedError,
+    SandboxNotFoundError,
+)
 from diracx.core.models import SandboxInfo, SandboxType, UserInfo
-from diracx.db.sql.utils import BaseSQLDB, utcnow
+from diracx.db.sql.utils.base import BaseSQLDB
+from diracx.db.sql.utils.functions import days_since, utcnow
 
 from .schema import Base as SandboxMetadataDBBase
-from .schema import sb_EntityMapping, sb_Owners, sb_SandBoxes
+from .schema import SandBoxes, SBEntityMapping, SBOwners
+
+logger = logging.getLogger(__name__)
 
 
 class SandboxMetadataDB(BaseSQLDB):
     metadata = SandboxMetadataDBBase.metadata
 
-    async def upsert_owner(self, user: UserInfo) -> int:
+    # Temporary table to store the sandboxes to delete, see `select_and_delete_expired`
+    _temp_table = Table(
+        "sb_to_delete",
+        MetaData(),
+        Column("SBId", BigInteger, primary_key=True),
+        prefixes=["TEMPORARY"],
+    )
+
+    async def get_owner_id(self, user: UserInfo) -> int | None:
         """Get the id of the owner from the database."""
-        # TODO: Follow https://github.com/DIRACGrid/diracx/issues/49
-        stmt = sqlalchemy.select(sb_Owners.OwnerID).where(
-            sb_Owners.Owner == user.preferred_username,
-            sb_Owners.OwnerGroup == user.dirac_group,
-            sb_Owners.VO == user.vo,
+        stmt = select(SBOwners.OwnerID).where(
+            SBOwners.Owner == user.preferred_username,
+            SBOwners.OwnerGroup == user.dirac_group,
+            SBOwners.VO == user.vo,
         )
-        result = await self.conn.execute(stmt)
-        if owner_id := result.scalar_one_or_none():
-            return owner_id
+        return (await self.conn.execute(stmt)).scalar_one_or_none()
 
-        stmt = sqlalchemy.insert(sb_Owners).values(
+    async def get_sandbox_owner_id(self, pfn: str, se_name: str) -> int | None:
+        """Get the id of the owner of a sandbox."""
+        stmt = select(SBOwners.OwnerID).where(
+            SBOwners.OwnerID == SandBoxes.OwnerId,
+            SandBoxes.SEName == se_name,
+            SandBoxes.SEPFN == pfn,
+        )
+        return (await self.conn.execute(stmt)).scalar_one_or_none()
+
+    async def insert_owner(self, user: UserInfo) -> int:
+        stmt = insert(SBOwners).values(
             Owner=user.preferred_username,
             OwnerGroup=user.dirac_group,
             VO=user.vo,
@@ -48,12 +89,10 @@ class SandboxMetadataDB(BaseSQLDB):
         return "/" + "/".join(parts)
 
     async def insert_sandbox(
-        self, se_name: str, user: UserInfo, pfn: str, size: int
+        self, owner_id: int, se_name: str, pfn: str, size: int
     ) -> None:
         """Add a new sandbox in SandboxMetadataDB."""
-        # TODO: Follow https://github.com/DIRACGrid/diracx/issues/49
-        owner_id = await self.upsert_owner(user)
-        stmt = sqlalchemy.insert(sb_SandBoxes).values(
+        stmt = insert(SandBoxes).values(
             OwnerId=owner_id,
             SEName=se_name,
             SEPFN=pfn,
@@ -62,28 +101,36 @@ class SandboxMetadataDB(BaseSQLDB):
             LastAccessTime=utcnow(),
         )
         try:
-            result = await self.conn.execute(stmt)
-        except sqlalchemy.exc.IntegrityError:
-            await self.update_sandbox_last_access_time(se_name, pfn)
-        else:
-            assert result.rowcount == 1
+            await self.conn.execute(stmt)
+        except IntegrityError as e:
+            raise SandboxAlreadyInsertedError(pfn, se_name) from e
 
     async def update_sandbox_last_access_time(self, se_name: str, pfn: str) -> None:
         stmt = (
-            sqlalchemy.update(sb_SandBoxes)
-            .where(sb_SandBoxes.SEName == se_name, sb_SandBoxes.SEPFN == pfn)
+            update(SandBoxes)
+            .where(SandBoxes.SEName == se_name, SandBoxes.SEPFN == pfn)
             .values(LastAccessTime=utcnow())
         )
         result = await self.conn.execute(stmt)
-        assert result.rowcount == 1
+        if result.rowcount == 0:
+            # If the update didn't affect any row, the sandbox doesn't exist
+            raise SandboxNotFoundError(pfn, se_name)
+        elif result.rowcount != 1:
+            raise NotImplementedError(
+                "More than one sandbox was updated. This should not happen."
+            )
 
-    async def sandbox_is_assigned(self, pfn: str, se_name: str) -> bool:
+    async def sandbox_is_assigned(self, pfn: str, se_name: str) -> bool | None:
         """Checks if a sandbox exists and has been assigned."""
-        stmt: sqlalchemy.Executable = sqlalchemy.select(sb_SandBoxes.Assigned).where(
-            sb_SandBoxes.SEName == se_name, sb_SandBoxes.SEPFN == pfn
+        stmt: Executable = select(SandBoxes.Assigned).where(
+            SandBoxes.SEName == se_name, SandBoxes.SEPFN == pfn
         )
         result = await self.conn.execute(stmt)
-        is_assigned = result.scalar_one()
+        try:
+            is_assigned = result.scalar_one()
+        except NoResultFound as e:
+            raise SandboxNotFoundError(pfn, se_name) from e
+
         return is_assigned
 
     @staticmethod
@@ -97,11 +144,11 @@ class SandboxMetadataDB(BaseSQLDB):
         """Get the sandbox assign to job."""
         entity_id = self.jobid_to_entity_id(job_id)
         stmt = (
-            sqlalchemy.select(sb_SandBoxes.SEPFN)
-            .where(sb_SandBoxes.SBId == sb_EntityMapping.SBId)
+            select(SandBoxes.SEPFN)
+            .where(SandBoxes.SBId == SBEntityMapping.SBId)
             .where(
-                sb_EntityMapping.EntityId == entity_id,
-                sb_EntityMapping.Type == sb_type,
+                SBEntityMapping.EntityId == entity_id,
+                SBEntityMapping.Type == sb_type,
             )
         )
         result = await self.conn.execute(stmt)
@@ -114,58 +161,118 @@ class SandboxMetadataDB(BaseSQLDB):
         sb_type: SandboxType,
         se_name: str,
     ) -> None:
-        """Mapp sandbox and jobs."""
+        """Map sandbox and jobs."""
         for job_id in jobs_ids:
             # Define the entity id as 'Entity:entity_id' due to the DB definition:
             entity_id = self.jobid_to_entity_id(job_id)
-            select_sb_id = sqlalchemy.select(
-                sb_SandBoxes.SBId,
-                sqlalchemy.literal(entity_id).label("EntityId"),
-                sqlalchemy.literal(sb_type).label("Type"),
+            select_sb_id = select(
+                SandBoxes.SBId,
+                literal(entity_id).label("EntityId"),
+                literal(sb_type).label("Type"),
             ).where(
-                sb_SandBoxes.SEName == se_name,
-                sb_SandBoxes.SEPFN == pfn,
+                SandBoxes.SEName == se_name,
+                SandBoxes.SEPFN == pfn,
             )
-            stmt = sqlalchemy.insert(sb_EntityMapping).from_select(
+            stmt = insert(SBEntityMapping).from_select(
                 ["SBId", "EntityId", "Type"], select_sb_id
             )
-            await self.conn.execute(stmt)
+            try:
+                await self.conn.execute(stmt)
+            except IntegrityError as e:
+                raise SandboxAlreadyAssignedError(pfn, se_name) from e
 
-            stmt = (
-                sqlalchemy.update(sb_SandBoxes)
-                .where(sb_SandBoxes.SEPFN == pfn)
-                .values(Assigned=True)
-            )
+            stmt = update(SandBoxes).where(SandBoxes.SEPFN == pfn).values(Assigned=True)
             result = await self.conn.execute(stmt)
+            if result.rowcount == 0:
+                # If the update didn't affect any row, the sandbox doesn't exist
+                # It means the previous insert didn't have any effect
+                raise SandboxNotFoundError(pfn, se_name)
+
             assert result.rowcount == 1
 
     async def unassign_sandboxes_to_jobs(self, jobs_ids: list[int]) -> None:
         """Delete mapping between jobs and sandboxes."""
         for job_id in jobs_ids:
             entity_id = self.jobid_to_entity_id(job_id)
-            sb_sel_stmt = sqlalchemy.select(sb_SandBoxes.SBId)
+            sb_sel_stmt = select(SandBoxes.SBId)
             sb_sel_stmt = sb_sel_stmt.join(
-                sb_EntityMapping, sb_EntityMapping.SBId == sb_SandBoxes.SBId
+                SBEntityMapping, SBEntityMapping.SBId == SandBoxes.SBId
             )
-            sb_sel_stmt = sb_sel_stmt.where(sb_EntityMapping.EntityId == entity_id)
+            sb_sel_stmt = sb_sel_stmt.where(SBEntityMapping.EntityId == entity_id)
 
             result = await self.conn.execute(sb_sel_stmt)
             sb_ids = [row.SBId for row in result]
 
-            del_stmt = sqlalchemy.delete(sb_EntityMapping).where(
-                sb_EntityMapping.EntityId == entity_id
+            del_stmt = delete(SBEntityMapping).where(
+                SBEntityMapping.EntityId == entity_id
             )
             await self.conn.execute(del_stmt)
 
-            sb_entity_sel_stmt = sqlalchemy.select(sb_EntityMapping.SBId).where(
-                sb_EntityMapping.SBId.in_(sb_ids)
+            sb_entity_sel_stmt = select(SBEntityMapping.SBId).where(
+                SBEntityMapping.SBId.in_(sb_ids)
             )
             result = await self.conn.execute(sb_entity_sel_stmt)
             remaining_sb_ids = [row.SBId for row in result]
             if not remaining_sb_ids:
                 unassign_stmt = (
-                    sqlalchemy.update(sb_SandBoxes)
-                    .where(sb_SandBoxes.SBId.in_(sb_ids))
+                    update(SandBoxes)
+                    .where(SandBoxes.SBId.in_(sb_ids))
                     .values(Assigned=False)
                 )
                 await self.conn.execute(unassign_stmt)
+
+    @asynccontextmanager
+    async def delete_unused_sandboxes(
+        self, *, limit: int | None = None
+    ) -> AsyncGenerator[AsyncGenerator[str, None], None]:
+        """Get the sandbox PFNs to delete.
+
+        The result of this function can be used as an async context manager
+        to yield the PFNs of the sandboxes to delete. The context manager
+        will automatically remove the sandboxes from the database upon exit.
+
+        Args:
+            limit: If not None, the maximum number of sandboxes to delete.
+
+        """
+        conditions = [
+            # If it has assigned to a job but is no longer mapped it can be removed
+            and_(
+                SandBoxes.Assigned,
+                ~exists().where(SBEntityMapping.SBId == SandBoxes.SBId),
+            ),
+            # If the sandbox is still unassigned after 15 days, remove it
+            and_(~SandBoxes.Assigned, days_since(SandBoxes.LastAccessTime) >= 15),
+        ]
+        # Sandboxes which are not on S3 will be handled by legacy DIRAC
+        condition = and_(SandBoxes.SEPFN.like("/S3/%"), or_(*conditions))
+
+        # Copy the in-flight rows to a temporary table
+        await self.conn.run_sync(partial(self._temp_table.create, checkfirst=True))
+        select_stmt = select(SandBoxes.SBId).where(condition)
+        if limit:
+            select_stmt = select_stmt.limit(limit)
+        insert_stmt = insert(self._temp_table).from_select(["SBId"], select_stmt)
+        await self.conn.execute(insert_stmt)
+
+        try:
+            # Select the sandbox PFNs from the temporary table and yield them
+            select_stmt = select(SandBoxes.SEPFN).join(
+                self._temp_table, self._temp_table.c.SBId == SandBoxes.SBId
+            )
+
+            async def yield_pfns() -> AsyncGenerator[str, None]:
+                async for row in await self.conn.stream(select_stmt):
+                    yield row.SEPFN
+
+            yield yield_pfns()
+
+            # Delete the sandboxes from the main table
+            delete_stmt = delete(SandBoxes).where(
+                SandBoxes.SBId.in_(select(self._temp_table.c.SBId))
+            )
+            result = await self.conn.execute(delete_stmt)
+            logger.info("Deleted %d expired/unassigned sandboxes", result.rowcount)
+
+        finally:
+            await self.conn.run_sync(partial(self._temp_table.drop, checkfirst=True))

diracx/db/sql/sandbox_metadata/schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from sqlalchemy import (
     BigInteger,
     Boolean,
@@ -14,16 +16,19 @@ from diracx.db.sql.utils import Column, DateNowColumn
 Base = declarative_base()
 
 
-class sb_Owners(Base):
+class SBOwners(Base):
     __tablename__ = "sb_Owners"
     OwnerID = Column(Integer, autoincrement=True)
     Owner = Column(String(32))
     OwnerGroup = Column(String(32))
     VO = Column(String(64))
-    __table_args__ = (PrimaryKeyConstraint("OwnerID"),)
+    __table_args__ = (
+        PrimaryKeyConstraint("OwnerID"),
+        UniqueConstraint("Owner", "OwnerGroup", "VO", name="unique_owner_group_vo"),
+    )
 
 
-class sb_SandBoxes(Base):
+class SandBoxes(Base):
     __tablename__ = "sb_SandBoxes"
     SBId = Column(Integer, autoincrement=True)
     OwnerId = Column(Integer)
@@ -40,7 +45,7 @@ class sb_SandBoxes(Base):
     )
 
 
-class sb_EntityMapping(Base):
+class SBEntityMapping(Base):
     __tablename__ = "sb_EntityMapping"
     SBId = Column(Integer)
     EntityId = Column(String(128))

diracx/db/sql/task_queue/db.py

@@ -7,8 +7,6 @@ from sqlalchemy import delete, func, select, update
 if TYPE_CHECKING:
     pass
 
-from diracx.core.properties import JOB_SHARING, SecurityProperty
-
 from ..utils import BaseSQLDB
 from .schema import (
     BannedSitesQueue,
@@ -49,126 +47,23 @@ class TaskQueueDB(BaseSQLDB):
         )
         return dict((await self.conn.execute(stmt)).one()._mapping)
 
-    async def remove_job(self, job_id: int):
-        """Remove a job from the task queues."""
-        stmt = delete(JobsQueue).where(JobsQueue.JobId == job_id)
-        await self.conn.execute(stmt)
-
-    async def remove_jobs(self, job_ids: list[int]):
-        """Remove jobs from the task queues."""
-        stmt = delete(JobsQueue).where(JobsQueue.JobId.in_(job_ids))
-        await self.conn.execute(stmt)
-
-    async def delete_task_queue_if_empty(
-        self,
-        tq_id: int,
-        tq_owner: str,
-        tq_group: str,
-        job_share: int,
-        group_properties: set[SecurityProperty],
-        enable_shares_correction: bool,
-        allow_background_tqs: bool,
-    ):
-        """Try to delete a task queue if it's empty."""
-        # Check if the task queue is empty
-        stmt = (
-            select(TaskQueues.TQId)
-            .where(TaskQueues.Enabled >= 1)
-            .where(TaskQueues.TQId == tq_id)
-            .where(~TaskQueues.TQId.in_(select(JobsQueue.TQId)))
-        )
-        rows = await self.conn.execute(stmt)
-        if not rows.rowcount:
-            return
-
-        # Deleting the task queue (the other tables will be deleted in cascade)
-        stmt = delete(TaskQueues).where(TaskQueues.TQId == tq_id)
-        await self.conn.execute(stmt)
-
-        await self.recalculate_tq_shares_for_entity(
-            tq_owner,
-            tq_group,
-            job_share,
-            group_properties,
-            enable_shares_correction,
-            allow_background_tqs,
-        )
-
-    async def recalculate_tq_shares_for_entity(
-        self,
-        owner: str,
-        group: str,
-        job_share: int,
-        group_properties: set[SecurityProperty],
-        enable_shares_correction: bool,
-        allow_background_tqs: bool,
-    ):
-        """Recalculate the shares for a user/userGroup combo."""
-        if JOB_SHARING in group_properties:
-            # If group has JobSharing just set prio for that entry, user is irrelevant
-            return await self.__set_priorities_for_entity(
-                owner, group, job_share, group_properties, allow_background_tqs
-            )
-
+    async def get_task_queue_owners_by_group(self, group: str) -> dict[str, int]:
+        """Get the owners for a task queue and group."""
         stmt = (
             select(TaskQueues.Owner, func.count(TaskQueues.Owner))
             .where(TaskQueues.OwnerGroup == group)
             .group_by(TaskQueues.Owner)
         )
         rows = await self.conn.execute(stmt)
-        # make the rows a list of tuples
         # Get owners in this group and the amount of times they appear
-        # TODO: I guess the rows are already a list of tupes
+        # TODO: I guess the rows are already a list of tuples
         # maybe refactor
-        data = [(r[0], r[1]) for r in rows if r]
-        numOwners = len(data)
-        # If there are no owners do now
-        if numOwners == 0:
-            return
-        # Split the share amongst the number of owners
-        entities_shares = {row[0]: job_share / numOwners for row in data}
-
-        # TODO: implement the following
-        # If corrector is enabled let it work it's magic
-        # if enable_shares_correction:
-        #     entities_shares = await self.__shares_corrector.correct_shares(
-        #         entitiesShares, group=group
-        #     )
-
-        # Keep updating
-        owners = dict(data)
-        # IF the user is already known and has more than 1 tq, the rest of the users don't need to be modified
-        # (The number of owners didn't change)
-        if owner in owners and owners[owner] > 1:
-            await self.__set_priorities_for_entity(
-                owner,
-                group,
-                entities_shares[owner],
-                group_properties,
-                allow_background_tqs,
-            )
-            return
-        # Oops the number of owners may have changed so we recalculate the prio for all owners in the group
-        for owner in owners:
-            await self.__set_priorities_for_entity(
-                owner,
-                group,
-                entities_shares[owner],
-                group_properties,
-                allow_background_tqs,
-            )
-
-    async def __set_priorities_for_entity(
-        self,
-        owner: str,
-        group: str,
-        share,
-        properties: set[SecurityProperty],
-        allow_background_tqs: bool,
-    ):
-        """Set the priority for a user/userGroup combo given a splitted share."""
-        from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import calculate_priority
+        return {r[0]: r[1] for r in rows if r}
 
+    async def get_task_queue_priorities(
+        self, group: str, owner: str | None = None
+    ) -> dict[int, float]:
+        """Get the priorities for a list of task queues."""
         stmt = (
             select(
                 TaskQueues.TQId,
@@ -178,24 +73,48 @@ class TaskQueueDB(BaseSQLDB):
             .where(TaskQueues.OwnerGroup == group)
             .group_by(TaskQueues.TQId)
         )
-        if JOB_SHARING not in properties:
+        if owner:
             stmt = stmt.where(TaskQueues.Owner == owner)
         rows = await self.conn.execute(stmt)
-        tq_dict: dict[int, float] = {tq_id: priority for tq_id, priority in rows}
+        return {tq_id: priority for tq_id, priority in rows}
 
-        if not tq_dict:
-            return
+    async def remove_jobs(self, job_ids: list[int]):
+        """Remove jobs from the task queues."""
+        stmt = delete(JobsQueue).where(JobsQueue.JobId.in_(job_ids))
+        await self.conn.execute(stmt)
 
-        rows = await self.retrieve_task_queues(list(tq_dict))
+    async def is_task_queue_empty(self, tq_id: int) -> bool:
+        """Check if a task queue is empty."""
+        stmt = (
+            select(TaskQueues.TQId)
+            .where(TaskQueues.Enabled >= 1)
+            .where(TaskQueues.TQId == tq_id)
+            .where(~TaskQueues.TQId.in_(select(JobsQueue.TQId)))
+        )
+        rows = await self.conn.execute(stmt)
+        return not rows.rowcount
 
-        prio_dict = calculate_priority(tq_dict, rows, share, allow_background_tqs)
+    async def delete_task_queue(
+        self,
+        tq_id: int,
+    ):
+        """Delete a task queue."""
+        # Deleting the task queue (the other tables will be deleted in cascade)
+        stmt = delete(TaskQueues).where(TaskQueues.TQId == tq_id)
+        await self.conn.execute(stmt)
 
-        # Execute updates
-        for prio, tqs in prio_dict.items():
-            update_stmt = (
-                update(TaskQueues).where(TaskQueues.TQId.in_(tqs)).values(Priority=prio)
-            )
-            await self.conn.execute(update_stmt)
+    async def set_priorities_for_entity(
+        self,
+        tq_ids: list[int],
+        priority: float,
+    ):
+        """Set the priority for a user/userGroup combo given a split share."""
+        update_stmt = (
+            update(TaskQueues)
+            .where(TaskQueues.TQId.in_(tq_ids))
+            .values(Priority=priority)
+        )
+        await self.conn.execute(update_stmt)
 
     async def retrieve_task_queues(self, tq_id_list=None):
         """Get all the task queues."""

diracx/db/sql/task_queue/schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from sqlalchemy import (
     BigInteger,
     Boolean,