diff --git a/src/DIRAC/Core/Utilities/MySQL.py b/src/DIRAC/Core/Utilities/MySQL.py index c5ca54cf27b..2f835dd906d 100755 --- a/src/DIRAC/Core/Utilities/MySQL.py +++ b/src/DIRAC/Core/Utilities/MySQL.py @@ -942,11 +942,6 @@ def _to_value(self, param): """ return str(param[0]) - def _to_string(self, param): - """ - """ - return param[0].tostring() - def _getConnection(self, retries=MAXCONNECTRETRY): """ Return a new connection to the DB, diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py b/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py index 0b72d5fdbe2..910bffffbab 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py +++ b/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py @@ -10,7 +10,7 @@ # DIRAC Components from DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent import PilotStatusAgent -from DIRAC import gLogger +from DIRAC import gLogger, S_OK # Mock objects mockReply = MagicMock() @@ -37,6 +37,8 @@ def test_clearWaitingPilots(mocker): module_str = "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.buildCondition" mocker.patch(module_str, side_effect=mockNone) mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._query", side_effect=mockOK) + mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._escapeString", + lambda s, c: S_OK('"%s"' % s)) # To bypass "connection.escape_string" pilotStatusAgent = PilotStatusAgent() pilotStatusAgent._AgentModule__configDefaults = mockAM @@ -79,6 +81,8 @@ def test_handleOldPilots(mocker, mockReplyInput, expected): mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.__init__", side_effect=mockNone) mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.JobDB.__init__", side_effect=mockNone) mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._query", side_effect=mockOK) + mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._escapeString", + lambda s, c: S_OK('"%s"' % s)) # To bypass "connection.escape_string" pilotStatusAgent = PilotStatusAgent() pilotStatusAgent._AgentModule__configDefaults = mockAM diff --git a/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py b/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py index fcf9b84a2ec..c28233a87f6 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py @@ -149,13 +149,7 @@ def getAttributesForJobList(self, jobIDList, attrList=None): jobID = retValues[0] jobDict = {'JobID': jobID} # Make a dict from the list of attributes names and values - for name, value in zip(attr_tmp_list, retValues[1:]): - try: - value = value.tostring() - except Exception: - value = str(value) - jobDict[name] = value - retDict[int(jobID)] = jobDict + retDict[int(jobID)] = {k: v.decode() for k, v in zip(attr_tmp_list, retValues[1:])} return S_OK(retDict) except Exception as e: return S_ERROR('JobDB.getAttributesForJobList: Failed\n%s' % repr(e)) @@ -203,11 +197,7 @@ def getJobParameters(self, jobID, paramList=None): if result['OK']: if result['Value']: for res_jobID, res_name, res_value in result['Value']: - try: - res_value = res_value.tostring() - except Exception: - pass - resultDict.setdefault(int(res_jobID), {})[res_name] = res_value + resultDict.setdefault(int(res_jobID), {})[res_name] = res_value.decode() return S_OK(resultDict) # there's a slim chance that this is an empty dictionary else: @@ -219,11 +209,7 @@ def getJobParameters(self, jobID, paramList=None): return result for res_jobID, res_name, res_value in result['Value']: - try: - res_value = res_value.tostring() - except Exception: - pass - resultDict.setdefault(int(res_jobID), {})[res_name] = res_value + resultDict.setdefault(int(res_jobID), {})[res_name] = res_value.decode() return S_OK(resultDict) # there's a slim chance that this is an empty dictionary @@ -262,11 +248,7 @@ def getAtticJobParameters(self, jobID, paramList=None, rescheduleCounter=-1): if result['OK']: if result['Value']: for name, value, counter in result['Value']: - try: - value = value.tostring() - except Exception: - pass - resultDict.setdefault(counter, {})[name] = value + resultDict.setdefault(counter, {})[name] = value.decode() return S_OK(resultDict) else: @@ -406,8 +388,6 @@ def getJobOptParameters(self, jobID, paramList=None): return ret jobID = ret['Value'] - resultDict = {} - if paramList: paramNameList = [] for x in paramList: @@ -421,18 +401,9 @@ def getJobOptParameters(self, jobID, paramList=None): cmd = "SELECT Name, Value from OptimizerParameters WHERE JobID=%s" % jobID result = self._query(cmd) - if result['OK']: - if result['Value']: - for name, value in result['Value']: - try: - value = value.tostring() - except Exception: - pass - resultDict[name] = value - - return S_OK(resultDict) - else: + if not result['OK']: return S_ERROR('JobDB.getJobOptParameters: failed to retrieve parameters') + return S_OK({name: value.decode() for name, value in result.get('Value', {})}) ############################################################################# diff --git a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py index c1c2cf08940..15d15cc230c 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py @@ -213,10 +213,12 @@ def deletePilots(self, pilotIDs, conn=False): failed = [] + result = self._escapeValues(pilotIDs) + if not result['OK']: + return S_ERROR('Failed to remove pilot: %s' % result['Value']) + stringIDs = ','.join(result['Value']) for table in ['PilotOutput', 'JobToPilotMapping', 'PilotAgents']: - idString = ','.join([str(pid) for pid in pilotIDs]) - req = "DELETE FROM %s WHERE PilotID in ( %s )" % (table, idString) - result = self._update(req, conn=conn) + result = self._update("DELETE FROM %s WHERE PilotID in (%s)" % (table, stringIDs), conn=conn) if not result['OK']: failed.append(table) @@ -271,27 +273,19 @@ def getPilotInfo(self, pilotRef=False, parentId=False, conn=False, paramNames=[] parameters = ['PilotJobReference', 'OwnerDN', 'OwnerGroup', 'GridType', 'Broker', 'Status', 'DestinationSite', 'BenchMark', 'ParentID', 'OutputReady', 'AccountingSent', 'SubmissionTime', 'PilotID', 'LastUpdateTime', 'TaskQueueID', 'GridSite', 'PilotStamp', - 'Queue'] - if paramNames: - parameters = paramNames + 'Queue'] if not paramNames else paramNames cmd = "SELECT %s FROM PilotAgents" % ", ".join(parameters) condSQL = [] - if pilotRef: - if isinstance(pilotRef, list): - condSQL.append("PilotJobReference IN (%s)" % ",".join(['"%s"' % x for x in pilotRef])) - else: - condSQL.append("PilotJobReference = '%s'" % pilotRef) - if pilotID: - if isinstance(pilotID, list): - condSQL.append("PilotID IN (%s)" % ",".join(['%s' % x for x in pilotID])) - else: - condSQL.append("PilotID = '%s'" % pilotID) - if parentId: - if isinstance(parentId, list): - condSQL.append("ParentID IN (%s)" % ",".join(['%s' % x for x in parentId])) - else: - condSQL.append("ParentID = %s" % parentId) + for key, value in [('PilotJobReference', pilotRef), ('PilotID', pilotID), ('ParentID', parentId)]: + resList = [] + for v in value if isinstance(value, list) else [value] if value else []: + result = self._escapeString(v) + if not result['OK']: + return result + resList.append(result['Value']) + if resList: + condSQL.append("%s IN (%s)" % (key, ",".join(resList))) if condSQL: cmd = "%s WHERE %s" % (cmd, " AND ".join(condSQL)) diff --git a/tests/Integration/Resources/ProxyProvider/Test_DIRACCAProxyProvider.py b/tests/Integration/Resources/ProxyProvider/Test_DIRACCAProxyProvider.py index bc8b7299513..8eca9805ec0 100644 --- a/tests/Integration/Resources/ProxyProvider/Test_DIRACCAProxyProvider.py +++ b/tests/Integration/Resources/ProxyProvider/Test_DIRACCAProxyProvider.py @@ -37,9 +37,6 @@ } } } -""" % (os.path.join(certsPath, 'ca/ca.cert.pem'), os.path.join(certsPath, 'ca/ca.key.pem')) - -userCFG = """ Registry { Users @@ -61,7 +58,7 @@ } } } -""" +""" % (os.path.join(certsPath, 'ca/ca.cert.pem'), os.path.join(certsPath, 'ca/ca.key.pem')) class DIRACCAPPTest(unittest.TestCase): @@ -70,20 +67,15 @@ class DIRACCAPPTest(unittest.TestCase): @classmethod def setUpClass(cls): - pass + cfg = CFG() + cfg.loadFromBuffer(diracTestCACFG) + gConfig.loadCFG(cfg) @classmethod def tearDownClass(cls): pass def setUp(self): - - cfg = CFG() - cfg.loadFromBuffer(diracTestCACFG) - gConfig.loadCFG(cfg) - cfg.loadFromBuffer(userCFG) - gConfig.loadCFG(cfg) - result = ProxyProviderFactory().getProxyProvider('DIRAC_TEST_CA') self.assertTrue(result['OK'], '\n%s' % result.get('Message') or 'Error message is absent.') self.pp = result['Value'] diff --git a/tests/Integration/WorkloadManagementSystem/Test_PilotsClient.py b/tests/Integration/WorkloadManagementSystem/Test_PilotsClient.py index a1590370433..08602c68863 100644 --- a/tests/Integration/WorkloadManagementSystem/Test_PilotsClient.py +++ b/tests/Integration/WorkloadManagementSystem/Test_PilotsClient.py @@ -28,74 +28,78 @@ def test_PilotsDB(): pilots = PilotManagerClient() + # This will allow you to run the test again if necessary + for jobID in ['aPilot', 'anotherPilot']: + pilots.deletePilots(jobID) + res = pilots.addPilotTQReference(['aPilot'], 1, '/a/ownerDN', 'a/owner/Group') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.getCurrentPilotCounters({}) - assert res['OK'] is True, res['Message'] - assert res['Value'] == {'Submitted': 1}, res['Value'] + assert res['OK'], res['Message'] + assert res['Value'] == {'Submitted': 1} res = pilots.deletePilots('aPilot') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.getCurrentPilotCounters({}) - assert res['OK'] is True, res['Message'] - assert res['Value'] == {}, res['Value'] + assert res['OK'], res['Message'] + assert res['Value'] == {} res = pilots.addPilotTQReference(['anotherPilot'], 1, '/a/ownerDN', 'a/owner/Group') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.storePilotOutput('anotherPilot', 'This is an output', 'this is an error') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.getPilotOutput('anotherPilot') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] assert res['Value'] == {'OwnerDN': '/a/ownerDN', 'OwnerGroup': 'a/owner/Group', 'StdErr': 'this is an error', 'FileList': [], 'StdOut': 'This is an output'} res = pilots.getPilotInfo('anotherPilot') - assert res['OK'] is True, res['Message'] - assert res['Value']['anotherPilot']['AccountingSent'] == 'False', res['Value'] - assert res['Value']['anotherPilot']['PilotJobReference'] == 'anotherPilot', res['Value'] + assert res['OK'], res['Message'] + assert res['Value']['anotherPilot']['AccountingSent'] == 'False' + assert res['Value']['anotherPilot']['PilotJobReference'] == 'anotherPilot' res = pilots.selectPilots({}) - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.getPilotSummary('', '') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] assert res['Value']['Total']['Submitted'] == 1 res = pilots.getPilotMonitorWeb({}, [], 0, 100) - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] assert res['Value']['TotalRecords'] == 1 res = pilots.getPilotMonitorSelectors() - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] assert res['Value'] == {'GridType': ['DIRAC'], 'OwnerGroup': ['a/owner/Group'], 'DestinationSite': ['NotAssigned'], 'Broker': ['Unknown'], 'Status': ['Submitted'], 'OwnerDN': ['/a/ownerDN'], 'GridSite': ['Unknown'], - 'Owner': []}, res['Value'] + 'Owner': []} res = pilots.getPilotSummaryWeb({}, [], 0, 100) - assert res['OK'] is True, res['Message'] - assert res['Value']['TotalRecords'] == 1, res['Value'] + assert res['OK'], res['Message'] + assert res['Value']['TotalRecords'] == 1 res = pilots.setAccountingFlag('anotherPilot', 'True') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.setPilotStatus('anotherPilot', 'Running') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.getPilotInfo('anotherPilot') - assert res['OK'] is True, res['Message'] - assert res['Value']['anotherPilot']['AccountingSent'] == 'True', res['Value'] - assert res['Value']['anotherPilot']['Status'] == 'Running', res['Value'] + assert res['OK'], res['Message'] + assert res['Value']['anotherPilot']['AccountingSent'] == 'True' + assert res['Value']['anotherPilot']['Status'] == 'Running' res = pilots.setJobForPilot(123, 'anotherPilot') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.setPilotBenchmark('anotherPilot', 12.3) - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.countPilots({}) - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] # res = pilots.getCounters() # # getPilotStatistics res = pilots.deletePilots('anotherPilot') - assert res['OK'] is True, res['Message'] + assert res['OK'], res['Message'] res = pilots.getCurrentPilotCounters({}) - assert res['OK'] is True, res['Message'] - assert res['Value'] == {}, res['Value'] + assert res['OK'], res['Message'] + assert res['Value'] == {}