diracx-db 0.0.1a44__py3-none-any.whl → 0.0.1a45__py3-none-any.whl

diracx/db/os/job_parameters.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from DIRAC.Core.Utilities import TimeUtilities
+from datetime import UTC, datetime
 
 from diracx.db.os.utils import BaseOSDB
 
@@ -31,7 +31,7 @@ class JobParametersDB(BaseOSDB):
     def upsert(self, vo, doc_id, document):
         document = {
             "JobID": doc_id,
-            "timestamp": TimeUtilities.toEpochMilliSeconds(),
+            "timestamp": int(datetime.now(tz=UTC).timestamp() * 1000),
             **document,
         }
         return super().upsert(vo, doc_id, document)

diracx/db/sql/job/db.py

@@ -50,7 +50,7 @@ class JobDB(BaseSQLDB):
     }
 
     # TODO: this is copied from the DIRAC JobDB
-    # but is overwriten in LHCbDIRAC, so we need
+    # but is overwritten in LHCbDIRAC, so we need
     # to find a way to make it dynamic
     jdl_2_db_parameters = ["JobName", "JobType", "JobGroup"]
 

diracx/db/sql/job_logging/schema.py

@@ -44,7 +44,7 @@ class MagicEpochDateTime(TypeDecorator):
         """
         if value is None:
             return None
-        # Carefully convert from Decimal to datetime to avoid loosing precision
+        # Carefully convert from Decimal to datetime to avoid losing precision
         value += self.MAGIC_EPOC_NUMBER
         value_int = int(value)
         result = datetime.fromtimestamp(value_int, tz=UTC)

diracx/db/sql/sandbox_metadata/db.py

@@ -1,8 +1,24 @@
 from __future__ import annotations
 
-from typing import Any
+import logging
+from contextlib import asynccontextmanager
+from functools import partial
+from typing import Any, AsyncGenerator
 
-from sqlalchemy import Executable, delete, insert, literal, select, update
+from sqlalchemy import (
+    BigInteger,
+    Column,
+    Executable,
+    MetaData,
+    Table,
+    and_,
+    delete,
+    insert,
+    literal,
+    or_,
+    select,
+    update,
+)
 from sqlalchemy.exc import IntegrityError, NoResultFound
 
 from diracx.core.exceptions import (
@@ -12,15 +28,25 @@ from diracx.core.exceptions import (
 )
 from diracx.core.models import SandboxInfo, SandboxType, UserInfo
 from diracx.db.sql.utils.base import BaseSQLDB
-from diracx.db.sql.utils.functions import utcnow
+from diracx.db.sql.utils.functions import days_since, utcnow
 
 from .schema import Base as SandboxMetadataDBBase
 from .schema import SandBoxes, SBEntityMapping, SBOwners
 
+logger = logging.getLogger(__name__)
+
 
 class SandboxMetadataDB(BaseSQLDB):
     metadata = SandboxMetadataDBBase.metadata
 
+    # 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."""
         stmt = select(SBOwners.OwnerID).where(
@@ -85,7 +111,13 @@ class SandboxMetadataDB(BaseSQLDB):
             .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 | None:
         """Checks if a sandbox exists and has been assigned."""
@@ -128,7 +160,7 @@ 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)
@@ -187,3 +219,56 @@ class SandboxMetadataDB(BaseSQLDB):
                     .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(SandBoxes.SBId == SBEntityMapping.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/task_queue/db.py

@@ -56,7 +56,7 @@ class TaskQueueDB(BaseSQLDB):
         )
         rows = await self.conn.execute(stmt)
         # 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
         return {r[0]: r[1] for r in rows if r}
 
@@ -108,7 +108,7 @@ class TaskQueueDB(BaseSQLDB):
         tq_ids: list[int],
         priority: float,
     ):
-        """Set the priority for a user/userGroup combo given a splitted share."""
+        """Set the priority for a user/userGroup combo given a split share."""
         update_stmt = (
             update(TaskQueues)
             .where(TaskQueues.TQId.in_(tq_ids))

diracx/db/sql/utils/base.py

@@ -51,7 +51,7 @@ class BaseSQLDB(metaclass=ABCMeta):
     of the same database. To list the available implementations use
     `BaseSQLDB.available_implementations(db_name)`. The first entry in this list
     will be the preferred implementation and it can be initialized by calling
-    it's `__init__` function with a URL perviously obtained from
+    it's `__init__` function with a URL previously obtained from
     `BaseSQLDB.available_urls`.
 
     To control the lifetime of the SQLAlchemy engine used for connecting to the

diracx/db/sql/utils/functions.py

@@ -103,6 +103,37 @@ def sqlite_date_trunc(element, compiler, **kw):
     )
 
 
+class days_since(expression.FunctionElement):  # noqa: N801
+    """Sqlalchemy function to get the number of days since a given date.
+
+    Primarily used to be able to query for a specific resolution of a date e.g.
+
+        select * from table where days_since(date_column) = 0
+        select * from table where days_since(date_column) = 1
+    """
+
+    type = DateTime()
+    inherit_cache = False
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+
+@compiles(days_since, "postgresql")
+def pg_days_since(element, compiler, **kw):
+    return f"EXTRACT(DAY FROM (now() - {compiler.process(element.clauses)}))"
+
+
+@compiles(days_since, "mysql")
+def mysql_days_since(element, compiler, **kw):
+    return f"DATEDIFF(NOW(), {compiler.process(element.clauses)})"
+
+
+@compiles(days_since, "sqlite")
+def sqlite_days_since(element, compiler, **kw):
+    return f"julianday('now') - julianday({compiler.process(element.clauses)})"
+
+
 def substract_date(**kwargs: float) -> datetime:
     return datetime.now(tz=timezone.utc) - timedelta(**kwargs)
 

{diracx_db-0.0.1a44.dist-info → diracx_db-0.0.1a45.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diracx-db
-Version: 0.0.1a44
+Version: 0.0.1a45
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research
@@ -9,11 +9,11 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: System :: Distributed Computing
 Requires-Python: >=3.11
-Description-Content-Type: text/markdown
 Requires-Dist: diracx-core
 Requires-Dist: opensearch-py[async]
 Requires-Dist: pydantic>=2.10
 Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
 Requires-Dist: uuid-utils
 Provides-Extra: testing
-Requires-Dist: diracx-testing; extra == "testing"
+Requires-Dist: diracx-testing; extra == 'testing'
+Requires-Dist: freezegun; extra == 'testing'

{diracx_db-0.0.1a44.dist-info → diracx_db-0.0.1a45.dist-info}/RECORD

@@ -3,7 +3,7 @@ diracx/db/__main__.py,sha256=tU4tp3OAClYCiPMxlRj524sZGBx9oy4CoWHd8pMuEEs,1715
 diracx/db/exceptions.py,sha256=1nn-SZLG-nQwkxbvHjZqXhE5ouzWj1f3qhSda2B4ZEg,83
 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=YKW5Mx0_TB6cPJoLi08SJWckjXsedRDLq9QbVSQb5CM,1228
+diracx/db/os/job_parameters.py,sha256=loAc-bo3u-RMAp_H1g8VRt8T-rCCsXp_d9aCvg5OS-A,1225
 diracx/db/os/utils.py,sha256=V4T-taos64SFNcorfIr7mq5l5y88K6TzyCj1YqWk8VI,11562
 diracx/db/sql/__init__.py,sha256=JYu0b0IVhoXy3lX2m2r2dmAjsRS7IbECBUMEDvX0Te4,391
 diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -13,26 +13,25 @@ diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 diracx/db/sql/dummy/db.py,sha256=IW4FzG7ERKbhZvC32KL7Rodu2u-zKAf8BryO4VAdJew,1650
 diracx/db/sql/dummy/schema.py,sha256=9zI53pKlzc6qBezsyjkatOQrNZdGCjwgjQ8Iz_pyAXs,789
 diracx/db/sql/job/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/job/db.py,sha256=PTivVrSl4hrOTgy77WECunPJUJiYzGhSlqO0wYwEuE8,11909
+diracx/db/sql/job/db.py,sha256=TnEc0fckiuMJAZg2v1_Pbwfn7kDPDam6TXp9ySuiddk,11910
 diracx/db/sql/job/schema.py,sha256=eFgZshe6NEzOM2qI0HI9Y3abrqDMoQIwa9L0vZugHcU,5431
 diracx/db/sql/job_logging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/job_logging/db.py,sha256=BYzlPuvdvHR7wdzQEVWMH_V5kL0bLBZtQkcugnSGbjs,5497
-diracx/db/sql/job_logging/schema.py,sha256=W2VeE7czBXF6uP2hmjBDur90S4BlidzFXX2V_WZkfzU,2525
+diracx/db/sql/job_logging/schema.py,sha256=k6uBw-RHAcJ5GEleNpiWoXEJBhCiNG-y4xAgBKHZjjM,2524
 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=KeWnFSpYOTrT3-_rOCFjbjNnPNXKnUZiJVsu4vv5U2U,2149
 diracx/db/sql/sandbox_metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/sandbox_metadata/db.py,sha256=1Nd0gNbPu2F21DYGQkFCLhTeKgUA26GE_vxDcdzUzwQ,7052
+diracx/db/sql/sandbox_metadata/db.py,sha256=DAmuk-PDGKq3eLV2EovDGnJI0GDiH5u8C74ARIy8MWo,10171
 diracx/db/sql/sandbox_metadata/schema.py,sha256=V5gV2PHwzTbBz_th9ribLfE7Lqk8YGemDmvqq4jWQJ4,1530
 diracx/db/sql/task_queue/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/task_queue/db.py,sha256=Ymd9cBasiSAPmgu5846B8KmguRiJ_GoJhuOev9YwW2o,6248
+diracx/db/sql/task_queue/db.py,sha256=2qul1D2tX2uCI92N591WK5xWHakG0pNibzDwKQ7W-I8,6246
 diracx/db/sql/task_queue/schema.py,sha256=5efAgvNYRkLlaJ2NzRInRfmVa3tyIzQu2l0oRPy4Kzw,3258
 diracx/db/sql/utils/__init__.py,sha256=QkvpqBuIAgkAOywAssYzdxSzUQVZlSUumK7mPxotXfM,547
-diracx/db/sql/utils/base.py,sha256=7UxHBNLOSjdrIdslMKW4C_c5H9-6Y1BEimxscri2poE,12367
-diracx/db/sql/utils/functions.py,sha256=iLqlUIQ6SrDUtDEnZ5szaFbdcINJW15KNbCdGXss6kc,3055
+diracx/db/sql/utils/base.py,sha256=HYQuX16mgg9LAMtAEmbTmJFIN0OSMe1Hcb57dtl7LCc,12367
+diracx/db/sql/utils/functions.py,sha256=_E4tc9Gti6LuSh7QEyoqPJSvCuByVqvRenOXCzxsulE,4014
 diracx/db/sql/utils/types.py,sha256=yU-tXsu6hFGPsr9ba1n3ZjGPnHQI_06lbpkTeDCWJtg,1287
-diracx_db-0.0.1a44.dist-info/METADATA,sha256=GTQNVgaV7dVycm3qPuHFl32-EkLmeUr4ahtCgPyNbfQ,670
-diracx_db-0.0.1a44.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-diracx_db-0.0.1a44.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
-diracx_db-0.0.1a44.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
-diracx_db-0.0.1a44.dist-info/RECORD,,
+diracx_db-0.0.1a45.dist-info/METADATA,sha256=ET2Uo-DegfUc6qRqvicSGR7PpWTtc4yXJKqoZQRbXyQ,675
+diracx_db-0.0.1a45.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+diracx_db-0.0.1a45.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
+diracx_db-0.0.1a45.dist-info/RECORD,,

{diracx_db-0.0.1a44.dist-info → diracx_db-0.0.1a45.dist-info}/WHEEL

@@ -1,5 +1,4 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.0.0)
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any
-

diracx_db-0.0.1a44.dist-info/top_level.txt

@@ -1 +0,0 @@
-diracx