Add an equivalent forwardDISET for DiracX

src/DIRAC/Core/Security/DiracX.py

@@ -3,6 +3,7 @@
 __all__ = (
     "DiracXClient",
     "diracxTokenFromPEM",
+    "FutureClient",
 )
 
 import base64
@@ -67,6 +68,14 @@ def diracxTokenFromPEM(pemPath) -> dict[str, Any] | None:
         return json.loads(base64.b64decode(match).decode("utf-8"))
 
 
+class FutureClient:
+    """This is just a empty class to make sure that all the FutureClients
+    inherit from a common class.
+    """
+
+    ...
+
+
 @contextmanager
 def DiracXClient() -> _DiracClient:
     """Get a DiracX client instance with the current user's credentials"""

src/DIRAC/RequestManagementSystem/Agent/RequestOperations/ForwardDISET.py

@@ -10,13 +10,15 @@
 
     DISET forwarding operation handler
 """
+import importlib
 
 # imports
 from DIRAC import S_ERROR, S_OK, gConfig
 from DIRAC.ConfigurationSystem.Client.ConfigurationData import gConfigurationData
 from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getDNForUsername
 from DIRAC.Core.Base.Client import executeRPCStub
 from DIRAC.Core.Utilities import DEncode
+from DIRAC.Core.Security.DiracX import FutureClient
 from DIRAC.RequestManagementSystem.private.OperationHandlerBase import OperationHandlerBase
 
 ########################################################################
@@ -45,27 +47,54 @@ def __call__(self):
         """execute RPC stub"""
         # # decode arguments
         try:
-            decode, length = DEncode.decode(self.operation.Arguments)
-            self.log.debug(f"decoded len={length} val={decode}")
+            stub, length = DEncode.decode(self.operation.Arguments)
+            self.log.debug(f"decoded len={length} val={stub}")
         except ValueError as error:
             self.log.exception(error)
             self.operation.Error = str(error)
             self.operation.Status = "Failed"
             return S_ERROR(str(error))
 
-        # Ensure the forwarded request is done on behalf of the request owner
-        res = getDNForUsername(self.request.Owner)
-        if not res["OK"]:
-            return res
-        decode[0][1]["delegatedDN"] = res["Value"][0]
-        decode[0][1]["delegatedGroup"] = self.request.OwnerGroup
-
-        # ForwardDiset is supposed to be used with a host certificate
-        useServerCertificate = gConfig.useServerCertificate()
-        gConfigurationData.setOptionInCFG("/DIRAC/Security/UseServerCertificate", "true")
-        forward = executeRPCStub(decode)
-        if not useServerCertificate:
-            gConfigurationData.setOptionInCFG("/DIRAC/Security/UseServerCertificate", "false")
+        # This is the DISET rpcStub
+        if isinstance(stub, tuple):
+            # Ensure the forwarded request is done on behalf of the request owner
+            res = getDNForUsername(self.request.Owner)
+            if not res["OK"]:
+                return res
+            stub[0][1]["delegatedDN"] = res["Value"][0]
+            stub[0][1]["delegatedGroup"] = self.request.OwnerGroup
+
+            # ForwardDiset is supposed to be used with a host certificate
+            useServerCertificate = gConfig.useServerCertificate()
+            gConfigurationData.setOptionInCFG("/DIRAC/Security/UseServerCertificate", "true")
+            forward = executeRPCStub(stub)
+            if not useServerCertificate:
+                gConfigurationData.setOptionInCFG("/DIRAC/Security/UseServerCertificate", "false")
+        # DiracX stub
+        elif isinstance(stub, dict):
+            className = stub.get("dCls")
+            modName = stub.get("dMod")
+            methName = stub.get("dMeth")
+            methArgs = stub.get("args")
+            methKwArgs = stub.get("kwargs")
+            # Load the module
+            mod = importlib.import_module(modName)
+            # import the class
+            cl = getattr(mod, className)
+
+            # Check that cl is a subclass of JSerializable,
+            # and that we are not putting ourselves in trouble...
+            if not (isinstance(cl, type) and issubclass(cl, FutureClient)):
+                raise TypeError("Only subclasses of FutureClient can be decoded")
+
+            # Instantiate the object
+            obj = cl()
+            meth = getattr(obj, methName)
+            forward = meth(*methArgs, **methKwArgs)
+
+            pass
+        else:
+            raise TypeError("Unknwon type of stub")
 
         if not forward["OK"]:
             self.log.error("unable to execute operation", f"'{self.operation.Type}' : {forward['Message']}")

src/DIRAC/WorkloadManagementSystem/FutureClient/JobStateUpdateClient.py

@@ -2,11 +2,35 @@
 from datetime import datetime, timezone
 
 
-from DIRAC.Core.Security.DiracX import DiracXClient
-from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
+from DIRAC.Core.Security.DiracX import DiracXClient, FutureClient
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue, isReturnStructure
 from DIRAC.Core.Utilities.TimeUtilities import fromString
 
 
+def getStub(meth):
+    """Decorator to add an rpc like stub to DiracX adapter method
+    to be called by the ForwardDISET operation
+
+    """
+
+    @functools.wraps(meth)
+    def inner(self, *args, **kwargs):
+        dCls = self.__class__.__name__
+        dMod = self.__module__
+        res = meth(self, *args, **kwargs)
+        if isReturnStructure(res):
+            res["rpcStub"] = {
+                "dCls": dCls,
+                "dMod": dMod,
+                "dMeth": meth.__name__,
+                "args": args,
+                "kwargs": kwargs,
+            }
+        return res
+
+    return inner
+
+
 def stripValueIfOK(func):
     """Decorator to remove S_OK["Value"] from the return value of a function if it is OK.
 
@@ -26,7 +50,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-class JobStateUpdateClient:
+class JobStateUpdateClient(FutureClient):
     @stripValueIfOK
     @convertToReturnValue
     def sendHeartBeat(self, jobID: str | int, dynamicData: dict, staticData: dict):