feat(framework): make logging with context var more generic

docs/source/DeveloperGuide/AddingNewComponents/Utilities/gLogger/gLogger/Advanced/index.rst

@@ -217,6 +217,68 @@ This option can not be modified in the children of *gLogger*, even by
 *gLogger* itself after the configuration, so the children receive
 the *gLogger* configuration.
 
+Add variables to different *Logging* objects depending on the context
+---------------------------------------------------------------------
+
+In complex cases, it can be useful to have loggers that change depending on
+the execution context, without having to pass logger instances explicitly
+through multiple layers of function calls.
+
+Python's `contextvars` module provides context-local storage, which can be used
+to store and retrieve context-specific data, such as logger instances.
+
+gLogger supports the use of context variables to manage loggers in a flexible way.
+
+Provide a Context Logger
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+When you have a *Logging* instance that you want to use in a specific context,
+you can set it in the context variable:
+
+::
+
+    # Create a logger instance
+    logger = gLogger.getSubLogger("MyContextLogger")
+
+    # Set it in the context variable
+    contextLogger.set(logger)
+
+Then, the instances within the context block will use the shared *Logging* object
+set in the context variable:
+
+::
+
+    with setContextLogger(contextualLogger):
+        # Any logging within this block will use contextualLogger
+        obj = MyClass()
+        obj.do_something()  # This will use contextualLogger
+
+Consume a Context Logger
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+In functions or classes that need to log messages, you can retrieve the logger
+from the context variable:
+
+::
+
+    class MyClass:
+        def __init__(self):
+            # Get the default logger if no context logger is set
+            self._defaultLogger = gLogger.getSubLogger("MyClass")
+
+        @property
+        def log(self):
+            # Return the context logger if set, otherwise the default logger
+            return contextLogger.get() or self._defaultLogger
+
+        @log.setter
+        def log(self, value):
+            # Optionally, allow setting a new default logger
+            self._defaultLogger = value
+
+        def do_something(self):
+            self.log.notice("Doing something")
+
 Some examples and summaries
 ---------------------------
 

src/DIRAC/FrameworkSystem/Client/Logger.py

@@ -1,7 +1,11 @@
+from DIRAC.FrameworkSystem.private.standardLogging.LoggingContext import contextLogger, setContextLogger
 from DIRAC.FrameworkSystem.private.standardLogging.LoggingRoot import LoggingRoot
 
 gLogger = LoggingRoot()
 
 
 def getLogger():
     return gLogger
+
+
+__all__ = ["contextLogger", "setContextLogger", "getLogger"]

src/DIRAC/FrameworkSystem/private/standardLogging/LoggingContext.py

@@ -0,0 +1,16 @@
+""" Logging context module"""
+
+# Context variable for the logger (adapted to the request of the pilot reference)
+import contextvars
+from contextlib import contextmanager
+
+contextLogger = contextvars.ContextVar("Logger", default=None)
+
+
+@contextmanager
+def setContextLogger(logger_name):
+    token = contextLogger.set(logger_name)
+    try:
+        yield
+    finally:
+        contextLogger.reset(token)

src/DIRAC/FrameworkSystem/private/standardLogging/test/Test_Logging_ContextVars.py

@@ -0,0 +1,71 @@
+""" Test the context variable logger """
+
+from DIRAC import gLogger
+from DIRAC.FrameworkSystem.private.standardLogging.Logging import Logging
+from DIRAC.FrameworkSystem.private.standardLogging.test.TestLogUtilities import gLoggerReset
+from DIRAC.FrameworkSystem.private.standardLogging.LoggingContext import contextLogger, setContextLogger
+
+
+class A:
+    def __init__(self):
+        # Get the logger from the context variable
+        self._defaultLogger = gLogger.getSubLogger("A")
+
+    # Use a property to get and set the logger, this is necessary to use the context variable
+    @property
+    def log(self):
+        return contextLogger.get() or self._defaultLogger
+
+    @log.setter
+    def log(self, value: Logging):
+        self._defaultLogger = value
+
+    def do_something(self):
+        self.log.notice("A is doing something")
+
+
+class B:
+    def __init__(self, a: A, pilotRef: str = None):
+        self.a = A()
+
+        # Get the logger from the context variable
+        if pilotRef:
+            self.log = gLogger.getLocalSubLogger(f"[{pilotRef}]B")
+            contextLogger.set(self.log)
+        else:
+            self.log = gLogger.getSubLogger("B")
+
+    def do_something_else(self):
+        with setContextLogger(self.log):
+            self.a.do_something()
+            self.log.notice("B is doing something else")
+
+
+def test_contextvar_logger():
+    capturedBackend, log, sublog = gLoggerReset()
+
+    # Create an instance of A
+    a = A()
+
+    # Create an instance of B and call its method without setting the pilotRef
+    # Log signature coming from A and B should be different
+    b1 = B(a)
+    b1.do_something_else()
+    assert "Framework/B NOTICE: A is doing something" in capturedBackend.getvalue()
+    assert "Framework/B NOTICE: B is doing something else" in capturedBackend.getvalue()
+
+    # Create an instance of B and call its method with setting the pilotRef
+    # Log signature coming from A and B should be similar because of the pilotRef
+    capturedBackend.truncate(0)
+
+    b2 = B(a, "pilotRef")
+    b2.do_something_else()
+    assert "Framework/[pilotRef]B NOTICE: A is doing something" in capturedBackend.getvalue()
+    assert "Framework/[pilotRef]B NOTICE: B is doing something else" in capturedBackend.getvalue()
+
+    # Now we check that the logger of b1 is not the same as the logger of b2 (b1 should still use its own logger)
+    capturedBackend.truncate(0)
+
+    b1.do_something_else()
+    assert "Framework/B NOTICE: A is doing something" in capturedBackend.getvalue()
+    assert "Framework/B NOTICE: B is doing something else" in capturedBackend.getvalue()

src/DIRAC/WorkloadManagementSystem/Client/Matcher.py

@@ -9,14 +9,14 @@
 from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
 from DIRAC.Core.Security import Properties
 from DIRAC.Core.Utilities.PrettyPrint import printDict
+from DIRAC.FrameworkSystem.Client.Logger import setContextLogger
 from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus
 from DIRAC.WorkloadManagementSystem.Client import JobStatus, PilotStatus
 from DIRAC.WorkloadManagementSystem.Client.Limiter import Limiter
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
 from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
 from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
 from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB, multiValueMatchFields, singleValueDefFields
-from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import setPilotRefLogger
 
 
 class PilotVersionError(Exception):
@@ -61,7 +61,7 @@ def __init__(self, pilotAgentsDB=None, jobDB=None, tqDB=None, jlDB=None, opsHelp
 
     def selectJob(self, resourceDescription, credDict):
         """Main job selection function to find the highest priority job matching the resource capacity"""
-        with setPilotRefLogger(self.log):
+        with setContextLogger(self.log):
             startTime = time.time()
 
             resourceDict = self._getResourceDict(resourceDescription, credDict)

src/DIRAC/WorkloadManagementSystem/DB/JobDB.py

@@ -21,6 +21,7 @@
 from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
 from DIRAC.Core.Utilities.DErrno import EWMSJMAN, EWMSSUBM, cmpError
 from DIRAC.Core.Utilities.ReturnValues import S_ERROR, S_OK
+from DIRAC.FrameworkSystem.Client.Logger import contextLogger
 from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus
 from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
@@ -32,7 +33,6 @@
     extractJDL,
     fixJDL,
 )
-from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger
 
 
 class JobDB(DB):
@@ -69,7 +69,7 @@ def __init__(self, parentLogger=None):
 
     @property
     def log(self):
-        return pilotRefLogger.get() or self._defaultLogger
+        return contextLogger.get() or self._defaultLogger
 
     @log.setter
     def log(self, value):

src/DIRAC/WorkloadManagementSystem/DB/JobLoggingDB.py

@@ -11,7 +11,7 @@
 from DIRAC import S_ERROR, S_OK
 from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Utilities import TimeUtilities
-from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger
+from DIRAC.FrameworkSystem.Client.Logger import contextLogger
 
 MAGIC_EPOC_NUMBER = 1270000000
 
@@ -29,7 +29,7 @@ def __init__(self, parentLogger=None):
 
     @property
     def log(self):
-        return pilotRefLogger.get() or self._defaultLogger
+        return contextLogger.get() or self._defaultLogger
 
     @log.setter
     def log(self, value):

src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py

@@ -30,9 +30,9 @@
 from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Utilities import DErrno
 from DIRAC.Core.Utilities.MySQL import _quotedList
+from DIRAC.FrameworkSystem.Client.Logger import contextLogger
 from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus
 from DIRAC.WorkloadManagementSystem.Client import PilotStatus
-from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger
 
 
 class PilotAgentsDB(DB):
@@ -43,7 +43,7 @@ def __init__(self, parentLogger=None):
 
     @property
     def log(self):
-        return pilotRefLogger.get() or self._defaultLogger
+        return contextLogger.get() or self._defaultLogger
 
     @log.setter
     def log(self, value):

src/DIRAC/WorkloadManagementSystem/DB/TaskQueueDB.py

@@ -11,8 +11,8 @@
 from DIRAC.Core.Security import Properties
 from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
+from DIRAC.FrameworkSystem.Client.Logger import contextLogger
 from DIRAC.WorkloadManagementSystem.private.SharesCorrector import SharesCorrector
-from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger
 
 DEFAULT_GROUP_SHARE = 1000
 TQ_MIN_SHARE = 0.001
@@ -68,7 +68,7 @@ def __init__(self, parentLogger=None):
 
     @property
     def log(self):
-        return pilotRefLogger.get() or self._defaultLogger
+        return contextLogger.get() or self._defaultLogger
 
     @log.setter
     def log(self, value):

src/DIRAC/WorkloadManagementSystem/Service/MatcherHandler.py

@@ -56,7 +56,7 @@ def export_requestJob(self, resourceDescription):
 
         resourceDescription["Setup"] = self.serviceInfoDict["clientSetup"]
         credDict = self.getRemoteCredentials()
-        pilotRef = resourceDescription.get("PilotReference")
+        pilotRef = resourceDescription.get("PilotReference", "Unknown")
 
         try:
             opsHelper = Operations(group=credDict["group"])