DIRAC 9.0.0a42__py3-none-any.whl → 9.0.7__py3-none-any.whl

DIRAC/DataManagementSystem/scripts/dirac_admin_allow_se.py

@@ -17,6 +17,7 @@ def main():
     remove = False
     site = ""
     mute = False
+    userName = ""
 
     Script.registerSwitch("r", "AllowRead", "     Allow only reading from the storage element")
     Script.registerSwitch("w", "AllowWrite", "     Allow only writing to the storage element")
@@ -25,6 +26,7 @@ def main():
     Script.registerSwitch("a", "All", "    Allow all access to the storage element")
     Script.registerSwitch("m", "Mute", "     Do not send email")
     Script.registerSwitch("S:", "Site=", "     Allow all SEs associated to site")
+    Script.registerSwitch("t:", "tokenOwner=", "     Optional Name of the token owner")
     # Registering arguments will automatically add their description to the help menu
     Script.registerArgument(["seGroupList: list of SEs or comma-separated SEs"])
 
@@ -48,6 +50,8 @@ def main():
             mute = True
         if switch[0].lower() in ("s", "site"):
             site = switch[1]
+        if switch[0] in ("t", "tokenOwner"):
+            userName = switch[1]
 
     # imports
     from DIRAC import gLogger
@@ -69,15 +73,16 @@ def main():
     ses = resolveSEGroup(ses)
     diracAdmin = DiracAdmin()
 
-    res = getProxyInfo()
-    if not res["OK"]:
-        gLogger.error("Failed to get proxy information", res["Message"])
-        DIRAC.exit(2)
-
-    userName = res["Value"].get("username")
     if not userName:
-        gLogger.error("Failed to get username for proxy")
-        DIRAC.exit(2)
+        res = getProxyInfo()
+        if not res["OK"]:
+            gLogger.error("Failed to get proxy information", res["Message"])
+            DIRAC.exit(2)
+
+        userName = res["Value"].get("username")
+        if not userName:
+            gLogger.error("Failed to get username for proxy")
+            DIRAC.exit(2)
 
     if site:
         res = getSites()

DIRAC/DataManagementSystem/scripts/dirac_admin_ban_se.py

@@ -18,6 +18,7 @@ def main():
     remove = True
     sites = []
     mute = False
+    userName = ""
 
     Script.registerSwitch("r", "BanRead", "     Ban only reading from the storage element")
     Script.registerSwitch("w", "BanWrite", "     Ban writing to the storage element")
@@ -28,6 +29,7 @@ def main():
     Script.registerSwitch(
         "S:", "Site=", "     Ban all SEs associate to site (note that if writing is allowed, check is always allowed)"
     )
+    Script.registerSwitch("t:", "tokenOwner=", "     Optional Name of the token owner")
     # Registering arguments will automatically add their description to the help menu
     Script.registerArgument(["seGroupList: list of SEs or comma-separated SEs"])
 
@@ -56,6 +58,8 @@ def main():
             mute = True
         if switch[0].lower() in ("s", "site"):
             sites = switch[1].split(",")
+        if switch[0] in ("t", "tokenOwner"):
+            userName = switch[1]
 
     # from DIRAC.ConfigurationSystem.Client.CSAPI           import CSAPI
     from DIRAC import gLogger
@@ -68,15 +72,16 @@ def main():
     ses = resolveSEGroup(ses)
     diracAdmin = DiracAdmin()
 
-    res = getProxyInfo()
-    if not res["OK"]:
-        gLogger.error("Failed to get proxy information", res["Message"])
-        DIRAC.exit(2)
-
-    userName = res["Value"].get("username")
     if not userName:
-        gLogger.error("Failed to get username for proxy")
-        DIRAC.exit(2)
+        res = getProxyInfo()
+        if not res["OK"]:
+            gLogger.error("Failed to get proxy information", res["Message"])
+            DIRAC.exit(2)
+
+        userName = res["Value"].get("username")
+        if not userName:
+            gLogger.error("Failed to get username for proxy")
+            DIRAC.exit(2)
 
     for site in sites:
         res = DMSHelpers().getSEsForSite(site)

DIRAC/DataManagementSystem/scripts/dirac_dms_create_moving_request.py

@@ -41,6 +41,7 @@ class CreateMovingRequest:
         self.flags = [
             ("C", "CheckMigration", "Ensure the LFNs are migrated to tape before removing any replicas"),
             ("X", "Execute", "Put Requests, else dryrun"),
+            ("", "SourceOnly", "Only treat files that are already at the Source-SE"),
         ]
         self.registerSwitchesAndParseCommandLine()
         self.getLFNList()
@@ -208,6 +209,7 @@ class CreateMovingRequest:
 
         replicate = Operation()
         replicate.Type = "ReplicateAndRegister"
+        replicate.SourceSE = ",".join(self.switches.get("SourceSE", []))
         replicate.TargetSE = self.switches.get("TargetSE")
         self.addLFNs(replicate, lfnChunk, addPFN=True)
         request.addOperation(replicate)

DIRAC/DataManagementSystem/scripts/dirac_dms_protocol_matrix.py

@@ -177,7 +177,6 @@ def main():
     for src, dst in ((x, y) for x in fromSE for y in targetSE):
         if ftsTab:
             try:
-                # breakpoint()
                 fts3TpcProto = fts3Plugin.selectTPCProtocols(sourceSEName=ses[src].name, destSEName=ses[dst].name)
                 res = ses[dst].generateTransferURLsBetweenSEs(lfn, ses[src], fts3TpcProto)
             except ValueError as e:

DIRAC/FrameworkSystem/Client/BundleDeliveryClient.py

@@ -143,13 +143,9 @@ class BundleDeliveryClient(Client):
         if "X509_CERT_DIR" in os.environ:
             X509_CERT_DIR = os.environ["X509_CERT_DIR"]
             del os.environ["X509_CERT_DIR"]
-        casLocation = Locations.getCAsLocation()
-        if not casLocation:
-            casLocation = Locations.getCAsDefaultLocation()
-        result = self.syncDir("CAs", casLocation)
         if X509_CERT_DIR:
             os.environ["X509_CERT_DIR"] = X509_CERT_DIR
-        return result
+        return self.syncDir("CAs", Locations.getCAsLocation())
 
     def syncCRLs(self):
         """Synchronize CRLs
@@ -160,10 +156,9 @@ class BundleDeliveryClient(Client):
         if "X509_CERT_DIR" in os.environ:
             X509_CERT_DIR = os.environ["X509_CERT_DIR"]
             del os.environ["X509_CERT_DIR"]
-        result = self.syncDir("CRLs", Locations.getCAsLocation())
         if X509_CERT_DIR:
             os.environ["X509_CERT_DIR"] = X509_CERT_DIR
-        return result
+        return self.syncDir("CRLs", Locations.getCAsLocation())
 
     def getCAs(self):
         """This method can be used to create the CAs. If the file can not be created,

DIRAC/FrameworkSystem/Client/ComponentInstaller.py

@@ -60,7 +60,6 @@ from collections import defaultdict
 
 import importlib_metadata as metadata
 import importlib_resources
-import MySQLdb
 from diraccfg import CFG
 from prompt_toolkit import prompt
 
@@ -69,7 +68,6 @@ from DIRAC import gConfig, gLogger, rootPath
 from DIRAC.ConfigurationSystem.Client import PathFinder
 from DIRAC.ConfigurationSystem.Client.CSAPI import CSAPI
 from DIRAC.ConfigurationSystem.Client.Helpers import (
-    CSGlobals,
     cfgInstallPath,
     cfgInstallSection,
     cfgPath,
@@ -85,6 +83,7 @@ from DIRAC.Core.Security.Properties import (
     PRODUCTION_MANAGEMENT,
     PROXY_MANAGEMENT,
     SERVICE_ADMINISTRATOR,
+    SITE_MANAGER,
     TRUSTED_HOST,
 )
 from DIRAC.Core.Utilities.Extensions import (
@@ -96,7 +95,6 @@ from DIRAC.Core.Utilities.Extensions import (
     findServices,
 )
 from DIRAC.Core.Utilities.File import mkDir, mkLink
-from DIRAC.Core.Utilities.MySQL import MySQL
 from DIRAC.Core.Utilities.PrettyPrint import printTable
 from DIRAC.Core.Utilities.ReturnValues import S_ERROR, S_OK
 from DIRAC.Core.Utilities.Subprocess import systemCall
@@ -432,6 +430,8 @@ class ComponentInstaller:
         defaultHostProperties = [
             TRUSTED_HOST,
             CS_ADMINISTRATOR,
+            SERVICE_ADMINISTRATOR,
+            SITE_MANAGER,
             JOB_ADMINISTRATOR,
             FULL_DELEGATION,
             PROXY_MANAGEMENT,
@@ -2053,6 +2053,8 @@ class ComponentInstaller:
         """
         Install requested DB in MySQL server
         """
+        import MySQLdb
+
         dbName = MySQLdb.escape_string(dbName.encode()).decode()
         if not self.mysqlRootPwd:
             rootPwdPath = cfgInstallPath("Database", "RootPwd")
@@ -2106,7 +2108,8 @@ class ComponentInstaller:
 
         perms = (
             "SELECT,INSERT,LOCK TABLES,UPDATE,DELETE,CREATE,DROP,ALTER,REFERENCES,"
-            "CREATE VIEW,SHOW VIEW,INDEX,TRIGGER,ALTER ROUTINE,CREATE ROUTINE"
+            "CREATE VIEW,SHOW VIEW,INDEX,TRIGGER,ALTER ROUTINE,CREATE ROUTINE,"
+            "CREATE TEMPORARY TABLES"
         )
         cmd = f"GRANT {perms} ON `{dbName}`.* TO '{self.mysqlUser}'@'%'"
         result = self.execMySQL(cmd)
@@ -2199,6 +2202,8 @@ class ComponentInstaller:
         """
         Execute MySQL Command
         """
+        from DIRAC.Core.Utilities.MySQL import MySQL
+
         if not self.mysqlRootPwd:
             return S_ERROR("MySQL root password is not defined")
         if dbName not in self.db:

DIRAC/FrameworkSystem/Client/ProxyManagerClient.py

@@ -418,7 +418,7 @@ class ProxyManagerClient(metaclass=DIRACSingleton.DIRACSingleton):
                 proxyToConnect=proxyToConnect,
             )
 
-    def dumpProxyToFile(self, chain, destinationFile=None, requiredTimeLeft=600):
+    def dumpProxyToFile(self, chain, destinationFile=None, requiredTimeLeft=600, includeToken=True):
         """Dump a proxy to a file. It's cached so multiple calls won't generate extra files
 
         :param X509Chain chain: proxy as a chain
@@ -442,7 +442,10 @@ class ProxyManagerClient(metaclass=DIRACSingleton.DIRACSingleton):
         filename = retVal["Value"]
         if not (result := chain.getDIRACGroup())["OK"]:
             return result
-        if not (result := addTokenToPEM(filename, result["Value"]))["OK"]:  # pylint: disable=unsubscriptable-object
+        if (
+            includeToken
+            and not (result := addTokenToPEM(filename, result["Value"]))["OK"]  # pylint: disable=unsubscriptable-object
+        ):
             return result
         self.__filesCache.add(cHash, chain.getRemainingSecs()["Value"], filename)
         return S_OK(filename)

DIRAC/FrameworkSystem/Client/SystemAdministratorClientCLI.py

@@ -11,7 +11,6 @@ import sys
 import time
 
 from DIRAC import gConfig, gLogger
-from DIRAC.ConfigurationSystem.Client.Helpers import CSGlobals
 from DIRAC.Core.Base.CLI import CLI, colorize
 from DIRAC.Core.Security.ProxyInfo import getProxyInfo
 from DIRAC.Core.Utilities import List
@@ -623,6 +622,11 @@ class SystemAdministratorClientCLI(CLI):
           install agent <system> <agent> [-m <ModuleName>] [-p <Option>=<Value>] [-p <Option>=<Value>] ...
           install executor <system> <executor> [-m <ModuleName>] [-p <Option>=<Value>] [-p <Option>=<Value>] ...
         """
+        result = getProxyInfo()
+        if not result["OK"]:
+            self._errMsg(result["Message"])
+        user = result["Value"]["username"]
+
         argss = args.split()
         hostSetup = extension = None
         if not argss:
@@ -673,7 +677,7 @@ class SystemAdministratorClientCLI(CLI):
 
                 if database != "InstalledComponentsDB":
                     result = MonitoringUtilities.monitorInstallation(
-                        "DB", system.replace("System", ""), database, cpu=cpu, hostname=hostname
+                        "DB", system.replace("System", ""), database, cpu=cpu, hostname=hostname, user=user
                     )
                     if not result["OK"]:
                         self._errMsg(result["Message"])
@@ -786,14 +790,14 @@ class SystemAdministratorClientCLI(CLI):
                     return
 
                 result = MonitoringUtilities.monitorInstallation(
-                    "DB", system, "InstalledComponentsDB", cpu=cpu, hostname=hostname
+                    "DB", system, "InstalledComponentsDB", cpu=cpu, hostname=hostname, user=user
                 )
                 if not result["OK"]:
                     self._errMsg(f"Error registering installation into database: {result['Message']}")
                     return
 
             result = MonitoringUtilities.monitorInstallation(
-                option, system, component, module, cpu=cpu, hostname=hostname
+                option, system, component, module, cpu=cpu, hostname=hostname, user=user
             )
             if not result["OK"]:
                 self._errMsg(f"Error registering installation into database: {result['Message']}")
@@ -820,6 +824,7 @@ class SystemAdministratorClientCLI(CLI):
         result = getProxyInfo()
         if not result["OK"]:
             self._errMsg(result["Message"])
+        user = result["Value"]["username"]
 
         option = argss[0]
         if option == "db":
@@ -842,7 +847,7 @@ class SystemAdministratorClientCLI(CLI):
                 self._errMsg(result["Message"])
                 return
             system = result["Value"][component]["System"]
-            result = MonitoringUtilities.monitorUninstallation(system, component, hostname=hostname, cpu=cpu)
+            result = MonitoringUtilities.monitorUninstallation(system, component, hostname=hostname, cpu=cpu, user=user)
             if not result["OK"]:
                 self._errMsg(result["Message"])
                 return
@@ -937,7 +942,7 @@ class SystemAdministratorClientCLI(CLI):
             else:
                 cpu = result["Value"]["CPUModel"]
             hostname = self.host
-            result = MonitoringUtilities.monitorUninstallation(system, component, hostname=hostname, cpu=cpu)
+            result = MonitoringUtilities.monitorUninstallation(system, component, hostname=hostname, cpu=cpu, user=user)
             if not result["OK"]:
                 return result
 

DIRAC/FrameworkSystem/ConfigTemplate.cfg

@@ -167,6 +167,7 @@ Services
       componentExists = authenticated
       getComponents = authenticated
       hostExists = authenticated
+      installationExists = authenticated
       getHosts = authenticated
       installationExists = authenticated
       getInstallations = authenticated
@@ -184,6 +185,7 @@ Services
       componentExists = authenticated
       getComponents = authenticated
       hostExists = authenticated
+      installationExists = authenticated
       getHosts = authenticated
       installationExists = authenticated
       getInstallations = authenticated

DIRAC/FrameworkSystem/DB/AuthDB.py

@@ -22,7 +22,7 @@ Model = declarative_base()
 
 class RefreshToken(Model):
     __tablename__ = "RefreshToken"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
     jti = Column(String(255), nullable=False, primary_key=True)
     issued_at = Column(Integer, nullable=False, default=0)
     access_token = Column(Text, nullable=False)
@@ -31,7 +31,7 @@ class RefreshToken(Model):
 
 class JWK(Model):
     __tablename__ = "JWK"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
     kid = Column(String(255), unique=True, primary_key=True, nullable=False)
     key = Column(Text, nullable=False)
     expires_at = Column(Integer, nullable=False, default=0)
@@ -39,7 +39,7 @@ class JWK(Model):
 
 class AuthSession(Model):
     __tablename__ = "AuthSession"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
     id = Column(String(255), unique=True, primary_key=True, nullable=False)
     uri = Column(String(255))
     state = Column(String(255))

DIRAC/FrameworkSystem/DB/InstalledComponentsDB.py

@@ -27,7 +27,7 @@ class Component(componentsBase):
     """
 
     __tablename__ = "Components"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
 
     componentID = Column("ComponentID", Integer, primary_key=True)
     system = Column("DIRACSystem", String(32), nullable=False)
@@ -87,7 +87,7 @@ class Host(componentsBase):
     """
 
     __tablename__ = "Hosts"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
 
     hostID = Column("HostID", Integer, primary_key=True)
     hostName = Column("HostName", String(32), nullable=False)
@@ -138,7 +138,7 @@ class InstalledComponent(componentsBase):
     """
 
     __tablename__ = "InstalledComponents"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
 
     componentID = Column("ComponentID", Integer, ForeignKey("Components.ComponentID"), primary_key=True)
     hostID = Column("HostID", Integer, ForeignKey("Hosts.HostID"), primary_key=True)
@@ -217,7 +217,7 @@ class HostLogging(componentsBase):
     """
 
     __tablename__ = "HostLogging"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
 
     hostName = Column("HostName", String(32), nullable=False, primary_key=True)
     # status

DIRAC/FrameworkSystem/DB/ProxyDB.py

@@ -9,6 +9,9 @@
     * ProxyDB_Log -- table with logs.
 """
 import textwrap
+from threading import Lock
+
+from cachetools import TTLCache, cachedmethod
 
 from DIRAC import S_ERROR, S_OK, gLogger
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
@@ -22,6 +25,10 @@ from DIRAC.Resources.ProxyProvider.ProxyProviderFactory import ProxyProviderFact
 
 DEFAULT_MAIL_FROM = "proxymanager@diracgrid.org"
 
+# Module-level cache for getProxyStrength method (shared across ProxyDB instances)
+_get_proxy_strength_cache = TTLCache(maxsize=1000, ttl=600)
+_get_proxy_strength_lock = Lock()
+
 
 class ProxyDB(DB):
     NOTIFICATION_TIMES = [2592000, 1296000]
@@ -395,6 +402,7 @@ class ProxyDB(DB):
             return S_ERROR(", ".join(errMsgs))
         return result
 
+    @cachedmethod(lambda self: _get_proxy_strength_cache, lock=lambda self: _get_proxy_strength_lock)
     def getProxyStrength(self, userDN, userGroup=None, vomsAttr=None):
         """Load the proxy in cache corresponding to the criteria, and check its strength
 
@@ -597,13 +605,13 @@ class ProxyDB(DB):
         :return: S_OK(dict)/S_ERROR() -- dict contain attribute and VOMS VO
         """
         if requiredVOMSAttribute:
-            return S_OK({"attribute": requiredVOMSAttribute, "VOMSVO": Registry.getVOMSVOForGroup(userGroup)})
+            return S_OK({"attribute": requiredVOMSAttribute, "VO": Registry.getVOForGroup(userGroup)})
 
         csVOMSMapping = Registry.getVOMSAttributeForGroup(userGroup)
         if not csVOMSMapping:
             return S_ERROR(f"No mapping defined for group {userGroup} in the CS")
 
-        return S_OK({"attribute": csVOMSMapping, "VOMSVO": Registry.getVOMSVOForGroup(userGroup)})
+        return S_OK({"attribute": csVOMSMapping, "VO": Registry.getVOForGroup(userGroup)})
 
     def getVOMSProxy(self, userDN, userGroup, requiredLifeTime=None, requestedVOMSAttr=None):
         """Get proxy string from the Proxy Repository for use with userDN
@@ -620,7 +628,7 @@ class ProxyDB(DB):
         if not retVal["OK"]:
             return retVal
         vomsAttr = retVal["Value"]["attribute"]
-        vomsVO = retVal["Value"]["VOMSVO"]
+        vomsVO = retVal["Value"]["VO"]
 
         # Look in the cache
         retVal = self.__getPemAndTimeLeft(userDN, userGroup, vomsAttr)

DIRAC/FrameworkSystem/DB/TokenDB.py

@@ -24,7 +24,7 @@ class Token(Model, OAuth2TokenMixin):
     """This class describes token fields"""
 
     __tablename__ = "Token"
-    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8"}
+    __table_args__ = {"mysql_engine": "InnoDB", "mysql_charset": "utf8mb4"}
     # access_token too large for varchar(255)
     # 767 bytes is the stated prefix limitation for InnoDB tables in MySQL version 5.6
     # https://stackoverflow.com/questions/1827063/mysql-error-key-specification-without-a-key-length

DIRAC/FrameworkSystem/Service/ProxyManagerHandler.py

@@ -1,11 +1,12 @@
-""" ProxyManager is the implementation of the ProxyManagement service in the DISET framework
+"""ProxyManager is the implementation of the ProxyManagement service in the DISET framework
 
-    .. literalinclude:: ../ConfigTemplate.cfg
-      :start-after: ##BEGIN ProxyManager:
-      :end-before: ##END
-      :dedent: 2
-      :caption: ProxyManager options
+.. literalinclude:: ../ConfigTemplate.cfg
+  :start-after: ##BEGIN ProxyManager:
+  :end-before: ##END
+  :dedent: 2
+  :caption: ProxyManager options
 """
+
 from DIRAC import S_ERROR, S_OK, gLogger
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.Core.DISET.RequestHandler import RequestHandler, getServiceOption
@@ -325,6 +326,7 @@ class ProxyManagerHandlerMixin:
             credDict["group"],
             set(credDict.get("groupProperties", []) + credDict.get("properties", [])),
             expires_minutes=credDict["secondsLeft"] // 60 + 1,
+            source="ProxyManager",
         )
 
 

DIRAC/FrameworkSystem/Utilities/MonitoringUtilities.py

@@ -6,10 +6,9 @@ import socket
 
 from DIRAC import S_OK
 from DIRAC.FrameworkSystem.Client.ComponentMonitoringClient import ComponentMonitoringClient
-from DIRAC.Core.Security.ProxyInfo import getProxyInfo
 
 
-def monitorInstallation(componentType, system, component, module=None, cpu=None, hostname=None):
+def monitorInstallation(componentType, system, component, module=None, cpu=None, hostname=None, user=None):
     """
     Register the installation of a component in the InstalledComponentsDB
     """
@@ -19,14 +18,6 @@ def monitorInstallation(componentType, system, component, module=None, cpu=None,
         module = component
 
     # Retrieve user installing the component
-    user = None
-    result = getProxyInfo()
-    if result["OK"]:
-        proxyInfo = result["Value"]
-        if "username" in proxyInfo:
-            user = proxyInfo["username"]
-    else:
-        return result
     if not user:
         user = "unknown"
 
@@ -61,21 +52,13 @@ def monitorInstallation(componentType, system, component, module=None, cpu=None,
     return result
 
 
-def monitorUninstallation(system, component, cpu=None, hostname=None):
+def monitorUninstallation(system, component, cpu=None, hostname=None, user=None):
     """
     Register the uninstallation of a component in the InstalledComponentsDB
     """
     monitoringClient = ComponentMonitoringClient()
 
     # Retrieve user uninstalling the component
-    user = None
-    result = getProxyInfo()
-    if result["OK"]:
-        proxyInfo = result["Value"]
-        if "username" in proxyInfo:
-            user = proxyInfo["username"]
-    else:
-        return result
     if not user:
         user = "unknown"
 

DIRAC/FrameworkSystem/Utilities/TokenManagementUtilities.py

@@ -10,11 +10,12 @@ DEFAULT_RT_EXPIRATION_TIME = 24 * 3600
 DEFAULT_AT_EXPIRATION_TIME = 1200
 
 
-def getIdProviderClient(userGroup: str, idProviderClientName: str = None):
+def getIdProviderClient(userGroup: str, idProviderClientName: str = None, client_name_prefix: str = ""):
     """Get an IdProvider client
 
     :param userGroup: group name
     :param idProviderClientName: name of an identity provider in the DIRAC CS
+    :param client_name_prefix: prefix of the client in the CS options
     """
     # Get IdProvider credentials from CS
     if not idProviderClientName and userGroup:
@@ -23,7 +24,7 @@ def getIdProviderClient(userGroup: str, idProviderClientName: str = None):
         return S_ERROR(f"The {userGroup} group belongs to the VO that is not tied to any Identity Provider.")
 
     # Prepare the client instance of the appropriate IdP
-    return IdProviderFactory().getIdProvider(idProviderClientName)
+    return IdProviderFactory().getIdProvider(idProviderClientName, client_name_prefix=client_name_prefix)
 
 
 def getCachedKey(

DIRAC/FrameworkSystem/Utilities/diracx.py

@@ -1,28 +1,40 @@
 import requests
 
-from cachetools import TTLCache, cached
+from cachetools import TTLCache, LRUCache, cached
 from cachetools.keys import hashkey
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import Any
+from collections.abc import Generator
 from DIRAC import gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
-
+from contextlib import contextmanager
 
 from diracx.core.preferences import DiracxPreferences
 
 from diracx.core.utils import write_credentials
 
 from diracx.core.models import TokenResponse
-from diracx.client import DiracClient
+
+try:
+    from diracx.client.sync import SyncDiracClient
+except ImportError:
+    # TODO: Remove this once diracx is tagged
+    from diracx.client import DiracClient as SyncDiracClient
 
 # How long tokens are kept
 DEFAULT_TOKEN_CACHE_TTL = 5 * 60
 DEFAULT_TOKEN_CACHE_SIZE = 1024
 
+legacy_exchange_session = requests.Session()
+
+
+def get_token(
+    username: str, group: str, dirac_properties: set[str], *, expires_minutes: int | None = None, source: str = ""
+):
+    """Do a legacy exchange to get a DiracX access_token+refresh_token
 
-def get_token(username: str, group: str, dirac_properties: set[str], *, expires_minutes: int | None = None):
-    """Do a legacy exchange to get a DiracX access_token+refresh_token"""
+    The source parameter only purpose is to appear in the URL on diracx logs"""
     diracxUrl = gConfig.getValue("/DiracX/URL")
     if not diracxUrl:
         raise ValueError("Missing mandatory /DiracX/URL configuration")
@@ -33,12 +45,13 @@ def get_token(username: str, group: str, dirac_properties: set[str], *, expires_
     vo = Registry.getVOForGroup(group)
     scopes = [f"vo:{vo}", f"group:{group}"] + [f"property:{prop}" for prop in dirac_properties]
 
-    r = requests.get(
+    r = legacy_exchange_session.get(
         f"{diracxUrl}/api/auth/legacy-exchange",
         params={
             "preferred_username": username,
             "scope": " ".join(scopes),
             "expires_minutes": expires_minutes,
+            "source": source,
         },
         headers={"Authorization": f"Bearer {apiKey}"},
         timeout=10,
@@ -51,21 +64,25 @@ def get_token(username: str, group: str, dirac_properties: set[str], *, expires_
 
 @cached(
     TTLCache(maxsize=DEFAULT_TOKEN_CACHE_SIZE, ttl=DEFAULT_TOKEN_CACHE_TTL),
-    key=lambda a, b, c: hashkey(a, b, *sorted(c)),
+    key=lambda a, b, c, **_: hashkey(a, b, *sorted(c)),
 )
-def _get_token_file(username: str, group: str, dirac_properties: set[str]) -> Path:
+def _get_token_file(username: str, group: str, dirac_properties: set[str], *, source: str = "") -> Path:
     """Write token to a temporary file and return the path to that file"""
-    data = get_token(username, group, dirac_properties)
+    data = get_token(username, group, dirac_properties, source=source)
     token_location = Path(NamedTemporaryFile().name)
     write_credentials(TokenResponse(**data), location=token_location)
     return token_location
 
 
-def TheImpersonator(credDict: dict[str, Any]) -> DiracClient:
+diracx_client_cache = LRUCache(maxsize=64)
+
+
+@contextmanager
+def TheImpersonator(credDict: dict[str, Any], *, source: str = "") -> Generator[SyncDiracClient, None, None]:
     """
     Client to be used by DIRAC server needing to impersonate
     a user for diracx.
-    It queries a token, places it in a file, and returns the `DiracClient`
+    It queries a token, places it in a file, and returns the `SyncDiracClient`
     class
 
     Use as a context manager
@@ -78,7 +95,12 @@ def TheImpersonator(credDict: dict[str, Any]) -> DiracClient:
         credDict["username"],
         credDict["group"],
         set(credDict.get("groupProperties", []) + credDict.get("properties", [])),
+        source=source,
     )
-    pref = DiracxPreferences(url=diracxUrl, credentials_path=token_location)
-
-    return DiracClient(diracx_preferences=pref)
+    client = diracx_client_cache.get(token_location)
+    if client is None:
+        pref = DiracxPreferences(url=diracxUrl, credentials_path=token_location)
+        client = SyncDiracClient(diracx_preferences=pref)
+        client.__enter__()
+        diracx_client_cache[token_location] = client
+    yield client