diracx-db 0.0.1a17__py3-none-any.whl → 0.0.1a19__py3-none-any.whl

diracx/db/__main__.py

@@ -31,6 +31,7 @@ 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():

diracx/db/os/utils.py

@@ -7,9 +7,10 @@ import json
 import logging
 import os
 from abc import ABCMeta, abstractmethod
+from collections.abc import AsyncIterator
 from contextvars import ContextVar
 from datetime import datetime
-from typing import Any, AsyncIterator, Self
+from typing import Any, Self
 
 from opensearchpy import AsyncOpenSearch
 
@@ -29,6 +30,48 @@ class OpenSearchDBUnavailable(DBUnavailable, OpenSearchDBError):
 
 
 class BaseOSDB(metaclass=ABCMeta):
+    """This should be the base class of all the OpenSearch DiracX DBs.
+
+    The details covered here should be handled automatically by the service and
+    task machinery of DiracX and this documentation exists for informational
+    purposes.
+
+    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 `pyproject.toml` file and the connection parameters are taken from the
+    environment variables prefixed with `DIRACX_OS_DB_{DB_NAME}`.
+
+    If extensions to DiracX are being used, there can be multiple implementations
+    of the same database. To list the available implementations use
+    `BaseOSDB.available_implementations(db_name)`. The first entry in this list
+    will be the preferred implementation and it can be initialized by calling
+    its `__init__` function with the connection parameters previously obtained
+    from `BaseOSDB.available_urls`.
+
+    To control the lifetime of the OpenSearch client, the `BaseOSDB.client_context`
+    asynchronous context manager should be entered. When inside this context
+    manager, the client can be accessed with `BaseOSDB.client`.
+
+    Upon entering, the DB class can then be used as an asynchronous context
+    manager to perform operations. Currently this context manager has no effect
+    however it must be used as it may be used in future. When inside this
+    context manager, the DB connection can be accessed with `BaseOSDB.client`.
+
+    For example:
+
+    ```python
+    db_name = ...
+    conn_params = BaseOSDB.available_urls()[db_name]
+    MyDBClass = BaseOSDB.available_implementations(db_name)[0]
+
+    db = MyDBClass(conn_params)
+    async with db.client_context:
+        async with db:
+            # Do something with the OpenSearch client
+    ```
+    """
+
     # TODO: Make metadata an abstract property
     fields: dict
     index_prefix: str
@@ -77,13 +120,15 @@ class BaseOSDB(metaclass=ABCMeta):
     @classmethod
     def session(cls) -> Self:
         """This is just a fake method such that the Dependency overwrite has
-        a hash to use"""
+        a hash to use.
+        """
         raise NotImplementedError("This should never be called")
 
     @property
     def client(self) -> AsyncOpenSearch:
         """Just a getter for _client, making sure we entered
-        the context manager"""
+        the context manager.
+        """
         if self._client is None:
             raise RuntimeError(f"{self.__class__} was used before entering")
         return self._client
@@ -91,17 +136,18 @@ class BaseOSDB(metaclass=ABCMeta):
     @contextlib.asynccontextmanager
     async def client_context(self) -> AsyncIterator[None]:
         """Context manage to manage the client lifecycle.
-        This is called when starting fastapi
+        This is called when starting fastapi.
 
         """
         assert self._client is None, "client_context cannot be nested"
         async with AsyncOpenSearch(**self._connection_kwargs) as self._client:
-            yield
-        self._client = None
+            try:
+                yield
+            finally:
+                self._client = None
 
     async def ping(self):
-        """
-        Check whether the connection to the DB is still working.
+        """Check whether the connection to the DB is still working.
         We could enable the ``pre_ping`` in the engine, but this would
         be ran at every query.
         """
@@ -113,7 +159,7 @@ class BaseOSDB(metaclass=ABCMeta):
     async def __aenter__(self):
         """This is entered on every request.
         At the moment it does nothing, however, we keep it here
-        in case we ever want to use OpenSearch equivalent of a transaction
+        in case we ever want to use OpenSearch equivalent of a transaction.
         """
         assert not self._conn.get(), "BaseOSDB context cannot be nested"
         assert self._client is not None, "client_context hasn't been entered"
@@ -122,9 +168,7 @@ class BaseOSDB(metaclass=ABCMeta):
 
     async def __aexit__(self, exc_type, exc, tb):
         assert self._conn.get()
-        self._client = None
         self._conn.set(False)
-        return
 
     async def create_index_template(self) -> None:
         template_body = {
@@ -237,6 +281,11 @@ def apply_search_filters(db_fields, search):
                     operator, field_name, field_type, {"keyword", "long", "date"}
                 )
                 result["must"].append({"terms": {field_name: query["values"]}})
+            case "not in":
+                require_type(
+                    operator, field_name, field_type, {"keyword", "long", "date"}
+                )
+                result["must_not"].append({"terms": {field_name: query["values"]}})
             # TODO: Implement like and ilike
             # If the pattern is a simple "col like 'abc%'", we can use a prefix query
             # Else we need to use a wildcard query where we replace % with * and _ with ?

diracx/db/sql/__init__.py

@@ -1,7 +1,17 @@
 from __future__ import annotations
 
-__all__ = ("AuthDB", "JobDB", "JobLoggingDB", "SandboxMetadataDB", "TaskQueueDB")
+__all__ = (
+    "AuthDB",
+    "JobDB",
+    "JobLoggingDB",
+    "PilotAgentsDB",
+    "SandboxMetadataDB",
+    "TaskQueueDB",
+)
 
 from .auth.db import AuthDB
-from .jobs.db import JobDB, JobLoggingDB, TaskQueueDB
+from .job.db import JobDB
+from .job_logging.db import JobLoggingDB
+from .pilot_agents.db import PilotAgentsDB
 from .sandbox_metadata.db import SandboxMetadataDB
+from .task_queue.db import TaskQueueDB

diracx/db/sql/auth/db.py

@@ -35,7 +35,7 @@ class AuthDB(BaseSQLDB):
     async def device_flow_validate_user_code(
         self, user_code: str, max_validity: int
     ) -> str:
-        """Validate that the user_code can be used (Pending status, not expired)
+        """Validate that the user_code can be used (Pending status, not expired).
 
         Returns the scope field for the given user_code
 
@@ -51,9 +51,7 @@ class AuthDB(BaseSQLDB):
         return (await self.conn.execute(stmt)).scalar_one()
 
     async def get_device_flow(self, device_code: str, max_validity: int):
-        """
-        :raises: NoResultFound
-        """
+        """:raises: NoResultFound"""
         # The with_for_update
         # prevents that the token is retrieved
         # multiple time concurrently
@@ -94,9 +92,7 @@ class AuthDB(BaseSQLDB):
     async def device_flow_insert_id_token(
         self, user_code: str, id_token: dict[str, str], max_validity: int
     ) -> None:
-        """
-        :raises: AuthorizationError if no such code or status not pending
-        """
+        """:raises: AuthorizationError if no such code or status not pending"""
         stmt = update(DeviceFlows)
         stmt = stmt.where(
             DeviceFlows.user_code == user_code,
@@ -170,11 +166,9 @@ class AuthDB(BaseSQLDB):
     async def authorization_flow_insert_id_token(
         self, uuid: str, id_token: dict[str, str], max_validity: int
     ) -> tuple[str, str]:
+        """Returns code, redirect_uri
+        :raises: AuthorizationError if no such uuid or status not pending.
         """
-        returns code, redirect_uri
-        :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()
@@ -232,8 +226,7 @@ class AuthDB(BaseSQLDB):
         preferred_username: str,
         scope: str,
     ) -> tuple[str, datetime]:
-        """
-        Insert a refresh token in the DB as well as user attributes
+        """Insert a refresh token in the DB as well as user attributes
         required to generate access tokens.
         """
         # Generate a JWT ID
@@ -257,9 +250,7 @@ class AuthDB(BaseSQLDB):
         return jti, row.creation_time
 
     async def get_refresh_token(self, jti: str) -> dict:
-        """
-        Get refresh token details bound to a given JWT ID
-        """
+        """Get refresh token details bound to a given JWT ID."""
         # The with_for_update
         # prevents that the token is retrieved
         # multiple time concurrently
@@ -275,7 +266,7 @@ class AuthDB(BaseSQLDB):
         return res
 
     async def get_user_refresh_tokens(self, subject: str | None = None) -> list[dict]:
-        """Get a list of refresh token details based on a subject ID (not revoked)"""
+        """Get a list of refresh token details based on a subject ID (not revoked)."""
         # Get a list of refresh tokens
         stmt = select(RefreshTokens).with_for_update()
 
@@ -295,7 +286,7 @@ class AuthDB(BaseSQLDB):
         return refresh_tokens
 
     async def revoke_refresh_token(self, jti: str):
-        """Revoke a token given by its JWT ID"""
+        """Revoke a token given by its JWT ID."""
         await self.conn.execute(
             update(RefreshTokens)
             .where(RefreshTokens.jti == jti)
@@ -303,7 +294,7 @@ class AuthDB(BaseSQLDB):
         )
 
     async def revoke_user_refresh_tokens(self, subject):
-        """Revoke all the refresh tokens belonging to a user (subject ID)"""
+        """Revoke all the refresh tokens belonging to a user (subject ID)."""
         await self.conn.execute(
             update(RefreshTokens)
             .where(RefreshTokens.sub == subject)

diracx/db/sql/auth/schema.py

@@ -15,12 +15,11 @@ Base = declarative_base()
 
 
 class FlowStatus(Enum):
-    """
-    The normal flow is
+    """The normal flow is
     PENDING -> READY -> DONE
     Pending is upon insertion
     Ready/Error is set in response to IdP
-    Done means the user has been issued the dirac token
+    Done means the user has been issued the dirac token.
     """
 
     # The flow is ongoing
@@ -64,9 +63,8 @@ class AuthorizationFlows(Base):
 
 
 class RefreshTokenStatus(Enum):
-    """
-    The normal flow is
-    CREATED -> REVOKED
+    """The normal flow is
+    CREATED -> REVOKED.
 
     Note1: There is no EXPIRED status as it can be calculated from a creation time
     Note2: As part of the refresh token rotation mechanism, the revoked token should be retained
@@ -82,7 +80,7 @@ class RefreshTokenStatus(Enum):
 
 class RefreshTokens(Base):
     """Store attributes bound to a refresh token, as well as specific user attributes
-    that might be then used to generate access tokens
+    that might be then used to generate access tokens.
     """
 
     __tablename__ = "RefreshTokens"

diracx/db/sql/dummy/db.py

@@ -11,8 +11,7 @@ from .schema import Cars, Owners
 
 
 class DummyDB(BaseSQLDB):
-    """
-    This DummyDB is just to illustrate some important aspect of writing
+    """This DummyDB is just to illustrate some important aspect of writing
     DB classes in DiracX.
 
     It is mostly pure SQLAlchemy, with a few convention
@@ -27,7 +26,7 @@ class DummyDB(BaseSQLDB):
         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__, stmt, search)
+        stmt = apply_search_filters(Cars.__table__.columns.__getitem__, stmt, search)
         stmt = stmt.group_by(*columns)
 
         # Execute the query