diracx-db 0.0.1a11__py3-none-any.whl → 0.0.1a13__py3-none-any.whl

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

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+from typing import Any
+
 import sqlalchemy
-from sqlalchemy import delete
 
-from diracx.core.models import SandboxInfo, UserInfo
+from diracx.core.models import SandboxInfo, SandboxType, UserInfo
 from diracx.db.sql.utils import BaseSQLDB, utcnow
 
 from .schema import Base as SandboxMetadataDBBase
@@ -76,7 +77,7 @@ class SandboxMetadataDB(BaseSQLDB):
         result = await self.conn.execute(stmt)
         assert result.rowcount == 1
 
-    async def sandbox_is_assigned(self, se_name: str, pfn: str) -> bool:
+    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
@@ -84,13 +85,85 @@ class SandboxMetadataDB(BaseSQLDB):
         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)
+
+    @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,
+            )
         )
-        await self.conn.execute(stmt)
+        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.dist-info → diracx_db-0.0.1a13.dist-info}/METADATA

@@ -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.dist-info → diracx_db-0.0.1a13.dist-info}/RECORD

@@ -8,20 +8,20 @@ diracx/db/os/utils.py,sha256=mau0_2uRi-I3geefmKQRWFKo4JcIkIUADvnwBiQX700,9129
 diracx/db/sql/__init__.py,sha256=R6tk5lo1EHbt8joGDesesYHcc1swIq9T4AaSixhh7lA,252
 diracx/db/sql/utils.py,sha256=BuXjIuXN-_v8YkCoMoMhw2tHVUqG6lTBx-e4VEYWE8o,7857
 diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/auth/db.py,sha256=7GAdwD7g2YR86yQkP9ecuIbFd1h_wxVKnF_GfXfZqLA,9915
-diracx/db/sql/auth/schema.py,sha256=wutCjZ_uz21J0HHZjwoOXq3cLdlNY2lCR390yIJ_T60,2891
+diracx/db/sql/auth/db.py,sha256=mKjy5B8orw0yu6nOwxyzbBqyeE-J9iYq6fKjuELmr9g,10273
+diracx/db/sql/auth/schema.py,sha256=JCkSa2IRzqMHTpaSc9aB9h33XsFyEM_Ohsenex6xagY,2835
 diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/dummy/db.py,sha256=5PIPv6aKY7CGIwmvnGKowjVr9ZQWpbjFSd2PIX7YOUw,1627
 diracx/db/sql/dummy/schema.py,sha256=uEkGDNVZbmJecytkHY1CO-M1MiKxe5w1_h0joJMPC9E,680
 diracx/db/sql/jobs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/jobs/db.py,sha256=Y_2mx5kPTeuz6nxXVwGLzTssKsIH6nfnoTvWvilSgxA,29876
 diracx/db/sql/jobs/schema.py,sha256=YkxIdjTkvLlEZ9IQt86nj80eMvOPbcrfk9aisjmNpqY,9275
-diracx/db/sql/jobs/status_utility.py,sha256=0kAt623nh1O5wgsgktctdCmHEynO1nU0vn-7zakNeOA,10525
+diracx/db/sql/jobs/status_utility.py,sha256=_3Wdd11ShA4Z6HKr0_D_o8-zPZhdzgFpZSYAyYkH4Q0,10525
 diracx/db/sql/sandbox_metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/sandbox_metadata/db.py,sha256=HjlbnsT4cRMuFAcTL_sK3IqCehA7zISzR_d7xIGZoNk,3498
+diracx/db/sql/sandbox_metadata/db.py,sha256=0EDFMfOW_O3pEPTShqBCME9z4j-JKpyYM6-BBccr27E,6303
 diracx/db/sql/sandbox_metadata/schema.py,sha256=rngYYkJxBhjETBHGLD1CTipDGe44mRYR0wdaFoAJwp0,1400
-diracx_db-0.0.1a11.dist-info/METADATA,sha256=xRS26odR-83dZ1cj6lfGBEuzufrq59vKglTFvVFvVd4,681
-diracx_db-0.0.1a11.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-diracx_db-0.0.1a11.dist-info/entry_points.txt,sha256=xEFGu_zgmPgQPlUeFtdahQfQIboJ1ugFOK8eMio9gtw,271
-diracx_db-0.0.1a11.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
-diracx_db-0.0.1a11.dist-info/RECORD,,
+diracx_db-0.0.1a13.dist-info/METADATA,sha256=jmbXQvJykcvn3vGnxvO8GUGP3D3yjL-cXZwqXXJkzP4,681
+diracx_db-0.0.1a13.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+diracx_db-0.0.1a13.dist-info/entry_points.txt,sha256=xEFGu_zgmPgQPlUeFtdahQfQIboJ1ugFOK8eMio9gtw,271
+diracx_db-0.0.1a13.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
+diracx_db-0.0.1a13.dist-info/RECORD,,

{diracx_db-0.0.1a11.dist-info → diracx_db-0.0.1a13.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.42.0)
+Generator: bdist_wheel (0.43.0)
 Root-Is-Purelib: true
 Tag: py3-none-any