diracx-db 0.0.1a26__py3-none-any.whl → 0.0.1a28__py3-none-any.whl

diracx/db/os/utils.py

@@ -38,7 +38,7 @@ class BaseOSDB(metaclass=ABCMeta):
 
     The available OpenSearch databases are discovered by calling `BaseOSDB.available_urls`.
     This method returns a dictionary of database names to connection parameters.
-    The available databases are determined by the `diracx.db.os` entrypoint in
+    The available databases are determined by the `diracx.dbs.os` entrypoint in
     the `pyproject.toml` file and the connection parameters are taken from the
     environment variables prefixed with `DIRACX_OS_DB_{DB_NAME}`.
 
@@ -92,7 +92,9 @@ class BaseOSDB(metaclass=ABCMeta):
         """Return the available implementations of the DB in reverse priority order."""
         db_classes: list[type[BaseOSDB]] = [
             entry_point.load()
-            for entry_point in select_from_extension(group="diracx.db.os", name=db_name)
+            for entry_point in select_from_extension(
+                group="diracx.dbs.os", name=db_name
+            )
         ]
         if not db_classes:
             raise NotImplementedError(f"Could not find any matches for {db_name=}")
@@ -106,7 +108,7 @@ class BaseOSDB(metaclass=ABCMeta):
         prefixed with ``DIRACX_OS_DB_{DB_NAME}``.
         """
         conn_kwargs: dict[str, dict[str, Any]] = {}
-        for entry_point in select_from_extension(group="diracx.db.os"):
+        for entry_point in select_from_extension(group="diracx.dbs.os"):
             db_name = entry_point.name
             var_name = f"DIRACX_OS_DB_{entry_point.name.upper()}"
             if var_name in os.environ:

diracx/db/sql/auth/db.py

@@ -1,19 +1,16 @@
 from __future__ import annotations
 
-import hashlib
 import secrets
-from datetime import datetime
-from uuid import uuid4
+from uuid import UUID, uuid4
 
 from sqlalchemy import insert, select, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
 
 from diracx.core.exceptions import (
     AuthorizationError,
-    ExpiredFlowError,
-    PendingAuthorizationError,
+    TokenNotFoundError,
 )
-from diracx.db.sql.utils import BaseSQLDB, substract_date
+from diracx.db.sql.utils import BaseSQLDB, hash, substract_date
 
 from .schema import (
     AuthorizationFlows,
@@ -50,44 +47,25 @@ class AuthDB(BaseSQLDB):
 
         return (await self.conn.execute(stmt)).scalar_one()
 
-    async def get_device_flow(self, device_code: str, max_validity: int):
+    async def get_device_flow(self, device_code: str):
         """:raises: NoResultFound"""
         # The with_for_update
         # prevents that the token is retrieved
         # multiple time concurrently
-        stmt = select(
-            DeviceFlows,
-            (DeviceFlows.creation_time < substract_date(seconds=max_validity)).label(
-                "IsExpired"
-            ),
-        ).with_for_update()
+        stmt = select(DeviceFlows).with_for_update()
         stmt = stmt.where(
-            DeviceFlows.device_code == hashlib.sha256(device_code.encode()).hexdigest(),
+            DeviceFlows.device_code == hash(device_code),
         )
-        res = dict((await self.conn.execute(stmt)).one()._mapping)
-
-        if res["IsExpired"]:
-            raise ExpiredFlowError()
-
-        if res["Status"] == FlowStatus.READY:
-            # Update the status to Done before returning
-            await self.conn.execute(
-                update(DeviceFlows)
-                .where(
-                    DeviceFlows.device_code
-                    == hashlib.sha256(device_code.encode()).hexdigest()
-                )
-                .values(status=FlowStatus.DONE)
-            )
-            return res
-
-        if res["Status"] == FlowStatus.DONE:
-            raise AuthorizationError("Code was already used")
-
-        if res["Status"] == FlowStatus.PENDING:
-            raise PendingAuthorizationError()
+        return dict((await self.conn.execute(stmt)).one()._mapping)
 
-        raise AuthorizationError("Bad state in device flow")
+    async def update_device_flow_status(
+        self, device_code: str, status: FlowStatus
+    ) -> None:
+        stmt = update(DeviceFlows).where(
+            DeviceFlows.device_code == hash(device_code),
+        )
+        stmt = stmt.values(status=status)
+        await self.conn.execute(stmt)
 
     async def device_flow_insert_id_token(
         self, user_code: str, id_token: dict[str, str], max_validity: int
@@ -121,7 +99,7 @@ class AuthDB(BaseSQLDB):
             device_code = secrets.token_urlsafe()
 
             # Hash the the device_code to avoid leaking information
-            hashed_device_code = hashlib.sha256(device_code.encode()).hexdigest()
+            hashed_device_code = hash(device_code)
 
             stmt = insert(DeviceFlows).values(
                 client_id=client_id,
@@ -171,7 +149,7 @@ class AuthDB(BaseSQLDB):
         """
         # Hash the code to avoid leaking information
         code = secrets.token_urlsafe()
-        hashed_code = hashlib.sha256(code.encode()).hexdigest()
+        hashed_code = hash(code)
 
         stmt = update(AuthorizationFlows)
 
@@ -193,7 +171,8 @@ class AuthDB(BaseSQLDB):
         return code, row.RedirectURI
 
     async def get_authorization_flow(self, code: str, max_validity: int):
-        hashed_code = hashlib.sha256(code.encode()).hexdigest()
+        """Get the authorization flow details based on the code."""
+        hashed_code = hash(code)
         # The with_for_update
         # prevents that the token is retrieved
         # multiple time concurrently
@@ -203,54 +182,41 @@ class AuthDB(BaseSQLDB):
             AuthorizationFlows.creation_time > substract_date(seconds=max_validity),
         )
 
-        res = dict((await self.conn.execute(stmt)).one()._mapping)
-
-        if res["Status"] == FlowStatus.READY:
-            # Update the status to Done before returning
-            await self.conn.execute(
-                update(AuthorizationFlows)
-                .where(AuthorizationFlows.code == hashed_code)
-                .values(status=FlowStatus.DONE)
-            )
-
-            return res
+        return dict((await self.conn.execute(stmt)).one()._mapping)
 
-        if res["Status"] == FlowStatus.DONE:
-            raise AuthorizationError("Code was already used")
-
-        raise AuthorizationError("Bad state in authorization flow")
+    async def update_authorization_flow_status(
+        self, code: str, status: FlowStatus
+    ) -> None:
+        """Update the status of an authorization flow based on the code."""
+        hashed_code = hash(code)
+        await self.conn.execute(
+            update(AuthorizationFlows)
+            .where(AuthorizationFlows.code == hashed_code)
+            .values(status=status)
+        )
 
     async def insert_refresh_token(
         self,
+        jti: UUID,
         subject: str,
         preferred_username: str,
         scope: str,
-    ) -> tuple[str, datetime]:
+    ) -> None:
         """Insert a refresh token in the DB as well as user attributes
         required to generate access tokens.
         """
-        # Generate a JWT ID
-        jti = str(uuid4())
-
         # Insert values into the DB
         stmt = insert(RefreshTokens).values(
-            jti=jti,
+            jti=str(jti),
             sub=subject,
             preferred_username=preferred_username,
             scope=scope,
         )
         await self.conn.execute(stmt)
 
-        # Get the creation time of the new tuple: generated by the insert operation
-        stmt = select(RefreshTokens.creation_time)
-        stmt = stmt.where(RefreshTokens.jti == jti)
-        row = (await self.conn.execute(stmt)).one()
-
-        # Return the JWT ID and the creation time
-        return jti, row.CreationTime
-
-    async def get_refresh_token(self, jti: str) -> dict:
+    async def get_refresh_token(self, jti: UUID) -> dict:
         """Get refresh token details bound to a given JWT ID."""
+        jti = str(jti)
         # The with_for_update
         # prevents that the token is retrieved
         # multiple time concurrently
@@ -260,8 +226,8 @@ class AuthDB(BaseSQLDB):
         )
         try:
             res = dict((await self.conn.execute(stmt)).one()._mapping)
-        except NoResultFound:
-            return {}
+        except NoResultFound as e:
+            raise TokenNotFoundError(jti) from e
 
         return res
 
@@ -285,11 +251,11 @@ class AuthDB(BaseSQLDB):
 
         return refresh_tokens
 
-    async def revoke_refresh_token(self, jti: str):
+    async def revoke_refresh_token(self, jti: UUID):
         """Revoke a token given by its JWT ID."""
         await self.conn.execute(
             update(RefreshTokens)
-            .where(RefreshTokens.jti == jti)
+            .where(RefreshTokens.jti == str(jti))
             .values(status=RefreshTokenStatus.REVOKED)
         )
 

diracx/db/sql/job/db.py

@@ -4,14 +4,12 @@ from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any
 
 from sqlalchemy import bindparam, case, delete, func, insert, select, update
-from sqlalchemy.exc import IntegrityError, NoResultFound
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import BindParameter
 
-from diracx.core.exceptions import InvalidQueryError, JobNotFoundError
+from diracx.core.exceptions import InvalidQueryError
 from diracx.core.models import (
-    LimitedJobStatusReturn,
     SearchSpec,
     SortSpec,
 )
@@ -46,6 +44,7 @@ class JobDB(BaseSQLDB):
     jdl_2_db_parameters = ["JobName", "JobType", "JobGroup"]
 
     async def summary(self, group_by, search) -> list[dict[str, str | int]]:
+        """Get a summary of the jobs."""
         columns = _get_columns(Jobs.__table__, group_by)
 
         stmt = select(*columns, func.count(Jobs.job_id).label("count"))
@@ -69,6 +68,7 @@ class JobDB(BaseSQLDB):
         per_page: int = 100,
         page: int | None = None,
     ) -> tuple[int, list[dict[Any, Any]]]:
+        """Search for jobs in the database."""
         # Find which columns to select
         columns = _get_columns(Jobs.__table__, parameters)
 
@@ -98,7 +98,24 @@ class JobDB(BaseSQLDB):
             dict(row._mapping) async for row in (await self.conn.stream(stmt))
         ]
 
+    async def create_job(self, compressed_original_jdl: str):
+        """Used to insert a new job with original JDL. Returns inserted job id."""
+        result = await self.conn.execute(
+            JobJDLs.__table__.insert().values(
+                JDL="",
+                JobRequirements="",
+                OriginalJDL=compressed_original_jdl,
+            )
+        )
+        return result.lastrowid
+
+    async def delete_jobs(self, job_ids: list[int]):
+        """Delete jobs from the database."""
+        stmt = delete(JobJDLs).where(JobJDLs.job_id.in_(job_ids))
+        await self.conn.execute(stmt)
+
     async def insert_input_data(self, lfns: dict[int, list[str]]):
+        """Insert input data for jobs."""
         await self.conn.execute(
             InputData.__table__.insert(),
             [
@@ -111,27 +128,8 @@ class JobDB(BaseSQLDB):
             ],
         )
 
-    async def set_job_attributes(self, job_id, job_data):
-        """TODO: add myDate and force parameters."""
-        if "Status" in job_data:
-            job_data = job_data | {"LastUpdateTime": datetime.now(tz=timezone.utc)}
-        stmt = update(Jobs).where(Jobs.job_id == job_id).values(job_data)
-        await self.conn.execute(stmt)
-
-    async def create_job(self, original_jdl):
-        """Used to insert a new job with original JDL. Returns inserted job id."""
-        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
-
-        result = await self.conn.execute(
-            JobJDLs.__table__.insert().values(
-                JDL="",
-                JobRequirements="",
-                OriginalJDL=compressJDL(original_jdl),
-            )
-        )
-        return result.lastrowid
-
     async def insert_job_attributes(self, jobs_to_update: dict[int, dict]):
+        """Insert the job attributes."""
         await self.conn.execute(
             Jobs.__table__.insert(),
             [
@@ -145,8 +143,6 @@ class JobDB(BaseSQLDB):
 
     async def update_job_jdls(self, jdls_to_update: dict[int, str]):
         """Used to update the JDL, typically just after inserting the original JDL, or rescheduling, for example."""
-        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
-
         await self.conn.execute(
             JobJDLs.__table__.update().where(
                 JobJDLs.__table__.c.JobID == bindparam("b_JobID")
@@ -154,67 +150,15 @@ class JobDB(BaseSQLDB):
             [
                 {
                     "b_JobID": job_id,
-                    "JDL": compressJDL(jdl),
+                    "JDL": compressed_jdl,
                 }
-                for job_id, jdl in jdls_to_update.items()
+                for job_id, compressed_jdl in jdls_to_update.items()
             ],
         )
 
-    async def check_and_prepare_job(
-        self,
-        job_id,
-        class_ad_job,
-        class_ad_req,
-        owner,
-        owner_group,
-        job_attrs,
-        vo,
-    ):
-        """Check Consistency of Submitted JDL and set some defaults
-        Prepare subJDL with Job Requirements.
-        """
-        from DIRAC.Core.Utilities.DErrno import EWMSSUBM, cmpError
-        from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
-        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import (
-            checkAndPrepareJob,
-        )
-
-        ret_val = checkAndPrepareJob(
-            job_id,
-            class_ad_job,
-            class_ad_req,
-            owner,
-            owner_group,
-            job_attrs,
-            vo,
-        )
-
-        if not ret_val["OK"]:
-            if cmpError(ret_val, EWMSSUBM):
-                await self.set_job_attributes(job_id, job_attrs)
-
-            returnValueOrRaise(ret_val)
-
-    async def set_job_jdl(self, job_id, jdl):
-        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
-
-        stmt = (
-            update(JobJDLs).where(JobJDLs.job_id == job_id).values(JDL=compressJDL(jdl))
-        )
-        await self.conn.execute(stmt)
-
-    async def set_job_jdl_bulk(self, jdls):
-        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
-
-        await self.conn.execute(
-            JobJDLs.__table__.update().where(
-                JobJDLs.__table__.c.JobID == bindparam("b_JobID")
-            ),
-            [{"b_JobID": jid, "JDL": compressJDL(jdl)} for jid, jdl in jdls.items()],
-        )
-
-    async def set_job_attributes_bulk(self, job_data):
-        """TODO: add myDate and force parameters."""
+    async def set_job_attributes(self, job_data):
+        """Update the parameters of the given jobs."""
+        # TODO: add myDate and force parameters.
         for job_id in job_data.keys():
             if "Status" in job_data[job_id]:
                 job_data[job_id].update(
@@ -240,11 +184,8 @@ class JobDB(BaseSQLDB):
         )
         await self.conn.execute(stmt)
 
-    async def get_job_jdls(
-        self, job_ids, original: bool = False
-    ) -> dict[int | str, str]:
-        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import extractJDL
-
+    async def get_job_jdls(self, job_ids, original: bool = False) -> dict[int, str]:
+        """Get the JDLs for the given jobs."""
         if original:
             stmt = select(JobJDLs.job_id, JobJDLs.original_jdl).where(
                 JobJDLs.job_id.in_(job_ids)
@@ -254,37 +195,9 @@ class JobDB(BaseSQLDB):
                 JobJDLs.job_id.in_(job_ids)
             )
 
-        return {
-            jobid: extractJDL(jdl)
-            for jobid, jdl in (await self.conn.execute(stmt))
-            if jdl
-        }
-
-    async def get_job_status(self, job_id: int) -> LimitedJobStatusReturn:
-        try:
-            stmt = select(
-                Jobs.status, Jobs.minor_status, Jobs.application_status
-            ).where(Jobs.job_id == job_id)
-            return LimitedJobStatusReturn(
-                **dict((await self.conn.execute(stmt)).one()._mapping)
-            )
-        except NoResultFound as e:
-            raise JobNotFoundError(job_id) from e
-
-    async def set_job_command(self, job_id: int, command: str, arguments: str = ""):
-        """Store a command to be passed to the job together with the next heart beat."""
-        try:
-            stmt = insert(JobCommands).values(
-                JobID=job_id,
-                Command=command,
-                Arguments=arguments,
-                ReceptionTime=datetime.now(tz=timezone.utc),
-            )
-            await self.conn.execute(stmt)
-        except IntegrityError as e:
-            raise JobNotFoundError(job_id) from e
+        return {jobid: jdl for jobid, jdl in (await self.conn.execute(stmt)) if jdl}
 
-    async def set_job_command_bulk(self, commands):
+    async def set_job_commands(self, commands: list[tuple[int, str, str]]):
         """Store a command to be passed to the job together with the next heart beat."""
         await self.conn.execute(
             insert(JobCommands),
@@ -298,12 +211,6 @@ class JobDB(BaseSQLDB):
                 for job_id, command, arguments in commands
             ],
         )
-        # FIXME handle IntegrityError
-
-    async def delete_jobs(self, job_ids: list[int]):
-        """Delete jobs from the database."""
-        stmt = delete(JobJDLs).where(JobJDLs.job_id.in_(job_ids))
-        await self.conn.execute(stmt)
 
     async def set_properties(
         self, properties: dict[int, dict[str, Any]], update_timestamp: bool = False

diracx/db/sql/job_logging/db.py

@@ -1,20 +1,18 @@
 from __future__ import annotations
 
 import time
-from datetime import datetime, timezone
+from datetime import timezone
 from typing import TYPE_CHECKING
 
-from pydantic import BaseModel
-from sqlalchemy import delete, func, insert, select
+from sqlalchemy import delete, func, select
 
 if TYPE_CHECKING:
     pass
 
 from collections import defaultdict
 
-from diracx.core.exceptions import JobNotFoundError
 from diracx.core.models import (
-    JobStatus,
+    JobLoggingRecord,
     JobStatusReturn,
 )
 
@@ -27,61 +25,12 @@ from .schema import (
 MAGIC_EPOC_NUMBER = 1270000000
 
 
-class JobLoggingRecord(BaseModel):
-    job_id: int
-    status: JobStatus
-    minor_status: str
-    application_status: str
-    date: datetime
-    source: str
-
-
 class JobLoggingDB(BaseSQLDB):
     """Frontend for the JobLoggingDB. Provides the ability to store changes with timestamps."""
 
     metadata = JobLoggingDBBase.metadata
 
-    async def insert_record(
-        self,
-        job_id: int,
-        status: JobStatus,
-        minor_status: str,
-        application_status: str,
-        date: datetime,
-        source: str,
-    ):
-        """Add a new entry to the JobLoggingDB table. One, two or all the three status
-        components (status, minorStatus, applicationStatus) can be specified.
-        Optionally the time stamp of the status can
-        be provided in a form of a string in a format '%Y-%m-%d %H:%M:%S' or
-        as datetime.datetime object. If the time stamp is not provided the current
-        UTC time is used.
-        """
-        # First, fetch the maximum seq_num for the given job_id
-        seqnum_stmt = select(func.coalesce(func.max(LoggingInfo.seq_num) + 1, 1)).where(
-            LoggingInfo.job_id == job_id
-        )
-        seqnum = await self.conn.scalar(seqnum_stmt)
-
-        epoc = (
-            time.mktime(date.timetuple())
-            + date.microsecond / 1000000.0
-            - MAGIC_EPOC_NUMBER
-        )
-
-        stmt = insert(LoggingInfo).values(
-            job_id=int(job_id),
-            seq_num=seqnum,
-            status=status,
-            minor_status=minor_status,
-            application_status=application_status[:255],
-            status_time=date,
-            status_time_order=epoc,
-            source=source[:32],
-        )
-        await self.conn.execute(stmt)
-
-    async def bulk_insert_record(
+    async def insert_records(
         self,
         records: list[JobLoggingRecord],
     ):
@@ -103,15 +52,20 @@ class JobLoggingDB(BaseSQLDB):
             .group_by(LoggingInfo.job_id)
         )
 
-        seqnum = {jid: seqnum for jid, seqnum in (await self.conn.execute(seqnum_stmt))}
+        seqnums = {
+            jid: seqnum for jid, seqnum in (await self.conn.execute(seqnum_stmt))
+        }
         # IF a seqnum is not found, then assume it does not exist and the first sequence number is 1.
         # https://docs.sqlalchemy.org/en/20/orm/queryguide/dml.html#orm-bulk-insert-statements
-        await self.conn.execute(
-            LoggingInfo.__table__.insert(),
-            [
+        values = []
+        for record in records:
+            if record.job_id not in seqnums:
+                seqnums[record.job_id] = 1
+
+            values.append(
                 {
                     "JobID": record.job_id,
-                    "SeqNum": seqnum.get(record.job_id, 1),
+                    "SeqNum": seqnums[record.job_id],
                     "Status": record.status,
                     "MinorStatus": record.minor_status,
                     "ApplicationStatus": record.application_status[:255],
@@ -119,8 +73,12 @@ class JobLoggingDB(BaseSQLDB):
                     "StatusTimeOrder": get_epoc(record.date),
                     "StatusSource": record.source[:32],
                 }
-                for record in records
-            ],
+            )
+            seqnums[record.job_id] = seqnums[record.job_id] + 1
+
+        await self.conn.execute(
+            LoggingInfo.__table__.insert(),
+            values,
         )
 
     async def get_records(self, job_ids: list[int]) -> dict[int, JobStatusReturn]:
@@ -201,25 +159,7 @@ class JobLoggingDB(BaseSQLDB):
         stmt = delete(LoggingInfo).where(LoggingInfo.job_id.in_(job_ids))
         await self.conn.execute(stmt)
 
-    async def get_wms_time_stamps(self, job_id):
-        """Get TimeStamps for job MajorState transitions
-        return a {State:timestamp} dictionary.
-        """
-        result = {}
-        stmt = select(
-            LoggingInfo.status,
-            LoggingInfo.status_time_order,
-        ).where(LoggingInfo.job_id == job_id)
-        rows = await self.conn.execute(stmt)
-        if not rows.rowcount:
-            raise JobNotFoundError(job_id) from None
-
-        for event, etime in rows:
-            result[event] = str(etime + MAGIC_EPOC_NUMBER)
-
-        return result
-
-    async def get_wms_time_stamps_bulk(self, job_ids):
+    async def get_wms_time_stamps(self, job_ids):
         """Get TimeStamps for job MajorState transitions for multiple jobs at once
         return a {JobID: {State:timestamp}} dictionary.
         """