diracx-db 0.0.1a11__tar.gz → 0.0.1a13__tar.gz

{diracx-db-0.0.1a11 → diracx-db-0.0.1a13}/PKG-INFO

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

{diracx-db-0.0.1a11 → diracx-db-0.0.1a13}/src/diracx/db/sql/auth/db.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import hashlib
 import secrets
 from datetime import datetime
 from uuid import uuid4
@@ -63,7 +64,7 @@ class AuthDB(BaseSQLDB):
             ),
         ).with_for_update()
         stmt = stmt.where(
-            DeviceFlows.device_code == device_code,
+            DeviceFlows.device_code == hashlib.sha256(device_code.encode()).hexdigest(),
         )
         res = dict((await self.conn.execute(stmt)).one()._mapping)
 
@@ -74,7 +75,10 @@ class AuthDB(BaseSQLDB):
             # Update the status to Done before returning
             await self.conn.execute(
                 update(DeviceFlows)
-                .where(DeviceFlows.device_code == device_code)
+                .where(
+                    DeviceFlows.device_code
+                    == hashlib.sha256(device_code.encode()).hexdigest()
+                )
                 .values(status=FlowStatus.DONE)
             )
             return res
@@ -110,7 +114,6 @@ class AuthDB(BaseSQLDB):
         self,
         client_id: str,
         scope: str,
-        audience: str,
     ) -> tuple[str, str]:
         # Because the user_code might be short, there is a risk of conflicts
         # This is why we retry multiple times
@@ -119,14 +122,16 @@ class AuthDB(BaseSQLDB):
                 secrets.choice(USER_CODE_ALPHABET)
                 for _ in range(DeviceFlows.user_code.type.length)  # type: ignore
             )
-            # user_code = "2QRKPY"
             device_code = secrets.token_urlsafe()
+
+            # Hash the the device_code to avoid leaking information
+            hashed_device_code = hashlib.sha256(device_code.encode()).hexdigest()
+
             stmt = insert(DeviceFlows).values(
                 client_id=client_id,
                 scope=scope,
-                audience=audience,
                 user_code=user_code,
-                device_code=device_code,
+                device_code=hashed_device_code,
             )
             try:
                 await self.conn.execute(stmt)
@@ -143,7 +148,6 @@ class AuthDB(BaseSQLDB):
         self,
         client_id: str,
         scope: str,
-        audience: str,
         code_challenge: str,
         code_challenge_method: str,
         redirect_uri: str,
@@ -154,7 +158,6 @@ class AuthDB(BaseSQLDB):
             uuid=uuid,
             client_id=client_id,
             scope=scope,
-            audience=audience,
             code_challenge=code_challenge,
             code_challenge_method=code_challenge_method,
             redirect_uri=redirect_uri,
@@ -172,7 +175,10 @@ class AuthDB(BaseSQLDB):
         :raises: AuthorizationError if no such uuid or status not pending
         """
 
+        # Hash the code to avoid leaking information
         code = secrets.token_urlsafe()
+        hashed_code = hashlib.sha256(code.encode()).hexdigest()
+
         stmt = update(AuthorizationFlows)
 
         stmt = stmt.where(
@@ -181,7 +187,7 @@ class AuthDB(BaseSQLDB):
             AuthorizationFlows.creation_time > substract_date(seconds=max_validity),
         )
 
-        stmt = stmt.values(id_token=id_token, code=code, status=FlowStatus.READY)
+        stmt = stmt.values(id_token=id_token, code=hashed_code, status=FlowStatus.READY)
         res = await self.conn.execute(stmt)
 
         if res.rowcount != 1:
@@ -190,15 +196,16 @@ class AuthDB(BaseSQLDB):
         stmt = select(AuthorizationFlows.code, AuthorizationFlows.redirect_uri)
         stmt = stmt.where(AuthorizationFlows.uuid == uuid)
         row = (await self.conn.execute(stmt)).one()
-        return row.code, row.redirect_uri
+        return code, row.redirect_uri
 
     async def get_authorization_flow(self, code: str, max_validity: int):
+        hashed_code = hashlib.sha256(code.encode()).hexdigest()
         # The with_for_update
         # prevents that the token is retrieved
         # multiple time concurrently
         stmt = select(AuthorizationFlows).with_for_update()
         stmt = stmt.where(
-            AuthorizationFlows.code == code,
+            AuthorizationFlows.code == hashed_code,
             AuthorizationFlows.creation_time > substract_date(seconds=max_validity),
         )
 
@@ -208,7 +215,7 @@ class AuthDB(BaseSQLDB):
             # Update the status to Done before returning
             await self.conn.execute(
                 update(AuthorizationFlows)
-                .where(AuthorizationFlows.code == code)
+                .where(AuthorizationFlows.code == hashed_code)
                 .values(status=FlowStatus.DONE)
             )
 

{diracx-db-0.0.1a11 → diracx-db-0.0.1a13}/src/diracx/db/sql/auth/schema.py

@@ -45,8 +45,7 @@ class DeviceFlows(Base):
     creation_time = DateNowColumn()
     client_id = Column(String(255))
     scope = Column(String(1024))
-    audience = Column(String(255))
-    device_code = Column(String(128), unique=True)  # hash it ?
+    device_code = Column(String(128), unique=True)  # Should be a hash
     id_token = NullColumn(JSON())
 
 
@@ -57,11 +56,10 @@ class AuthorizationFlows(Base):
     client_id = Column(String(255))
     creation_time = DateNowColumn()
     scope = Column(String(1024))
-    audience = Column(String(255))
     code_challenge = Column(String(255))
     code_challenge_method = Column(String(8))
     redirect_uri = Column(String(255))
-    code = NullColumn(String(255))  # hash it ?
+    code = NullColumn(String(255))  # Should be a hash
     id_token = NullColumn(JSON())
 
 

{diracx-db-0.0.1a11 → diracx-db-0.0.1a13}/src/diracx/db/sql/jobs/status_utility.py

@@ -272,7 +272,7 @@ async def remove_jobs(
 
     # TODO: this was also  not done in the JobManagerHandler, but it was done in the JobCleaningAgent
     # I think it should be done here as well
-    await sandbox_metadata_db.unassign_sandbox_from_jobs(job_ids)
+    await sandbox_metadata_db.unassign_sandboxes_to_jobs(job_ids)
 
     # Remove the job from TaskQueueDB
     await _remove_jobs_from_task_queue(job_ids, config, task_queue_db, background_task)

diracx-db-0.0.1a13/src/diracx/db/sql/sandbox_metadata/db.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+from typing import Any
+
+import sqlalchemy
+
+from diracx.core.models import SandboxInfo, SandboxType, UserInfo
+from diracx.db.sql.utils import BaseSQLDB, utcnow
+
+from .schema import Base as SandboxMetadataDBBase
+from .schema import sb_EntityMapping, sb_Owners, sb_SandBoxes
+
+
+class SandboxMetadataDB(BaseSQLDB):
+    metadata = SandboxMetadataDBBase.metadata
+
+    async def upsert_owner(self, user: UserInfo) -> int:
+        """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,
+        )
+        result = await self.conn.execute(stmt)
+        if owner_id := result.scalar_one_or_none():
+            return owner_id
+
+        stmt = sqlalchemy.insert(sb_Owners).values(
+            Owner=user.preferred_username,
+            OwnerGroup=user.dirac_group,
+            VO=user.vo,
+        )
+        result = await self.conn.execute(stmt)
+        return result.lastrowid
+
+    @staticmethod
+    def get_pfn(bucket_name: str, user: UserInfo, sandbox_info: SandboxInfo) -> str:
+        """Get the sandbox's user namespaced and content addressed PFN"""
+        parts = [
+            "S3",
+            bucket_name,
+            user.vo,
+            user.dirac_group,
+            user.preferred_username,
+            f"{sandbox_info.checksum_algorithm}:{sandbox_info.checksum}.{sandbox_info.format}",
+        ]
+        return "/" + "/".join(parts)
+
+    async def insert_sandbox(
+        self, se_name: str, user: UserInfo, 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(
+            OwnerId=owner_id,
+            SEName=se_name,
+            SEPFN=pfn,
+            Bytes=size,
+            RegistrationTime=utcnow(),
+            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
+
+    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)
+            .values(LastAccessTime=utcnow())
+        )
+        result = await self.conn.execute(stmt)
+        assert result.rowcount == 1
+
+    async def sandbox_is_assigned(self, pfn: str, se_name: str) -> bool:
+        """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
+        )
+        result = await self.conn.execute(stmt)
+        is_assigned = result.scalar_one()
+        return is_assigned
+
+    @staticmethod
+    def jobid_to_entity_id(job_id: int) -> str:
+        """Define the entity id as 'Entity:entity_id' due to the DB definition"""
+        return f"Job:{job_id}"
+
+    async def get_sandbox_assigned_to_job(
+        self, job_id: int, sb_type: SandboxType
+    ) -> list[Any]:
+        """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)
+            .where(
+                sb_EntityMapping.EntityId == entity_id,
+                sb_EntityMapping.Type == sb_type,
+            )
+        )
+        result = await self.conn.execute(stmt)
+        return [result.scalar()]
+
+    async def assign_sandbox_to_jobs(
+        self,
+        jobs_ids: list[int],
+        pfn: str,
+        sb_type: SandboxType,
+        se_name: str,
+    ) -> None:
+        """Mapp 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"),
+            ).where(
+                sb_SandBoxes.SEName == se_name,
+                sb_SandBoxes.SEPFN == pfn,
+            )
+            stmt = sqlalchemy.insert(sb_EntityMapping).from_select(
+                ["SBId", "EntityId", "Type"], select_sb_id
+            )
+            await self.conn.execute(stmt)
+
+            stmt = (
+                sqlalchemy.update(sb_SandBoxes)
+                .where(sb_SandBoxes.SEPFN == pfn)
+                .values(Assigned=True)
+            )
+            result = await self.conn.execute(stmt)
+            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,
+            ).where(sb_EntityMapping.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
+            )
+            await self.conn.execute(del_stmt)
+
+            sb_entity_sel_stmt = sqlalchemy.select(sb_EntityMapping.SBId).where(
+                sb_EntityMapping.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))
+                    .values(Assigned=False)
+                )
+                await self.conn.execute(unassign_stmt)

{diracx-db-0.0.1a11 → diracx-db-0.0.1a13}/src/diracx_db.egg-info/PKG-INFO

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

{diracx-db-0.0.1a11 → diracx-db-0.0.1a13}/tests/auth/test_authorization_flow.py

@@ -23,7 +23,7 @@ async def test_insert_id_token(auth_db: AuthDB):
     # First insert
     async with auth_db as auth_db:
         uuid = await auth_db.insert_authorization_flow(
-            "client_id", "scope", "audience", "code_challenge", "S256", "redirect_uri"
+            "client_id", "scope", "code_challenge", "S256", "redirect_uri"
         )
 
     id_token = {"sub": "myIdToken"}
@@ -68,12 +68,11 @@ async def test_insert(auth_db: AuthDB):
     # First insert
     async with auth_db as auth_db:
         uuid1 = await auth_db.insert_authorization_flow(
-            "client_id", "scope", "audience", "code_challenge", "S256", "redirect_uri"
+            "client_id", "scope", "code_challenge", "S256", "redirect_uri"
         )
         uuid2 = await auth_db.insert_authorization_flow(
             "client_id2",
             "scope2",
-            "audience2",
             "code_challenge2",
             "S256",
             "redirect_uri2",

{diracx-db-0.0.1a11 → diracx-db-0.0.1a13}/tests/auth/test_device_flow.py

@@ -28,20 +28,22 @@ async def test_device_user_code_collision(auth_db: AuthDB, monkeypatch):
     # First insert should work
     async with auth_db as auth_db:
         code, device = await auth_db.insert_device_flow(
-            "client_id", "scope", "audience"
+            "client_id",
+            "scope",
         )
         assert code == "A" * USER_CODE_LENGTH
         assert device
 
     async with auth_db as auth_db:
         with pytest.raises(NotImplementedError, match="insert new device flow"):
-            await auth_db.insert_device_flow("client_id", "scope", "audience")
+            await auth_db.insert_device_flow("client_id", "scope")
 
     monkeypatch.setattr(secrets, "choice", lambda _: "B")
 
     async with auth_db as auth_db:
         code, device = await auth_db.insert_device_flow(
-            "client_id", "scope", "audience"
+            "client_id",
+            "scope",
         )
         assert code == "B" * USER_CODE_LENGTH
         assert device
@@ -59,10 +61,12 @@ async def test_device_flow_lookup(auth_db: AuthDB, monkeypatch):
     # First insert
     async with auth_db as auth_db:
         user_code1, device_code1 = await auth_db.insert_device_flow(
-            "client_id1", "scope1", "audience1"
+            "client_id1",
+            "scope1",
         )
         user_code2, device_code2 = await auth_db.insert_device_flow(
-            "client_id2", "scope2", "audience2"
+            "client_id2",
+            "scope2",
         )
 
         assert user_code1 != user_code2
@@ -123,7 +127,8 @@ async def test_device_flow_insert_id_token(auth_db: AuthDB):
     # First insert
     async with auth_db as auth_db:
         user_code, device_code = await auth_db.insert_device_flow(
-            "client_id", "scope", "audience"
+            "client_id",
+            "scope",
         )
 
     # Make sure it exists, and is Pending

diracx-db-0.0.1a13/tests/test_sandbox_metadata.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import asyncio
+import secrets
+from datetime import datetime
+
+import pytest
+import sqlalchemy
+
+from diracx.core.models import SandboxInfo, UserInfo
+from diracx.db.sql.sandbox_metadata.db import SandboxMetadataDB
+from diracx.db.sql.sandbox_metadata.schema import sb_EntityMapping, sb_SandBoxes
+
+
+@pytest.fixture
+async def sandbox_metadata_db(tmp_path):
+    sandbox_metadata_db = SandboxMetadataDB("sqlite+aiosqlite:///:memory:")
+    async with sandbox_metadata_db.engine_context():
+        async with sandbox_metadata_db.engine.begin() as conn:
+            await conn.run_sync(sandbox_metadata_db.metadata.create_all)
+        yield sandbox_metadata_db
+
+
+def test_get_pfn(sandbox_metadata_db: SandboxMetadataDB):
+    user_info = UserInfo(
+        sub="vo:sub", preferred_username="user1", dirac_group="group1", vo="vo"
+    )
+    sandbox_info = SandboxInfo(
+        checksum="checksum",
+        checksum_algorithm="sha256",
+        format="tar.bz2",
+        size=100,
+    )
+    pfn = sandbox_metadata_db.get_pfn("bucket1", user_info, sandbox_info)
+    assert pfn == "/S3/bucket1/vo/group1/user1/sha256:checksum.tar.bz2"
+
+
+async def test_insert_sandbox(sandbox_metadata_db: SandboxMetadataDB):
+    user_info = UserInfo(
+        sub="vo:sub", preferred_username="user1", dirac_group="group1", vo="vo"
+    )
+    pfn1 = secrets.token_hex()
+
+    # Make sure the sandbox doesn't already exist
+    db_contents = await _dump_db(sandbox_metadata_db)
+    assert pfn1 not in db_contents
+    async with sandbox_metadata_db:
+        with pytest.raises(sqlalchemy.exc.NoResultFound):
+            await sandbox_metadata_db.sandbox_is_assigned(pfn1, "SandboxSE")
+
+    # Insert the sandbox
+    async with sandbox_metadata_db:
+        await sandbox_metadata_db.insert_sandbox("SandboxSE", user_info, pfn1, 100)
+    db_contents = await _dump_db(sandbox_metadata_db)
+    owner_id1, last_access_time1 = db_contents[pfn1]
+
+    # Inserting again should update the last access time
+    await asyncio.sleep(1)  # The timestamp only has second precision
+    async with sandbox_metadata_db:
+        await sandbox_metadata_db.insert_sandbox("SandboxSE", user_info, pfn1, 100)
+    db_contents = await _dump_db(sandbox_metadata_db)
+    owner_id2, last_access_time2 = db_contents[pfn1]
+    assert owner_id1 == owner_id2
+    assert last_access_time2 > last_access_time1
+
+    # The sandbox still hasn't been assigned
+    async with sandbox_metadata_db:
+        assert not await sandbox_metadata_db.sandbox_is_assigned(pfn1, "SandboxSE")
+
+    # Inserting again should update the last access time
+    await asyncio.sleep(1)  # The timestamp only has second precision
+    last_access_time3 = (await _dump_db(sandbox_metadata_db))[pfn1][1]
+    assert last_access_time2 == last_access_time3
+    async with sandbox_metadata_db:
+        await sandbox_metadata_db.update_sandbox_last_access_time("SandboxSE", pfn1)
+    last_access_time4 = (await _dump_db(sandbox_metadata_db))[pfn1][1]
+    assert last_access_time2 < last_access_time4
+
+
+async def _dump_db(
+    sandbox_metadata_db: SandboxMetadataDB,
+) -> dict[str, tuple[int, datetime]]:
+    """Dump the contents of the sandbox metadata database
+
+    Returns a dict[pfn: str, (owner_id: int, last_access_time: datetime)]
+    """
+    async with sandbox_metadata_db:
+        stmt = sqlalchemy.select(
+            sb_SandBoxes.SEPFN, sb_SandBoxes.OwnerId, sb_SandBoxes.LastAccessTime
+        )
+        res = await sandbox_metadata_db.conn.execute(stmt)
+        return {row.SEPFN: (row.OwnerId, row.LastAccessTime) for row in res}
+
+
+async def test_assign_and_unsassign_sandbox_to_jobs(
+    sandbox_metadata_db: SandboxMetadataDB,
+):
+    pfn = secrets.token_hex()
+    user_info = UserInfo(
+        sub="vo:sub", preferred_username="user1", dirac_group="group1", vo="vo"
+    )
+    dummy_jobid = 666
+    sandbox_se = "SandboxSE"
+    # Insert the sandbox
+    async with sandbox_metadata_db:
+        await sandbox_metadata_db.insert_sandbox(sandbox_se, user_info, pfn, 100)
+
+    async with sandbox_metadata_db:
+        stmt = sqlalchemy.select(sb_SandBoxes.SBId, sb_SandBoxes.SEPFN)
+        res = await sandbox_metadata_db.conn.execute(stmt)
+    db_contents = {row.SEPFN: row.SBId for row in res}
+    sb_id_1 = db_contents[pfn]
+    # The sandbox still hasn't been assigned
+    async with sandbox_metadata_db:
+        assert not await sandbox_metadata_db.sandbox_is_assigned(pfn, sandbox_se)
+
+    # Check there is no mapping
+    async with sandbox_metadata_db:
+        stmt = sqlalchemy.select(
+            sb_EntityMapping.SBId, sb_EntityMapping.EntityId, sb_EntityMapping.Type
+        )
+        res = await sandbox_metadata_db.conn.execute(stmt)
+    db_contents = {row.SBId: (row.EntityId, row.Type) for row in res}
+    assert db_contents == {}
+
+    # Assign sandbox with dummy jobid
+    async with sandbox_metadata_db:
+        await sandbox_metadata_db.assign_sandbox_to_jobs(
+            jobs_ids=[dummy_jobid], pfn=pfn, sb_type="Output", se_name=sandbox_se
+        )
+    # Check if sandbox and job are mapped
+    async with sandbox_metadata_db:
+        stmt = sqlalchemy.select(
+            sb_EntityMapping.SBId, sb_EntityMapping.EntityId, sb_EntityMapping.Type
+        )
+        res = await sandbox_metadata_db.conn.execute(stmt)
+    db_contents = {row.SBId: (row.EntityId, row.Type) for row in res}
+
+    entity_id_1, sb_type = db_contents[sb_id_1]
+    assert entity_id_1 == f"Job:{dummy_jobid}"
+    assert sb_type == "Output"
+
+    async with sandbox_metadata_db:
+        stmt = sqlalchemy.select(sb_SandBoxes.SBId, sb_SandBoxes.SEPFN)
+        res = await sandbox_metadata_db.conn.execute(stmt)
+    db_contents = {row.SEPFN: row.SBId for row in res}
+    sb_id_1 = db_contents[pfn]
+    # The sandbox should be assigned
+    async with sandbox_metadata_db:
+        assert await sandbox_metadata_db.sandbox_is_assigned(pfn, sandbox_se)
+
+    # Unassign the sandbox to job
+    async with sandbox_metadata_db:
+        await sandbox_metadata_db.unassign_sandboxes_to_jobs([dummy_jobid])
+
+    # Entity should not exists anymore
+    async with sandbox_metadata_db:
+        stmt = sqlalchemy.select(sb_EntityMapping.SBId).where(
+            sb_EntityMapping.EntityId == entity_id_1
+        )
+        res = await sandbox_metadata_db.conn.execute(stmt)
+    entity_sb_id = [row.SBId for row in res]
+    assert entity_sb_id == []
+
+    # Should not be assigned anymore
+    async with sandbox_metadata_db:
+        assert await sandbox_metadata_db.sandbox_is_assigned(pfn, sandbox_se) is False
+    # Check the mapping has been deleted
+    async with sandbox_metadata_db:
+        stmt = sqlalchemy.select(sb_EntityMapping.SBId)
+        res = await sandbox_metadata_db.conn.execute(stmt)
+    res_sb_id = [row.SBId for row in res]
+    assert sb_id_1 not in res_sb_id

diracx-db-0.0.1a11/src/diracx/db/sql/sandbox_metadata/db.py

@@ -1,96 +0,0 @@
-from __future__ import annotations
-
-import sqlalchemy
-from sqlalchemy import delete
-
-from diracx.core.models import SandboxInfo, UserInfo
-from diracx.db.sql.utils import BaseSQLDB, utcnow
-
-from .schema import Base as SandboxMetadataDBBase
-from .schema import sb_EntityMapping, sb_Owners, sb_SandBoxes
-
-
-class SandboxMetadataDB(BaseSQLDB):
-    metadata = SandboxMetadataDBBase.metadata
-
-    async def upsert_owner(self, user: UserInfo) -> int:
-        """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,
-        )
-        result = await self.conn.execute(stmt)
-        if owner_id := result.scalar_one_or_none():
-            return owner_id
-
-        stmt = sqlalchemy.insert(sb_Owners).values(
-            Owner=user.preferred_username,
-            OwnerGroup=user.dirac_group,
-            VO=user.vo,
-        )
-        result = await self.conn.execute(stmt)
-        return result.lastrowid
-
-    @staticmethod
-    def get_pfn(bucket_name: str, user: UserInfo, sandbox_info: SandboxInfo) -> str:
-        """Get the sandbox's user namespaced and content addressed PFN"""
-        parts = [
-            "S3",
-            bucket_name,
-            user.vo,
-            user.dirac_group,
-            user.preferred_username,
-            f"{sandbox_info.checksum_algorithm}:{sandbox_info.checksum}.{sandbox_info.format}",
-        ]
-        return "/" + "/".join(parts)
-
-    async def insert_sandbox(
-        self, se_name: str, user: UserInfo, 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(
-            OwnerId=owner_id,
-            SEName=se_name,
-            SEPFN=pfn,
-            Bytes=size,
-            RegistrationTime=utcnow(),
-            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
-
-    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)
-            .values(LastAccessTime=utcnow())
-        )
-        result = await self.conn.execute(stmt)
-        assert result.rowcount == 1
-
-    async def sandbox_is_assigned(self, se_name: str, pfn: str) -> bool:
-        """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
-        )
-        result = await self.conn.execute(stmt)
-        is_assigned = result.scalar_one()
-        return is_assigned
-        return True
-
-    async def unassign_sandbox_from_jobs(self, job_ids: list[int]):
-        """
-        Unassign sandbox from jobs
-        """
-        stmt = delete(sb_EntityMapping).where(
-            sb_EntityMapping.EntityId.in_(f"Job:{job_id}" for job_id in job_ids)
-        )
-        await self.conn.execute(stmt)

diracx-db-0.0.1a11/tests/test_sandbox_metadata.py

@@ -1,92 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import secrets
-from datetime import datetime
-
-import pytest
-import sqlalchemy
-
-from diracx.core.models import SandboxInfo, UserInfo
-from diracx.db.sql.sandbox_metadata.db import SandboxMetadataDB
-from diracx.db.sql.sandbox_metadata.schema import sb_SandBoxes
-
-
-@pytest.fixture
-async def sandbox_metadata_db(tmp_path):
-    sandbox_metadata_db = SandboxMetadataDB("sqlite+aiosqlite:///:memory:")
-    async with sandbox_metadata_db.engine_context():
-        async with sandbox_metadata_db.engine.begin() as conn:
-            await conn.run_sync(sandbox_metadata_db.metadata.create_all)
-        yield sandbox_metadata_db
-
-
-def test_get_pfn(sandbox_metadata_db: SandboxMetadataDB):
-    user_info = UserInfo(
-        sub="vo:sub", preferred_username="user1", dirac_group="group1", vo="vo"
-    )
-    sandbox_info = SandboxInfo(
-        checksum="checksum",
-        checksum_algorithm="sha256",
-        format="tar.bz2",
-        size=100,
-    )
-    pfn = sandbox_metadata_db.get_pfn("bucket1", user_info, sandbox_info)
-    assert pfn == "/S3/bucket1/vo/group1/user1/sha256:checksum.tar.bz2"
-
-
-async def test_insert_sandbox(sandbox_metadata_db: SandboxMetadataDB):
-    user_info = UserInfo(
-        sub="vo:sub", preferred_username="user1", dirac_group="group1", vo="vo"
-    )
-    pfn1 = secrets.token_hex()
-
-    # Make sure the sandbox doesn't already exist
-    db_contents = await _dump_db(sandbox_metadata_db)
-    assert pfn1 not in db_contents
-    async with sandbox_metadata_db:
-        with pytest.raises(sqlalchemy.exc.NoResultFound):
-            await sandbox_metadata_db.sandbox_is_assigned("SandboxSE", pfn1)
-
-    # Insert the sandbox
-    async with sandbox_metadata_db:
-        await sandbox_metadata_db.insert_sandbox("SandboxSE", user_info, pfn1, 100)
-    db_contents = await _dump_db(sandbox_metadata_db)
-    owner_id1, last_access_time1 = db_contents[pfn1]
-
-    # Inserting again should update the last access time
-    await asyncio.sleep(1)  # The timestamp only has second precision
-    async with sandbox_metadata_db:
-        await sandbox_metadata_db.insert_sandbox("SandboxSE", user_info, pfn1, 100)
-    db_contents = await _dump_db(sandbox_metadata_db)
-    owner_id2, last_access_time2 = db_contents[pfn1]
-    assert owner_id1 == owner_id2
-    assert last_access_time2 > last_access_time1
-
-    # The sandbox still hasn't been assigned
-    async with sandbox_metadata_db:
-        assert not await sandbox_metadata_db.sandbox_is_assigned("SandboxSE", pfn1)
-
-    # Inserting again should update the last access time
-    await asyncio.sleep(1)  # The timestamp only has second precision
-    last_access_time3 = (await _dump_db(sandbox_metadata_db))[pfn1][1]
-    assert last_access_time2 == last_access_time3
-    async with sandbox_metadata_db:
-        await sandbox_metadata_db.update_sandbox_last_access_time("SandboxSE", pfn1)
-    last_access_time4 = (await _dump_db(sandbox_metadata_db))[pfn1][1]
-    assert last_access_time2 < last_access_time4
-
-
-async def _dump_db(
-    sandbox_metadata_db: SandboxMetadataDB,
-) -> dict[str, tuple[int, datetime]]:
-    """Dump the contents of the sandbox metadata database
-
-    Returns a dict[pfn: str, (owner_id: int, last_access_time: datetime)]
-    """
-    async with sandbox_metadata_db:
-        stmt = sqlalchemy.select(
-            sb_SandBoxes.SEPFN, sb_SandBoxes.OwnerId, sb_SandBoxes.LastAccessTime
-        )
-        res = await sandbox_metadata_db.conn.execute(stmt)
-        return {row.SEPFN: (row.OwnerId, row.LastAccessTime) for row in res}