From dbfa5bb709e3fd22a739cc7f36171dce6f7017ca Mon Sep 17 00:00:00 2001 From: Christophe Haen Date: Wed, 27 May 2026 14:32:29 +0200 Subject: [PATCH] feat: support global OS prefix --- diracx-db/src/diracx/db/os/utils.py | 6 +++++- diracx-routers/src/diracx/routers/factory.py | 5 ++++- diracx-tasks/src/diracx/tasks/plumbing/factory.py | 5 ++++- diracx-testing/src/diracx/testing/mock_osdb.py | 4 +++- toto | 0 5 files changed, 16 insertions(+), 4 deletions(-) create mode 100644 toto diff --git a/diracx-db/src/diracx/db/os/utils.py b/diracx-db/src/diracx/db/os/utils.py index 1d80b9958..8d76d62e7 100644 --- a/diracx-db/src/diracx/db/os/utils.py +++ b/diracx-db/src/diracx/db/os/utils.py @@ -79,13 +79,17 @@ class BaseOSDB(metaclass=ABCMeta): @abstractmethod def index_name(self, vo: str, doc_id: int) -> str: ... - def __init__(self, connection_kwargs: dict[str, Any]) -> None: + def __init__( + self, connection_kwargs: dict[str, Any], *, global_prefix: str = "" + ) -> None: self._client: AsyncOpenSearch | None = None self._connection_kwargs = connection_kwargs # We use a ContextVar to make sure that self._conn # is specific to each context, and avoid parallel # route executions to overlap self._conn: ContextVar[bool] = ContextVar("_conn", default=False) + if global_prefix: + self.index_prefix = f"{global_prefix}_{self.index_prefix}" @classmethod def available_implementations(cls, db_name: str) -> list[type[BaseOSDB]]: diff --git a/diracx-routers/src/diracx/routers/factory.py b/diracx-routers/src/diracx/routers/factory.py index 342e6d5cd..33b56f556 100644 --- a/diracx-routers/src/diracx/routers/factory.py +++ b/diracx-routers/src/diracx/routers/factory.py @@ -216,11 +216,14 @@ def create_app_inner( raise Exception("No SQL database could be initialized, aborting") # Add the OpenSearch DBs to the application + os_global_prefix = os.environ.get("DIRACX_FACTORY_OS_GLOBAL_PREFIX", "") available_os_db_classes: set[type[BaseOSDB]] = set() for db_name, connection_kwargs in os_database_conn_kwargs.items(): os_db_classes = BaseOSDB.available_implementations(db_name) # The first DB is the highest priority one - os_db = os_db_classes[0](connection_kwargs=connection_kwargs) + os_db = os_db_classes[0]( + connection_kwargs=connection_kwargs, global_prefix=os_global_prefix + ) app.lifetime_functions.append(os_db.client_context) # Add overrides for all the DB classes, including those from extensions # This means vanilla DiracX routers get an instance of the extension's DB diff --git a/diracx-tasks/src/diracx/tasks/plumbing/factory.py b/diracx-tasks/src/diracx/tasks/plumbing/factory.py index 17187a056..5bce00335 100644 --- a/diracx-tasks/src/diracx/tasks/plumbing/factory.py +++ b/diracx-tasks/src/diracx/tasks/plumbing/factory.py @@ -305,9 +305,12 @@ async def setup_dependency_overrides( ) # --- OS databases --- + os_global_prefix = os.environ.get("DIRACX_FACTORY_OS_GLOBAL_PREFIX", "") for db_name, conn_kwargs in BaseOSDB.available_urls().items(): os_db_classes = BaseOSDB.available_implementations(db_name) - os_db = os_db_classes[0](connection_kwargs=conn_kwargs) + os_db = os_db_classes[0]( + connection_kwargs=conn_kwargs, global_prefix=os_global_prefix + ) await stack.enter_async_context(os_db.client_context()) for os_db_class in os_db_classes: overrides[os_db_class.session] = partial(_db_context, os_db) diff --git a/diracx-testing/src/diracx/testing/mock_osdb.py b/diracx-testing/src/diracx/testing/mock_osdb.py index 7cc20551a..103cc818e 100644 --- a/diracx-testing/src/diracx/testing/mock_osdb.py +++ b/diracx-testing/src/diracx/testing/mock_osdb.py @@ -36,7 +36,9 @@ class JobParametersDB(MockOSDBMixin, JobParametersDB): JobParametersDB = type("JobParametersDB", (MockOSDBMixin, JobParametersDB), {}) """ - def __init__(self, connection_kwargs: dict[str, Any]) -> None: + def __init__( + self, connection_kwargs: dict[str, Any], global_prefix: str = "" + ) -> None: from sqlalchemy import JSON, Column, DateTime, Integer, MetaData, String, Table # Dynamically create a subclass of BaseSQLDB so we get clearer errors diff --git a/toto b/toto new file mode 100644 index 000000000..e69de29bb