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

diracx/db/__main__.py

@@ -31,7 +31,6 @@ async def init_sql():
     from diracx.db.sql.utils import BaseSQLDB
 
     for db_name, db_url in BaseSQLDB.available_urls().items():
-
         logger.info("Initialising %s", db_name)
         db = BaseSQLDB.available_implementations(db_name)[0](db_url)
         async with db.engine_context():
@@ -40,6 +39,7 @@ async def init_sql():
                 if db._db_url.startswith("sqlite"):
                     await conn.exec_driver_sql("PRAGMA foreign_keys=ON")
                 await conn.run_sync(db.metadata.create_all)
+                await db.post_create(conn)
 
 
 async def init_os():

diracx/db/exceptions.py

@@ -1,2 +1,5 @@
-class DBUnavailable(Exception):
+from __future__ import annotations
+
+
+class DBUnavailableError(Exception):
     pass

diracx/db/os/job_parameters.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from datetime import UTC, datetime
+
 from diracx.db.os.utils import BaseOSDB
 
 
@@ -7,19 +9,35 @@ class JobParametersDB(BaseOSDB):
     fields = {
         "JobID": {"type": "long"},
         "timestamp": {"type": "date"},
+        "PilotAgent": {"type": "keyword"},
+        "Pilot_Reference": {"type": "keyword"},
+        "JobGroup": {"type": "keyword"},
         "CPUNormalizationFactor": {"type": "long"},
         "NormCPUTime(s)": {"type": "long"},
-        "Memory(kB)": {"type": "long"},
+        "Memory(MB)": {"type": "long"},
+        "LocalAccount": {"type": "keyword"},
         "TotalCPUTime(s)": {"type": "long"},
-        "MemoryUsed(kb)": {"type": "long"},
-        "HostName": {"type": "keyword"},
+        "PayloadPID": {"type": "long"},
+        "HostName": {"type": "text"},
         "GridCE": {"type": "keyword"},
+        "CEQueue": {"type": "keyword"},
+        "BatchSystem": {"type": "keyword"},
         "ModelName": {"type": "keyword"},
         "Status": {"type": "keyword"},
         "JobType": {"type": "keyword"},
     }
-    index_prefix = "mysetup_elasticjobparameters_index_"
+    # TODO: Does this need to be configurable?
+    index_prefix = "job_parameters"
+
+    def index_name(self, vo, doc_id: int) -> str:
+        split = int(int(doc_id) // 1e6)
+        # The index name must be lowercase or opensearchpy will throw.
+        return f"{self.index_prefix}_{vo.lower()}_{split}m"
 
-    def index_name(self, doc_id: int) -> str:
-        # TODO: Remove setup and replace "123.0m" with "120m"?
-        return f"{self.index_prefix}_{doc_id // 1e6:.1f}m"
+    def upsert(self, vo, doc_id, document):
+        document = {
+            "JobID": doc_id,
+            "timestamp": int(datetime.now(tz=UTC).timestamp() * 1000),
+            **document,
+        }
+        return super().upsert(vo, doc_id, document)

diracx/db/os/utils.py

@@ -16,7 +16,7 @@ from opensearchpy import AsyncOpenSearch
 
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.extensions import select_from_extension
-from diracx.db.exceptions import DBUnavailable
+from diracx.db.exceptions import DBUnavailableError
 
 logger = logging.getLogger(__name__)
 
@@ -25,7 +25,7 @@ class OpenSearchDBError(Exception):
     pass
 
 
-class OpenSearchDBUnavailable(DBUnavailable, OpenSearchDBError):
+class OpenSearchDBUnavailableError(DBUnavailableError, OpenSearchDBError):
     pass
 
 
@@ -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}`.
 
@@ -77,7 +77,7 @@ class BaseOSDB(metaclass=ABCMeta):
     index_prefix: str
 
     @abstractmethod
-    def index_name(self, doc_id: int) -> str: ...
+    def index_name(self, vo: str, doc_id: int) -> str: ...
 
     def __init__(self, connection_kwargs: dict[str, Any]) -> None:
         self._client: AsyncOpenSearch | None = None
@@ -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:
@@ -152,7 +154,7 @@ class BaseOSDB(metaclass=ABCMeta):
         be ran at every query.
         """
         if not await self.client.ping():
-            raise OpenSearchDBUnavailable(
+            raise OpenSearchDBUnavailableError(
                 f"Failed to connect to {self.__class__.__qualname__}"
             )
 
@@ -180,15 +182,20 @@ class BaseOSDB(metaclass=ABCMeta):
         )
         assert result["acknowledged"]
 
-    async def upsert(self, doc_id, document) -> None:
-        # TODO: Implement properly
+    async def upsert(self, vo: str, doc_id: int, document: Any) -> None:
+        index_name = self.index_name(vo, doc_id)
         response = await self.client.update(
-            index=self.index_name(doc_id),
+            index=index_name,
             id=doc_id,
             body={"doc": document, "doc_as_upsert": True},
             params=dict(retry_on_conflict=10),
         )
-        print(f"{response=}")
+        logger.debug(
+            "Upserted document %s in index %s with response: %s",
+            doc_id,
+            index_name,
+            response,
+        )
 
     async def search(
         self, parameters, search, sorts, *, per_page: int = 100, page: int | None = None

diracx/db/sql/auth/db.py

@@ -1,19 +1,21 @@
 from __future__ import annotations
 
-import hashlib
+import logging
 import secrets
-from datetime import datetime
-from uuid import uuid4
+from datetime import UTC, datetime
+from itertools import pairwise
 
-from sqlalchemy import insert, select, update
+from dateutil.rrule import MONTHLY, rrule
+from sqlalchemy import insert, select, text, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
+from sqlalchemy.ext.asyncio import AsyncConnection
+from uuid_utils import UUID, uuid7
 
 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, uuid7_from_datetime
 
 from .schema import (
     AuthorizationFlows,
@@ -28,10 +30,72 @@ from .schema import Base as AuthDBBase
 USER_CODE_ALPHABET = "BCDFGHJKLMNPQRSTVWXZ"
 MAX_RETRY = 5
 
+logger = logging.getLogger(__name__)
+
 
 class AuthDB(BaseSQLDB):
     metadata = AuthDBBase.metadata
 
+    @classmethod
+    async def post_create(cls, conn: AsyncConnection) -> None:
+        """Create partitions if it is a MySQL DB and it does not have
+        it yet and the table does not have any data yet.
+        We do this as a post_create step as sqlalchemy does not support
+        partition so well.
+        """
+        if conn.dialect.name == "mysql":
+            check_partition_query = text(
+                "SELECT PARTITION_NAME FROM information_schema.partitions "
+                "WHERE TABLE_NAME = 'RefreshTokens' AND PARTITION_NAME is not NULL"
+            )
+            partition_names = (await conn.execute(check_partition_query)).all()
+
+            if not partition_names:
+                # Create a monthly partition from today until 2 years
+                # The partition are named p_<year>_<month>
+                start_date = datetime.now(tz=UTC).replace(
+                    day=1, hour=0, minute=0, second=0, microsecond=0
+                )
+                end_date = start_date.replace(year=start_date.year + 2)
+
+                dates = [
+                    dt for dt in rrule(MONTHLY, dtstart=start_date, until=end_date)
+                ]
+
+                partition_list = []
+                for name, limit in pairwise(dates):
+                    partition_list.append(
+                        f"PARTITION p_{name.year}_{name.month} "
+                        f"VALUES LESS THAN ('{str(uuid7_from_datetime(limit, randomize=False)).replace('-', '')}')"
+                    )
+                partition_list.append("PARTITION p_future VALUES LESS THAN (MAXVALUE)")
+
+                alter_query = text(
+                    f"ALTER TABLE RefreshTokens PARTITION BY RANGE COLUMNS (JTI) ({','.join(partition_list)})"
+                )
+
+                check_table_empty_query = text("SELECT * FROM RefreshTokens LIMIT 1")
+                refresh_table_content = (
+                    await conn.execute(check_table_empty_query)
+                ).all()
+                if refresh_table_content:
+                    logger.warning(
+                        "RefreshTokens table not empty. Run the following query yourself"
+                    )
+                    logger.warning(alter_query)
+                    return
+
+                await conn.execute(alter_query)
+
+                partition_names = (
+                    await conn.execute(
+                        check_partition_query, {"table_name": "RefreshTokens"}
+                    )
+                ).all()
+                assert partition_names, (
+                    f"There should be partitions now {partition_names}"
+                )
+
     async def device_flow_validate_user_code(
         self, user_code: str, max_validity: int
     ) -> str:
@@ -50,44 +114,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(
-                "is_expired"
-            ),
-        ).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["is_expired"]:
-            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")
+        return dict((await self.conn.execute(stmt)).one()._mapping)
 
-        if res["status"] == FlowStatus.PENDING:
-            raise PendingAuthorizationError()
-
-        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 +166,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,
@@ -133,6 +178,10 @@ class AuthDB(BaseSQLDB):
                 await self.conn.execute(stmt)
 
             except IntegrityError:
+                logger.warning(
+                    "Device flow code collision detected, retrying (user_code=%s)",
+                    user_code,
+                )
                 continue
 
             return user_code, device_code
@@ -148,7 +197,7 @@ class AuthDB(BaseSQLDB):
         code_challenge_method: str,
         redirect_uri: str,
     ) -> str:
-        uuid = str(uuid4())
+        uuid = str(uuid7())
 
         stmt = insert(AuthorizationFlows).values(
             uuid=uuid,
@@ -171,7 +220,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)
 
@@ -190,10 +239,11 @@ class AuthDB(BaseSQLDB):
         stmt = select(AuthorizationFlows.code, AuthorizationFlows.redirect_uri)
         stmt = stmt.where(AuthorizationFlows.uuid == uuid)
         row = (await self.conn.execute(stmt)).one()
-        return code, row.redirect_uri
+        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 +253,39 @@ 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 dict((await self.conn.execute(stmt)).one()._mapping)
 
-            return res
-
-        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.creation_time
-
-    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 +295,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 +320,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/auth/schema.py

@@ -1,13 +1,21 @@
+from __future__ import annotations
+
 from enum import Enum, auto
 
 from sqlalchemy import (
     JSON,
+    Index,
     String,
     Uuid,
 )
 from sqlalchemy.orm import declarative_base
 
-from diracx.db.sql.utils import Column, DateNowColumn, EnumColumn, NullColumn
+from diracx.db.sql.utils import (
+    Column,
+    DateNowColumn,
+    EnumColumn,
+    NullColumn,
+)
 
 USER_CODE_LENGTH = 8
 
@@ -39,27 +47,27 @@ class FlowStatus(Enum):
 
 class DeviceFlows(Base):
     __tablename__ = "DeviceFlows"
-    user_code = Column(String(USER_CODE_LENGTH), primary_key=True)
-    status = EnumColumn(FlowStatus, server_default=FlowStatus.PENDING.name)
-    creation_time = DateNowColumn()
-    client_id = Column(String(255))
-    scope = Column(String(1024))
-    device_code = Column(String(128), unique=True)  # Should be a hash
-    id_token = NullColumn(JSON())
+    user_code = Column("UserCode", String(USER_CODE_LENGTH), primary_key=True)
+    status = EnumColumn("Status", FlowStatus, server_default=FlowStatus.PENDING.name)
+    creation_time = DateNowColumn("CreationTime")
+    client_id = Column("ClientID", String(255))
+    scope = Column("Scope", String(1024))
+    device_code = Column("DeviceCode", String(128), unique=True)  # Should be a hash
+    id_token = NullColumn("IDToken", JSON())
 
 
 class AuthorizationFlows(Base):
     __tablename__ = "AuthorizationFlows"
-    uuid = Column(Uuid(as_uuid=False), primary_key=True)
-    status = EnumColumn(FlowStatus, server_default=FlowStatus.PENDING.name)
-    client_id = Column(String(255))
-    creation_time = DateNowColumn()
-    scope = Column(String(1024))
-    code_challenge = Column(String(255))
-    code_challenge_method = Column(String(8))
-    redirect_uri = Column(String(255))
-    code = NullColumn(String(255))  # Should be a hash
-    id_token = NullColumn(JSON())
+    uuid = Column("UUID", Uuid(as_uuid=False), primary_key=True)
+    status = EnumColumn("Status", FlowStatus, server_default=FlowStatus.PENDING.name)
+    client_id = Column("ClientID", String(255))
+    creation_time = DateNowColumn("CreationTime")
+    scope = Column("Scope", String(1024))
+    code_challenge = Column("CodeChallenge", String(255))
+    code_challenge_method = Column("CodeChallengeMethod", String(8))
+    redirect_uri = Column("RedirectURI", String(255))
+    code = NullColumn("Code", String(255))  # Should be a hash
+    id_token = NullColumn("IDToken", JSON())
 
 
 class RefreshTokenStatus(Enum):
@@ -85,13 +93,13 @@ class RefreshTokens(Base):
 
     __tablename__ = "RefreshTokens"
     # Refresh token attributes
-    jti = Column(Uuid(as_uuid=False), primary_key=True)
+    jti = Column("JTI", Uuid(as_uuid=False), primary_key=True)
     status = EnumColumn(
-        RefreshTokenStatus, server_default=RefreshTokenStatus.CREATED.name
+        "Status", RefreshTokenStatus, server_default=RefreshTokenStatus.CREATED.name
     )
-    creation_time = DateNowColumn()
-    scope = Column(String(1024))
+    scope = Column("Scope", String(1024))
 
     # User attributes bound to the refresh token
-    sub = Column(String(1024))
-    preferred_username = Column(String(255))
+    sub = Column("Sub", String(256), index=True)
+
+    __table_args__ = (Index("index_status_sub", status, sub),)

diracx/db/sql/dummy/db.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-from uuid import UUID
+from sqlalchemy import insert
+from uuid_utils import UUID
 
-from sqlalchemy import func, insert, select
-
-from diracx.db.sql.utils import BaseSQLDB, apply_search_filters
+from diracx.db.sql.utils import BaseSQLDB
 
 from .schema import Base as DummyDBBase
 from .schema import Cars, Owners
@@ -23,18 +22,7 @@ class DummyDB(BaseSQLDB):
     metadata = DummyDBBase.metadata
 
     async def summary(self, group_by, search) -> list[dict[str, str | int]]:
-        columns = [Cars.__table__.columns[x] for x in group_by]
-
-        stmt = select(*columns, func.count(Cars.licensePlate).label("count"))
-        stmt = apply_search_filters(Cars.__table__.columns.__getitem__, stmt, search)
-        stmt = stmt.group_by(*columns)
-
-        # Execute the query
-        return [
-            dict(row._mapping)
-            async for row in (await self.conn.stream(stmt))
-            if row.count > 0  # type: ignore
-        ]
+        return await self._summary(Cars, group_by, search)
 
     async def insert_owner(self, name: str) -> int:
         stmt = insert(Owners).values(name=name)
@@ -44,7 +32,7 @@ class DummyDB(BaseSQLDB):
 
     async def insert_car(self, license_plate: UUID, model: str, owner_id: int) -> int:
         stmt = insert(Cars).values(
-            licensePlate=license_plate, model=model, ownerID=owner_id
+            license_plate=license_plate, model=model, owner_id=owner_id
         )
 
         result = await self.conn.execute(stmt)

diracx/db/sql/dummy/schema.py

@@ -1,5 +1,7 @@
 # The utils class define some boilerplate types that should be used
 # in place of the SQLAlchemy one. Have a look at them
+from __future__ import annotations
+
 from sqlalchemy import ForeignKey, Integer, String, Uuid
 from sqlalchemy.orm import declarative_base
 
@@ -10,13 +12,13 @@ Base = declarative_base()
 
 class Owners(Base):
     __tablename__ = "Owners"
-    ownerID = Column(Integer, primary_key=True, autoincrement=True)
-    creation_time = DateNowColumn()
-    name = Column(String(255))
+    owner_id = Column("OwnerID", Integer, primary_key=True, autoincrement=True)
+    creation_time = DateNowColumn("CreationTime")
+    name = Column("Name", String(255))
 
 
 class Cars(Base):
     __tablename__ = "Cars"
-    licensePlate = Column(Uuid(), primary_key=True)
-    model = Column(String(255))
-    ownerID = Column(Integer, ForeignKey(Owners.ownerID))
+    license_plate = Column("LicensePlate", Uuid(), primary_key=True)
+    model = Column("Model", String(255))
+    owner_id = Column("OwnerID", Integer, ForeignKey(Owners.owner_id))