""" StorageManagementDB is a front end to the Stager Database.

    There are five tables in the StorageManagementDB: Tasks, CacheReplicas, TaskReplicas, StageRequests.

    The Tasks table is the place holder for the tasks that have requested files to be staged.
    These can be from different systems and have different associated call back methods.
    The CacheReplicas table keeps the information on all the CacheReplicas in the system.
    It maps all the file information LFN, PFN, SE to an assigned ReplicaID.
    The TaskReplicas table maps the TaskIDs from the Tasks table to the ReplicaID from the CacheReplicas table.
    The StageRequests table contains each of the prestage request IDs for each of the replicas.
"""

__RCSID__ = "$Id$"

import inspect
import threading

from DIRAC import gLogger, S_OK, S_ERROR
from DIRAC.Core.Base.DB import DB
from DIRAC.Core.Utilities.List import intListToString, stringListToString

# Stage Request are issue with a length of "PinLength"
# However, once Staged, the entry in the StageRequest will set a PinExpiryTime only for "PinLength" / THROTTLING_STEPS
# As PinExpiryTime arrives, StageRequest and their corresponding CacheReplicas entries are cleaned
# This allows to throttle the submission of Stage Requests up to a maximum of "DiskCacheTB" per "PinLength"
# After "PinLength" / THROTTLING_STEPS seconds, entries are removed, so new requests for the same replica will trigger
# a new Stage Request to the SE, and thus an update of the Pinning on the SE.
#
#  - "PinLength" is an Option of the StageRequest Agent that defaults to THROTTLING_TIME
#  - "DiskCacheTB" is an Option of the StorageElement that defaults to 1 (TB)
#
THROTTLING_TIME = 86400
THROTTLING_STEPS = 12


class StorageManagementDB(DB):

  def __init__(self, systemInstance='Default'):
    DB.__init__(self, 'StorageManagementDB', 'StorageManagement/StorageManagementDB')
    self.lock = threading.Lock()
    self.TASKPARAMS = [
        'TaskID',
        'Status',
        'Source',
        'SubmitTime',
        'LastUpdate',
        'CompleteTime',
        'CallBackMethod',
        'SourceTaskID']
    self.REPLICAPARAMS = [
        'ReplicaID',
        'Type',
        'Status',
        'SE',
        'LFN',
        'PFN',
        'Size',
        'FileChecksum',
        'GUID',
        'SubmitTime',
        'LastUpdate',
        'Reason',
        'Links']
    self.STAGEPARAMS = [
        'ReplicaID',
        'StageStatus',
        'RequestID',
        'StageRequestSubmitTime',
        'StageRequestCompletedTime',
        'PinLength',
        'PinExpiryTime']
    self.STATES = ['Failed', 'New', 'Waiting', 'Offline', 'StageSubmitted', 'Staged']

  def __getConnection(self, connection):
    if connection:
      return connection
    res = self._getConnection()
    if res['OK']:
      return res['Value']
    gLogger.warn("Failed to get MySQL connection", res['Message'])
    return connection

  def _caller(self):
    return inspect.stack()[2][3]
  ################################################################
  #
  # State machine management
  #

  def updateTaskStatus(self, taskIDs, newTaskStatus, connection=False):
    return self.__updateTaskStatus(taskIDs, newTaskStatus, connection=connection)

  def __updateTaskStatus(self, taskIDs, newTaskStatus, force=False, connection=False):
    connection = self.__getConnection(connection)
    if not taskIDs:
      return S_OK(taskIDs)
    if force:
      toUpdate = taskIDs
    else:
      res = self._checkTaskUpdate(taskIDs, newTaskStatus, connection=connection)
      if not res['OK']:
        return res
      toUpdate = res['Value']
    if not toUpdate:
      return S_OK(toUpdate)

    # reqSelect = "SELECT * FROM Tasks WHERE TaskID IN (%s) AND Status != '%s';" % ( intListToString( toUpdate ), newTaskStatus )
    reqSelect = "SELECT TaskID FROM Tasks WHERE TaskID IN (%s) AND Status != '%s';" % (
        intListToString(toUpdate), newTaskStatus)
    resSelect = self._query(reqSelect, connection)
    if not resSelect['OK']:
      gLogger.error("%s.%s_DB: problem retrieving record:" % (self._caller(), '__updateTaskStatus'),
                    "%s. %s" % (reqSelect, resSelect['Message']))

    req = "UPDATE Tasks SET Status='%s',LastUpdate=UTC_TIMESTAMP() WHERE TaskID IN (%s) AND Status != '%s';" % (
        newTaskStatus, intListToString(toUpdate), newTaskStatus)
    res = self._update(req, connection)
    if not res['OK']:
      return res

    taskIDs = []
    for record in resSelect['Value']:
      taskIDs.append(record[0])
      gLogger.verbose(
          "%s.%s_DB: to_update Tasks =  %s" %
          (self._caller(), '__updateTaskStatus', record))

    if len(taskIDs) > 0:
      reqSelect1 = "SELECT * FROM Tasks WHERE TaskID IN (%s);" % intListToString(taskIDs)
      resSelect1 = self._query(reqSelect1, connection)
      if not resSelect1["OK"]:
        gLogger.warn("%s.%s_DB: problem retrieving records: %s. %s" %
                     (self._caller(), '__updateTaskStatus', reqSelect1, resSelect1['Message']))
      else:
        for record in resSelect1['Value']:
          gLogger.verbose(
              "%s.%s_DB: updated Tasks = %s" %
              (self._caller(), '__updateTaskStatus', record))

    return S_OK(toUpdate)

  def _checkTaskUpdate(self, taskIDs, newTaskState, connection=False):
    connection = self.__getConnection(connection)
    if not taskIDs:
      return S_OK(taskIDs)
    # * -> Failed
    if newTaskState == 'Failed':
      oldTaskState = []
    # StageCompleting -> Done
    elif newTaskState == 'Done':
      oldTaskState = ['StageCompleting']
    # StageSubmitted -> StageCompleting
    elif newTaskState == 'StageCompleting':
      oldTaskState = ['StageSubmitted']
    # Waiting -> StageSubmitted
    elif newTaskState == 'StageSubmitted':
      oldTaskState = ['Waiting', 'Offline']
    # New -> Waiting
    elif newTaskState == 'Waiting':
      oldTaskState = ['New']
    elif newTaskState == 'Offline':
      oldTaskState = ['Waiting']
    else:
      return S_ERROR("Task status not recognized")
    if not oldTaskState:
      toUpdate = taskIDs
    else:
      req = "SELECT TaskID FROM Tasks WHERE Status in (%s) AND TaskID IN (%s)" % (
          stringListToString(oldTaskState), intListToString(taskIDs))
      res = self._query(req, connection)
      if not res['OK']:
        return res
      toUpdate = [row[0] for row in res['Value']]
    return S_OK(toUpdate)

  def updateReplicaStatus(self, replicaIDs, newReplicaStatus, connection=False):
    connection = self.__getConnection(connection)
    if not replicaIDs:
      return S_OK(replicaIDs)
    res = self._checkReplicaUpdate(replicaIDs, newReplicaStatus)
    if not res['OK']:
      return res
    toUpdate = res['Value']
    if not toUpdate:
      return S_OK(toUpdate)
    # reqSelect = "SELECT * FROM CacheReplicas WHERE ReplicaID IN (%s) AND Status != '%s';" % ( intListToString( toUpdate ), newReplicaStatus )
    reqSelect = "SELECT ReplicaID FROM CacheReplicas WHERE ReplicaID IN (%s) AND Status != '%s';" % (
        intListToString(toUpdate), newReplicaStatus)
    resSelect = self._query(reqSelect, connection)
    if not resSelect['OK']:
      gLogger.error(
          "%s.%s_DB: problem retrieving record:" %
          (self._caller(), 'updateReplicaStatus'), "%s. %s" %
          (reqSelect, resSelect['Message']))

    req = "UPDATE CacheReplicas SET Status='%s',LastUpdate=UTC_TIMESTAMP() WHERE ReplicaID IN (%s) AND Status != '%s';" % (
        newReplicaStatus, intListToString(toUpdate), newReplicaStatus)
    res = self._update(req, connection)
    if not res['OK']:
      return res

    replicaIDs = []
    for record in resSelect['Value']:
      replicaIDs.append(record[0])
      gLogger.verbose("%s.%s_DB: to_update CacheReplicas =  %s" %
                      (self._caller(), 'updateReplicaStatus', record))
    if len(replicaIDs) > 0:
      reqSelect1 = "SELECT * FROM CacheReplicas WHERE ReplicaID IN (%s);" % intListToString(
          replicaIDs)
      resSelect1 = self._query(reqSelect1, connection)
      if not resSelect1['OK']:
        gLogger.warn("%s.%s_DB: problem retrieving records: %s. %s" %
                     (self._caller(), 'updateReplicaStatus', reqSelect1, resSelect1['Message']))
      else:
        for record in resSelect1['Value']:
          gLogger.verbose(
              "%s.%s_DB: updated CacheReplicas = %s" %
              (self._caller(), 'updateReplicaStatus', record))

    res = self._updateTasksForReplica(replicaIDs, connection=connection)
    if not res['OK']:
      return res
    return S_OK(toUpdate)

  def _updateTasksForReplica(self, replicaIDs, connection=False):
    tasksInStatus = {}
    for state in self.STATES:
      tasksInStatus[state] = []

    req = "SELECT T.TaskID,T.Status FROM Tasks AS T, TaskReplicas AS R WHERE R.ReplicaID IN ( %s ) AND R.TaskID = T.TaskID GROUP BY T.TaskID, T.Status;" % intListToString(
        replicaIDs)
    res = self._query(req, connection)
    if not res['OK']:
      return res

    for taskId, status in res['Value']:
      subreq = "SELECT DISTINCT(C.Status) FROM TaskReplicas AS R, CacheReplicas AS C WHERE R.TaskID=%s AND R.ReplicaID = C.ReplicaID;" % taskId
      subres = self._query(subreq, connection)
      if not subres['OK']:
        return subres

      cacheStatesForTask = [row[0] for row in subres['Value']]
      if not cacheStatesForTask:
        tasksInStatus['Failed'].append(taskId)
        continue

      wrongState = False
      for state in cacheStatesForTask:
        if state not in self.STATES:
          wrongState = True
          break
      if wrongState:
        tasksInStatus['Failed'].append(taskId)
        continue
      for state in self.STATES:
        if state in cacheStatesForTask:
          if status != state:
            tasksInStatus[state].append(taskId)
          break

    for newStatus in tasksInStatus:
      if tasksInStatus[newStatus]:
        res = self.__updateTaskStatus(
            tasksInStatus[newStatus],
            newStatus,
            True,
            connection=connection)
        if not res['OK']:
          gLogger.warn("Failed to update task associated to replicas", res['Message'])
          # return res
    return S_OK(tasksInStatus)

  def getAssociatedReplicas(self, replicaIDs):
    """ Retrieve the list of Replicas that belong to the same Tasks as the provided list
    """
    res = self._getReplicaIDTasks(replicaIDs)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.getAssociatedReplicas: Failed to get Tasks.',
          res['Message'])
      return res
    taskIDs = res['Value']

    return self.getCacheReplicas({'TaskID': taskIDs})

  def _checkReplicaUpdate(self, replicaIDs, newReplicaState, connection=False):
    connection = self.__getConnection(connection)
    if not replicaIDs:
      return S_OK(replicaIDs)
    # * -> Failed
    if newReplicaState == 'Failed':
      oldReplicaState = []
    # New -> Waiting
    elif newReplicaState == 'Waiting':
      oldReplicaState = ['New']
    # Waiting -> StageSubmitted
    elif newReplicaState == 'StageSubmitted':
      oldReplicaState = ['Waiting', 'Offline']
    # StageSubmitted -> Staged
    elif newReplicaState == 'Staged':
      oldReplicaState = ['StageSubmitted']
    elif newReplicaState == 'Offline':
      oldReplicaState = ['Waiting']
    else:
      return S_ERROR("Replica status not recognized")
    if not oldReplicaState:
      toUpdate = replicaIDs
    else:
      req = "SELECT ReplicaID FROM CacheReplicas WHERE Status IN (%s) AND ReplicaID IN (%s)" % (
          stringListToString(oldReplicaState), intListToString(replicaIDs))
      res = self._query(req, connection)
      if not res['OK']:
        return res
      toUpdate = [row[0] for row in res['Value']]
    return S_OK(toUpdate)

  def __getTaskStateFromReplicaState(self, replicaState):
    # For the moment the task state just references to the replicaState
    return replicaState

  def updateStageRequestStatus(self, replicaIDs, newStageStatus, connection=False):
    connection = self.__getConnection(connection)
    if not replicaIDs:
      return S_OK(replicaIDs)
    res = self._checkStageUpdate(replicaIDs, newStageStatus, connection=connection)
    if not res['OK']:
      return res
    toUpdate = res['Value']
    if not toUpdate:
      return S_OK(toUpdate)
    # reqSelect = "Select * FROM CacheReplicas WHERE ReplicaID IN (%s) AND Status != '%s';" % ( intListToString( toUpdate ), newStageStatus )
    reqSelect = "Select ReplicaID FROM CacheReplicas WHERE ReplicaID IN (%s) AND Status != '%s';" % (
        intListToString(toUpdate), newStageStatus)
    resSelect = self._query(reqSelect, connection)
    if not resSelect['OK']:
      gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                   (self._caller(), 'updateStageRequestStatus', reqSelect, resSelect['Message']))

    req = "UPDATE CacheReplicas SET Status='%s',LastUpdate=UTC_TIMESTAMP() WHERE ReplicaID IN (%s) AND Status != '%s';" % (
        newStageStatus, intListToString(toUpdate), newStageStatus)
    res = self._update(req, connection)
    if not res['OK']:
      return res

    replicaIDs = []
    for record in resSelect['Value']:
      replicaIDs.append(record[0])
      gLogger.verbose("%s.%s_DB: to_update CacheReplicas =  %s" %
                      (self._caller(), 'updateStageRequestStatus', record))

    reqSelect1 = "SELECT * FROM CacheReplicas WHERE ReplicaID IN (%s);" % intListToString(
        replicaIDs)
    resSelect1 = self._query(reqSelect1, connection)
    if not resSelect1['OK']:
      gLogger.warn("%s.%s_DB: problem retrieving records: %s. %s" %
                   (self._caller(), 'updateStageRequestStatus', reqSelect1, resSelect1['Message']))
    else:
      for record in resSelect1['Value']:
        gLogger.verbose(
            "%s.%s_DB: updated CacheReplicas = %s" %
            (self._caller(), 'updateStageRequestStatus', record))

    # Now update the replicas associated to the replicaIDs
    newReplicaStatus = self.__getReplicaStateFromStageState(newStageStatus)
    res = self.updateReplicaStatus(toUpdate, newReplicaStatus, connection=connection)
    if not res['OK']:
      gLogger.warn("Failed to update cache replicas associated to stage requests", res['Message'])
    return S_OK(toUpdate)

  def _checkStageUpdate(self, replicaIDs, newStageState, connection=False):
    connection = self.__getConnection(connection)
    if not replicaIDs:
      return S_OK(replicaIDs)
    # * -> Failed
    if newStageState == 'Failed':
      oldStageState = []
    elif newStageState == 'Staged':
      oldStageState = ['StageSubmitted']
    else:
      return S_ERROR("StageRequest status not recognized")
    if not oldStageState:
      toUpdate = replicaIDs
    else:
      req = "SELECT ReplicaID FROM StageRequests WHERE StageStatus = '%s' AND ReplicaID IN (%s)" % (
          oldStageState, intListToString(replicaIDs))
      res = self._query(req, connection)
      if not res['OK']:
        return res
      toUpdate = [row[0] for row in res['Value']]
    return S_OK(toUpdate)

  def __getReplicaStateFromStageState(self, stageState):
    # For the moment the replica state just references to the stage state
    return stageState

  #
  #                               End of state machine management
  #
  ################################################################

  ################################################################
  #
  # Monitoring of stage tasks
  #
  def getTaskStatus(self, taskID, connection=False):
    """ Obtain the task status from the Tasks table. """
    connection = self.__getConnection(connection)
    res = self.getTaskInfo(taskID, connection=connection)
    if not res['OK']:
      return res
    taskInfo = res['Value'][taskID]
    return S_OK(taskInfo['Status'])

  def getTaskInfo(self, taskID, connection=False):
    """ Obtain all the information from the Tasks table for a supplied task. """
    connection = self.__getConnection(connection)
    req = "SELECT TaskID,Status,Source,SubmitTime,CompleteTime,CallBackMethod,SourceTaskID from Tasks WHERE TaskID IN (%s);" % intListToString(
        taskID)
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.getTaskInfo: Failed to get task information.',
          res['Message'])
      return res
    resDict = {}
    for taskID, status, source, submitTime, completeTime, callBackMethod, sourceTaskID in res[
            'Value']:
      resDict[sourceTaskID] = {
          'Status': status,
          'Source': source,
          'SubmitTime': submitTime,
          'CompleteTime': completeTime,
          'CallBackMethod': callBackMethod,
          'SourceTaskID': sourceTaskID}
    if not resDict:
      gLogger.error('StorageManagementDB.getTaskInfo: The supplied task did not exist', taskID)
      return S_ERROR('The supplied task %s did not exist' % taskID)
    return S_OK(resDict)

  def _getTaskIDForJob(self, jobID, connection=False):
    # Stager taskID is retrieved from the source DIRAC jobID
    connection = self.__getConnection(connection)
    req = "SELECT TaskID from Tasks WHERE SourceTaskID=%s;" % int(jobID)
    res = self._query(req)
    if not res['OK']:
      gLogger.error("%s.%s_DB: problem retrieving record:" % (self._caller(), '_getTaskIDForJob'),
                    "%s. %s" % (req, res['Message']))
      return S_ERROR('The supplied JobID does not exist!')
    taskID = [row[0] for row in res['Value']]
    return S_OK(taskID)

  def getTaskSummary(self, jobID, connection=False):
    """ Obtain the task summary from the database. """
    connection = self.__getConnection(connection)
    res = self._getTaskIDForJob(jobID, connection=connection)
    if not res['OK']:
      return res
    if res['Value']:
      taskID = res['Value']
    else:
      return S_OK()
    res = self.getTaskInfo(taskID, connection=connection)
    if not res['OK']:
      return res
    taskInfo = res['Value']
    req = "SELECT R.LFN,R.SE,R.PFN,R.Size,R.Status,R.LastUpdate,R.Reason FROM CacheReplicas AS R, TaskReplicas AS TR WHERE TR.TaskID in (%s) AND TR.ReplicaID=R.ReplicaID;" % intListToString(
        taskID)
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.getTaskSummary: Failed to get Replica summary for task.',
          res['Message'])
      return res
    replicaInfo = {}
    for lfn, storageElement, pfn, fileSize, status, lastupdate, reason in res['Value']:
      replicaInfo[lfn] = {'StorageElement': storageElement, 'PFN': pfn, 'FileSize': fileSize,
                          'Status': status, 'LastUpdate': lastupdate, 'Reason': reason}
    resDict = {'TaskInfo': taskInfo, 'ReplicaInfo': replicaInfo}
    return S_OK(resDict)

  def getTasks(
          self,
          condDict={},
          older=None,
          newer=None,
          timeStamp='SubmitTime',
          orderAttribute=None,
          limit=None,
          connection=False):
    """ Get stage requests for the supplied selection with support for web standard structure """
    connection = self.__getConnection(connection)
    req = "SELECT %s FROM Tasks" % (intListToString(self.TASKPARAMS))
    if condDict or older or newer:
      if 'ReplicaID' in condDict:
        replicaIDs = condDict.pop('ReplicaID')
        if not isinstance(replicaIDs, (list, tuple)):
          replicaIDs = [replicaIDs]
        res = self._getReplicaIDTasks(replicaIDs, connection=connection)
        if not res['OK']:
          return res
        condDict['TaskID'] = res['Value']
      req = "%s %s" % (req,
                       self.buildCondition(
                           condDict,
                           older,
                           newer,
                           timeStamp,
                           orderAttribute,
                           limit))
    res = self._query(req, connection)
    if not res['OK']:
      return res
    tasks = res['Value']
    resultDict = {}
    for row in tasks:
      resultDict[row[0]] = dict(zip(self.TASKPARAMS[1:], row[1:]))
    result = S_OK(resultDict)
    result['Records'] = tasks
    result['ParameterNames'] = self.TASKPARAMS
    return result

  def getCacheReplicas(
          self,
          condDict={},
          older=None,
          newer=None,
          timeStamp='LastUpdate',
          orderAttribute=None,
          limit=None,
          connection=False):
    """ Get cache replicas for the supplied selection with support for the web standard structure """
    connection = self.__getConnection(connection)
    req = "SELECT %s FROM CacheReplicas" % (intListToString(self.REPLICAPARAMS))
    if condDict or older or newer:
      if 'TaskID' in condDict:
        taskIDs = condDict.pop('TaskID')
        if not isinstance(taskIDs, (list, tuple)):
          taskIDs = [taskIDs]
        res = self._getTaskReplicaIDs(taskIDs, connection=connection)
        if not res['OK']:
          return res
        if res['Value']:
          condDict['ReplicaID'] = res['Value']
        else:
          condDict['ReplicaID'] = [-1]
      # BUG: limit is ignored unless there is a nonempty condition dictionary OR
      # older OR newer is nonemtpy
      req = "%s %s" % (req,
                       self.buildCondition(
                           condDict,
                           older,
                           newer,
                           timeStamp,
                           orderAttribute,
                           limit))
    res = self._query(req, connection)
    if not res['OK']:
      return res
    cacheReplicas = res['Value']
    resultDict = {}
    for row in cacheReplicas:
      resultDict[row[0]] = dict(zip(self.REPLICAPARAMS[1:], row[1:]))
    result = S_OK(resultDict)
    result['Records'] = cacheReplicas
    result['ParameterNames'] = self.REPLICAPARAMS
    return result

  def getStageRequests(
          self,
          condDict={},
          older=None,
          newer=None,
          timeStamp='StageRequestSubmitTime',
          orderAttribute=None,
          limit=None,
          connection=False):
    """ Get stage requests for the supplied selection with support for web standard structure """
    connection = self.__getConnection(connection)
    req = "SELECT %s FROM StageRequests" % (intListToString(self.STAGEPARAMS))
    if condDict or older or newer:
      if 'TaskID' in condDict:
        taskIDs = condDict.pop('TaskID')
        if not isinstance(taskIDs, (list, tuple)):
          taskIDs = [taskIDs]
        res = self._getTaskReplicaIDs(taskIDs, connection=connection)
        if not res['OK']:
          return res
        if res['Value']:
          condDict['ReplicaID'] = res['Value']
        else:
          condDict['ReplicaID'] = [-1]
      req = "%s %s" % (req,
                       self.buildCondition(
                           condDict,
                           older,
                           newer,
                           timeStamp,
                           orderAttribute,
                           limit))
    res = self._query(req, connection)
    if not res['OK']:
      return res
    stageRequests = res['Value']
    resultDict = {}
    for row in stageRequests:
      resultDict[row[0]] = dict(zip(self.STAGEPARAMS[1:], row[1:]))
    result = S_OK(resultDict)
    result['Records'] = stageRequests
    result['ParameterNames'] = self.STAGEPARAMS
    return result

  def _getTaskReplicaIDs(self, taskIDs, connection=False):
    if not taskIDs:
      return S_OK([])
    req = "SELECT DISTINCT(ReplicaID) FROM TaskReplicas WHERE TaskID IN (%s);" % intListToString(
        taskIDs)
    res = self._query(req, connection)
    if not res['OK']:
      return res
    replicaIDs = [row[0] for row in res['Value']]
    return S_OK(replicaIDs)

  def _getReplicaIDTasks(self, replicaIDs, connection=False):
    if not replicaIDs:
      return S_OK([])
    req = "SELECT DISTINCT(TaskID) FROM TaskReplicas WHERE ReplicaID IN (%s);" % intListToString(
        replicaIDs)
    res = self._query(req, connection)
    if not res['OK']:
      return res
    taskIDs = [row[0] for row in res['Value']]

    return S_OK(taskIDs)

  #
  #                               End of monitoring of stage tasks
  #
  ################################################################

  ####################################################################
  #
  # Submission of stage requests
  #

  def setRequest(self, lfnDict, source, callbackMethod, sourceTaskID, connection=False):
    """ This method populates the StorageManagementDB Tasks table with the requested files. """
    connection = self.__getConnection(connection)
    if not lfnDict:
      return S_ERROR("No files supplied in request")
    # The first step is to create the task in the Tasks table
    res = self._createTask(source, callbackMethod, sourceTaskID, connection=connection)
    if not res['OK']:
      return res
    taskID = res['Value']
    # Get the Replicas which already exist in the CacheReplicas table
    allReplicaIDs = []
    taskStates = []
    for se, lfns in lfnDict.iteritems():
      if isinstance(lfns, basestring):
        lfns = [lfns]
      res = self._getExistingReplicas(se, lfns, connection=connection)
      if not res['OK']:
        return res
      existingReplicas = res['Value']
      # Insert the CacheReplicas that do not already exist
      for lfn in lfns:
        if lfn in existingReplicas:
          gLogger.verbose(
              'StorageManagementDB.setRequest: Replica already exists in CacheReplicas table %s @ %s' %
              (lfn, se))
          existingFileState = existingReplicas[lfn][1]
          taskState = self.__getTaskStateFromReplicaState(existingFileState)
        else:
          res = self._insertReplicaInformation(lfn, se, 'Stage', connection=connection)
          if not res['OK']:
            self._cleanTask(taskID, connection=connection)
            return res

          existingReplicas[lfn] = (res['Value'], 'New')
          newFileState = existingReplicas[lfn][1]
          taskState = self.__getTaskStateFromReplicaState(newFileState)
        if taskState not in taskStates:
          taskStates.append(taskState)

      allReplicaIDs.extend(existingReplicas.values())
    # Insert all the replicas into the TaskReplicas table
    res = self._insertTaskReplicaInformation(taskID, allReplicaIDs, connection=connection)
    if not res['OK']:
      self._cleanTask(taskID, connection=connection)
      return res
    # Check whether the the task status is Done based on the existing file states
    # If all the files for a particular Task are 'Staged', update the Task
    if taskStates == ['Staged']:
      # so if the tasks are for LFNs from the lfns dictionary, which are already staged,
      # they immediately change state New->Done. Fixed it to translate such tasks to 'Staged' state
      self.__updateTaskStatus([taskID], 'Staged', True, connection=connection)
    if 'Failed' in taskStates:
      self.__updateTaskStatus([taskID], 'Failed', True, connection=connection)
    return S_OK(taskID)

  def _cleanTask(self, taskID, connection=False):
    """ Remove a task and any related information """
    connection = self.__getConnection(connection)
    self.removeTasks([taskID], connection=connection)
    self.removeUnlinkedReplicas(connection=connection)

  def _createTask(self, source, callbackMethod, sourceTaskID, connection=False):
    """ Enter the task details into the Tasks table """
    connection = self.__getConnection(connection)
    req = "INSERT INTO Tasks (Source,SubmitTime,CallBackMethod,SourceTaskID) VALUES ('%s',UTC_TIMESTAMP(),'%s','%s');" % (
        source, callbackMethod, sourceTaskID)
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error("StorageManagementDB._createTask: Failed to create task.", res['Message'])
      return res
    # gLogger.info( "%s_DB:%s" % ('_createTask',req))
    taskID = res['lastRowId']
    reqSelect = "SELECT * FROM Tasks WHERE TaskID = %s;" % (taskID)
    resSelect = self._query(reqSelect, connection)
    if not resSelect['OK']:
      gLogger.info("%s.%s_DB: problem retrieving record: %s. %s" %
                   (self._caller(), '_createTask', reqSelect, resSelect['Message']))
    else:
      gLogger.verbose("%s.%s_DB: inserted Tasks = %s" %
                      (self._caller(), '_createTask', resSelect['Value'][0]))

    # gLogger.info("StorageManagementDB._createTask: Created task with ('%s','%s','%s') and obtained TaskID %s" % (source,callbackMethod,sourceTaskID,taskID))
    return S_OK(taskID)

  def _getExistingReplicas(self, storageElement, lfns, connection=False):
    """ Obtains the ReplicasIDs for the replicas already entered in the CacheReplicas table """
    connection = self.__getConnection(connection)
    req = "SELECT ReplicaID,LFN,Status FROM CacheReplicas WHERE SE = '%s' AND LFN IN (%s);" % (
        storageElement, stringListToString(lfns))
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB._getExistingReplicas: Failed to get existing replicas.',
          res['Message'])
      return res
    existingReplicas = {}
    for replicaID, lfn, status in res['Value']:
      existingReplicas[lfn] = (replicaID, status)
    return S_OK(existingReplicas)

  def _insertReplicaInformation(self, lfn, storageElement, rType, connection=False):
    """ Enter the replica into the CacheReplicas table """
    connection = self.__getConnection(connection)
    req = "INSERT INTO CacheReplicas (Type,SE,LFN,PFN,Size,FileChecksum,GUID,SubmitTime,LastUpdate) VALUES ('%s','%s','%s','',0,'','',UTC_TIMESTAMP(),UTC_TIMESTAMP());" % (
        rType, storageElement, lfn)
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error(
          "_insertReplicaInformation: Failed to insert to CacheReplicas table.",
          res['Message'])
      return res
    # gLogger.info( "%s_DB:%s" % ('_insertReplicaInformation',req))

    replicaID = res['lastRowId']
    reqSelect = "SELECT * FROM CacheReplicas WHERE ReplicaID = %s;" % (replicaID)
    resSelect = self._query(reqSelect, connection)
    if not resSelect['OK']:
      gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                   (self._caller(), '_insertReplicaInformation', reqSelect, resSelect['Message']))
    else:
      gLogger.verbose("%s.%s_DB: inserted CacheReplicas = %s" %
                      (self._caller(), '_insertReplicaInformation', resSelect['Value'][0]))
    # gLogger.verbose("_insertReplicaInformation: Inserted Replica ('%s','%s') and obtained ReplicaID %s" % (lfn,storageElement,replicaID))
    return S_OK(replicaID)

  def _insertTaskReplicaInformation(self, taskID, replicaIDs, connection=False):
    """ Enter the replicas into TaskReplicas table """
    connection = self.__getConnection(connection)
    req = "INSERT INTO TaskReplicas (TaskID,ReplicaID) VALUES "
    for replicaID, _status in replicaIDs:
      replicaString = "(%s,%s)," % (taskID, replicaID)
      req = "%s %s" % (req, replicaString)
    req = req.rstrip(',')
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB._insertTaskReplicaInformation: Failed to insert to TaskReplicas table.',
          res['Message'])
      return res
    # gLogger.info( "%s_DB:%s" % ('_insertTaskReplicaInformation',req))
    gLogger.verbose(
        "StorageManagementDB._insertTaskReplicaInformation: Successfully added %s CacheReplicas to Task %s." %
        (res['Value'], taskID))
    return S_OK()

  #
  #                               End of insertion methods
  #
  ################################################################

  ####################################################################

  def getStagedReplicas(self, connection=False):
    connection = self.__getConnection(connection)
    req = "SELECT TR.TaskID, R.Status, COUNT(*) from TaskReplicas as TR, CacheReplicas as R where TR.ReplicaID=R.ReplicaID GROUP BY TR.TaskID,R.Status;"
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.getStagedReplicas: Failed to get eligible TaskReplicas',
          res['Message'])
      return res
    goodTasks = []
    for taskID, status, _count in res['Value']:
      if taskID in goodTasks:
        continue
      elif status in ('Staged', 'StageSubmitted'):
        goodTasks.append(taskID)
    return self.getCacheReplicas({'Status': 'Staged', 'TaskID': goodTasks}, connection=connection)

  def getWaitingReplicas(self, connection=False):
    connection = self.__getConnection(connection)
    req = "SELECT TR.TaskID, R.Status, COUNT(*) from TaskReplicas as TR, CacheReplicas as R where TR.ReplicaID=R.ReplicaID GROUP BY TR.TaskID,R.Status;"
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.getWaitingReplicas: Failed to get eligible TaskReplicas',
          res['Message'])
      return res
    badTasks = []
    goodTasks = []
    for taskID, status, _count in res['Value']:
      if taskID in badTasks:
        continue
      elif status in ('New', 'Failed'):
        badTasks.append(taskID)
      elif status == 'Waiting':
        goodTasks.append(taskID)
    return self.getCacheReplicas({'Status': 'Waiting', 'TaskID': goodTasks}, connection=connection)

  ####################################################################

  def getOfflineReplicas(self, connection=False):
    connection = self.__getConnection(connection)
    req = "SELECT TR.TaskID, R.Status, COUNT(*) from TaskReplicas as TR, CacheReplicas as R where TR.ReplicaID=R.ReplicaID GROUP BY TR.TaskID,R.Status;"
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.getOfflineReplicas: Failed to get eligible TaskReplicas',
          res['Message'])
      return res
    badTasks = []
    goodTasks = []
    for taskID, status, _count in res['Value']:
      if taskID in badTasks:
        continue
      elif status in ('New', 'Failed'):
        badTasks.append(taskID)
      elif status == 'Offline':
        goodTasks.append(taskID)
    return self.getCacheReplicas({'Status': 'Offline', 'TaskID': goodTasks}, connection=connection)

  ####################################################################

  def getTasksWithStatus(self, status):
    """ This method retrieves the TaskID from the Tasks table with the supplied Status. """
    req = "SELECT TaskID,Source,CallBackMethod,SourceTaskID from Tasks WHERE Status = '%s';" % status
    res = self._query(req)
    if not res['OK']:
      return res
    taskIDs = {}
    for taskID, source, callback, sourceTask in res['Value']:
      taskIDs[taskID] = (source, callback, sourceTask)
    return S_OK(taskIDs)

  ####################################################################
  #
  # The state transition of the CacheReplicas from *->Failed
  #

  def updateReplicaFailure(self, terminalReplicaIDs):
    """ This method sets the status to Failure with the failure reason for the supplied Replicas. """
    res = self.updateReplicaStatus(terminalReplicaIDs.keys(), 'Failed')
    if not res['OK']:
      return res
    updated = res['Value']
    if not updated:
      return S_OK(updated)
    for replicaID in updated:
      reqSelect = "Select * FROM CacheReplicas WHERE ReplicaID = %d" % (replicaID)
      resSelect = self._query(reqSelect)
      if not resSelect['OK']:
        gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                     (self._caller(), 'updateReplicaFailure', reqSelect, resSelect['Message']))

      req = "UPDATE CacheReplicas SET Reason = '%s' WHERE ReplicaID = %d" % (
          terminalReplicaIDs[replicaID], replicaID)
      res = self._update(req)
      if not res['OK']:
        gLogger.error(
            'StorageManagementDB.updateReplicaFailure: Failed to update replica fail reason.',
            res['Message'])
        return res

      replicaIDs = []
      for record in resSelect['Value']:
        replicaIDs.append(record[0])
        gLogger.verbose(
            "%s.%s_DB: to_update CacheReplicas =  %s" %
            (self._caller(), 'updateReplicaFailure', record))

      reqSelect1 = "SELECT * FROM CacheReplicas WHERE ReplicaID IN (%s);" % intListToString(
          replicaIDs)
      resSelect1 = self._query(reqSelect1)
      if not resSelect1['OK']:
        gLogger.warn("%s.%s_DB: problem retrieving records: %s. %s" %
                     (self._caller(), 'updateReplicaFailure', reqSelect1, resSelect1['Message']))
      else:
        for record in resSelect1['Value']:
          gLogger.verbose(
              "%s.%s_DB: updated CacheReplicas = %s" %
              (self._caller(), 'updateReplicaFailure', record))

    return S_OK(updated)

  ####################################################################
  #
  # The state transition of the CacheReplicas from New->Waiting
  #

  def updateReplicaInformation(self, replicaTuples):
    """ This method set the replica size information and pfn for the requested storage element.  """
    for replicaID, pfn, size in replicaTuples:
      # reqSelect = "SELECT * FROM CacheReplicas WHERE ReplicaID = %s and Status != 'Cancelled';" % ( replicaID )
      reqSelect = "SELECT ReplicaID FROM CacheReplicas WHERE ReplicaID = %s and Status != 'Cancelled';" % (
          replicaID)
      resSelect = self._query(reqSelect)
      if not resSelect['OK']:
        gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                     (self._caller(), 'updateReplicaInformation', reqSelect, resSelect['Message']))

      req = "UPDATE CacheReplicas SET PFN = '%s', Size = %s, Status = 'Waiting' WHERE ReplicaID = %s and Status != 'Cancelled';" % (
          pfn, size, replicaID)
      res = self._update(req)
      if not res['OK']:
        gLogger.error(
            'StagerDB.updateReplicaInformation: Failed to insert replica information.',
            res['Message'])

      replicaIDs = []
      for record in resSelect['Value']:
        replicaIDs.append(record[0])
        gLogger.verbose(
            "%s.%s_DB: to_update CacheReplicas =  %s" %
            (self._caller(), 'updateReplicaInformation', record))

      reqSelect1 = "SELECT * FROM CacheReplicas WHERE ReplicaID IN (%s);" % intListToString(
          replicaIDs)
      resSelect1 = self._query(reqSelect1)
      if not resSelect1['OK']:
        gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                     (self._caller(), 'updateReplicaInformation', reqSelect1, resSelect1['Message']))
      else:
        for record in resSelect1['Value']:
          gLogger.verbose(
              "%s.%s_DB: updated CacheReplicas = %s" %
              (self._caller(), 'updateReplicaInformation', record))

      gLogger.debug(
          'StagerDB.updateReplicaInformation: Successfully updated CacheReplicas record With Status=Waiting, for ReplicaID= %s' %
          (replicaID))
    return S_OK()

  ####################################################################
  #
  # The state transition of the CacheReplicas from Waiting->StageSubmitted
  #

  def getSubmittedStagePins(self):
    # change the query to take into account pin expiry time
    req = "SELECT SE,COUNT(*),SUM(Size) from CacheReplicas WHERE Status NOT IN ('New','Waiting','Offline','Failed') GROUP BY SE;"
    # req = "SELECT SE,Count(*),SUM(Size) from CacheReplicas,StageRequests WHERE Status NOT IN ('New','Waiting','Failed') and CacheReplicas.ReplicaID=StageRequests.ReplicaID and PinExpiryTime>UTC_TIMESTAMP() GROUP BY SE;"
    res = self._query(req)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.getSubmittedStagePins: Failed to obtain submitted requests.',
          res['Message'])
      return res
    storageRequests = {}
    for storageElement, replicas, totalSize in res['Value']:
      storageRequests[storageElement] = {'Replicas': int(replicas), 'TotalSize': int(totalSize)}
    return S_OK(storageRequests)

  def insertStageRequest(self, requestDict, pinLifeTime):
    req = "INSERT INTO StageRequests (ReplicaID,RequestID,StageRequestSubmitTime,PinLength) VALUES "
    for requestID, replicaIDs in requestDict.iteritems():
      for replicaID in replicaIDs:
        replicaString = "(%s,'%s',UTC_TIMESTAMP(),%d)," % (replicaID, requestID, pinLifeTime)
        req = "%s %s" % (req, replicaString)
    req = req.rstrip(',')
    res = self._update(req)
    if not res['OK']:
      gLogger.error(
          'StorageManagementDB.insertStageRequest: Failed to insert to StageRequests table.',
          res['Message'])
      return res

    for requestID, replicaIDs in requestDict.iteritems():
      for replicaID in replicaIDs:
        # fix, no individual queries
        reqSelect = "SELECT * FROM StageRequests WHERE ReplicaID = %s AND RequestID = '%s';" % (
            replicaID, requestID)
        resSelect = self._query(reqSelect)
        if not resSelect['OK']:
          gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                       (self._caller(), 'insertStageRequest', reqSelect, resSelect['Message']))
        else:
          gLogger.verbose("%s.%s_DB: inserted StageRequests = %s" %
                          (self._caller(), 'insertStageRequest', resSelect['Value'][0]))

    # gLogger.info( "%s_DB: howmany = %s" % ('insertStageRequest',res))

    # gLogger.info( "%s_DB:%s" % ('insertStageRequest',req))
    gLogger.debug(
        "StorageManagementDB.insertStageRequest: Successfully added %s StageRequests with RequestID %s." %
        (res['Value'], requestID))
    return S_OK()

  ####################################################################
  #
  # The state transition of the CacheReplicas from StageSubmitted->Staged
  #

  def setStageComplete(self, replicaIDs):
    # Daniela: FIX wrong PinExpiryTime (84000->86400 seconds = 1 day)

    reqSelect = "SELECT * FROM StageRequests WHERE ReplicaID IN (%s);" % intListToString(replicaIDs)
    resSelect = self._query(reqSelect)
    if not resSelect['OK']:
      gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                   (self._caller(), 'setStageComplete', reqSelect, resSelect['Message']))
      return resSelect

    req = "UPDATE StageRequests SET StageStatus='Staged',StageRequestCompletedTime = UTC_TIMESTAMP(),PinExpiryTime = DATE_ADD(UTC_TIMESTAMP(),INTERVAL ( PinLength / %s ) SECOND) WHERE ReplicaID IN (%s);" % (THROTTLING_STEPS, intListToString(replicaIDs))
    res = self._update(req)
    if not res['OK']:
      gLogger.error(
          "StorageManagementDB.setStageComplete: Failed to set StageRequest completed.",
          res['Message'])
      return res

    for record in resSelect['Value']:
      gLogger.verbose("%s.%s_DB: to_update StageRequests =  %s" %
                      (self._caller(), 'setStageComplete', record))

    reqSelect1 = "SELECT * FROM StageRequests WHERE ReplicaID IN (%s);" % intListToString(
        replicaIDs)
    resSelect1 = self._query(reqSelect1)
    if not resSelect1['OK']:
      gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                   (self._caller(), 'setStageComplete', reqSelect1, resSelect1['Message']))

    for record in resSelect1['Value']:
      gLogger.verbose(
          "%s.%s_DB: updated StageRequests = %s" %
          (self._caller(), 'setStageComplete', record))

    gLogger.debug(
        "StorageManagementDB.setStageComplete: Successfully updated %s StageRequests table with StageStatus=Staged for ReplicaIDs: %s." %
        (res['Value'], replicaIDs))
    return res

  def wakeupOldRequests(self, replicaIDs, retryInterval, connection=False):
    """
    get only StageRequests with StageRequestSubmitTime older than 1 day AND are still not staged
    delete these requests
    reset Replicas with corresponding ReplicaIDs to Status='New'
    """
    try:
      retryInterval = max(retryInterval, 2)
      retryInterval = min(retryInterval, 24)
      retryInterval = int(retryInterval)
    except Exception:
      errorString = 'Wrong argument type'
      gLogger.exception(errorString)
      return S_ERROR(errorString)
    if(replicaIDs) > 0:
      req = "SELECT ReplicaID FROM StageRequests WHERE ReplicaID IN (%s) AND StageStatus='StageSubmitted' AND DATE_ADD( StageRequestSubmitTime, INTERVAL %s HOUR ) < UTC_TIMESTAMP();" % (
          intListToString(replicaIDs), retryInterval)
      res = self._query(req)
      if not res['OK']:
        gLogger.error(
            "StorageManagementDB.wakeupOldRequests: Failed to select old StageRequests.",
            res['Message'])
        return res

      old_replicaIDs = [row[0] for row in res['Value']]

      if len(old_replicaIDs) > 0:
        req = "UPDATE CacheReplicas SET Status='New',LastUpdate = UTC_TIMESTAMP(), Reason = 'wakeupOldRequests' WHERE ReplicaID in (%s);" % intListToString(old_replicaIDs)
        res = self._update(req, connection)
        if not res['OK']:
          gLogger.error(
              "StorageManagementDB.wakeupOldRequests: Failed to roll CacheReplicas back to Status=New.",
              res['Message'])
          return res

        req = "DELETE FROM StageRequests WHERE ReplicaID in (%s);" % intListToString(old_replicaIDs)
        res = self._update(req, connection)
        if not res['OK']:
          gLogger.error(
              "StorageManagementDB.wakeupOldRequests. Problem removing entries from StageRequests.")
          return res

    return S_OK()

  ####################################################################
  #
  # This code handles the finalization of stage tasks
  #
  # Daniela: useless method
  '''
  def updateStageCompletingTasks(self):
    """ This will select all the Tasks in StageCompleting status and check whether all the associated files are Staged. """
    req = "SELECT TR.TaskID,COUNT(if(R.Status NOT IN ('Staged'),1,NULL)) FROM Tasks AS T, TaskReplicas AS TR, CacheReplicas AS R WHERE T.Status='StageCompleting' AND T.TaskID=TR.TaskID AND TR.ReplicaID=R.ReplicaID GROUP BY TR.TaskID;"
    res = self._query(req)
    if not res['OK']:
      return res
    taskIDs = []
    for taskID,count in res['Value']:
      if int(count) == 0:
        taskIDs.append(taskID)
    if not taskIDs:
      return S_OK(taskIDs)
    req = "UPDATE Tasks SET Status = 'Staged' WHERE TaskID IN (%s);" % intListToString(taskIDs)
    res = self._update(req)
    if not res['OK']:
      return res
    return S_OK(taskIDs)
  '''

  def setTasksDone(self, taskIDs):
    """ This will update the status for a list of taskIDs to Done. """
    reqSelect = "SELECT * FROM Tasks WHERE TaskID IN (%s);" % intListToString(taskIDs)
    resSelect = self._query(reqSelect)
    if not resSelect['OK']:
      gLogger.error("%s.%s_DB: problem retrieving record: %s. %s" %
                    (self._caller(), 'setTasksDone', reqSelect, resSelect['Message']))

    req = "UPDATE Tasks SET Status = 'Done', CompleteTime = UTC_TIMESTAMP() WHERE TaskID IN (%s);" % intListToString(taskIDs)
    res = self._update(req)
    if not res['OK']:
      gLogger.error(
          "StorageManagementDB.setTasksDone: Failed to set Tasks status to Done.",
          res['Message'])
      return res

    for record in resSelect['Value']:
      gLogger.verbose("%s.%s_DB: to_update Tasks =  %s" % (self._caller(), 'setTasksDone', record))
      # fix, no individual queries
    reqSelect1 = "SELECT * FROM Tasks WHERE TaskID IN (%s);" % intListToString(taskIDs)
    resSelect1 = self._query(reqSelect1)
    if not resSelect1['OK']:
      gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                   (self._caller(), 'setTasksDone', reqSelect1, resSelect1['Message']))
    else:
      for record in resSelect1['Value']:
        gLogger.verbose("%s.%s_DB: updated Tasks = %s" % (self._caller(), 'setTasksDone', record))

    gLogger.debug(
        "StorageManagementDB.setTasksDone: Successfully updated %s Tasks with StageStatus=Done for taskIDs: %s." %
        (res['Value'], taskIDs))
    return res

  def killTasksBySourceTaskID(self, sourceTaskIDs, connection=False):
    """ Given SourceTaskIDs (jobs), this will cancel further staging of files for the corresponding tasks.
    The "cancel" is actually removing all stager DB records for these jobs.
    Care must be taken to NOT cancel staging of files that are requested also by other tasks. """
    connection = self.__getConnection(connection)

    # get the TaskIDs
    req = "SELECT TaskID from Tasks WHERE SourceTaskID IN (%s);" % intListToString(sourceTaskIDs)
    res = self._query(req)
    if not res['OK']:
      gLogger.error("%s.%s_DB: problem retrieving records: %s. %s" %
                    (self._caller(), 'killTasksBySourceTaskID', req, res['Message']))
    taskIDs = [row[0] for row in res['Value']]

    # ! Make sure to only cancel file staging for files with no relations with other tasks (jobs) but the killed ones
    if taskIDs:
      req = "SELECT DISTINCT(CR.ReplicaID) FROM TaskReplicas AS TR, CacheReplicas AS CR WHERE TR.TaskID IN (%s) AND CR.Links=1 and TR.ReplicaID=CR.ReplicaID;" % intListToString(taskIDs)
      res = self._query(req)
      if not res['OK']:
        gLogger.error("%s.%s_DB: problem retrieving records: %s. %s" %
                      (self._caller(), 'killTasksBySourceTaskID', req, res['Message']))

      replicaIDs = [row[0] for row in res['Value']]

      if replicaIDs:
        req = "DELETE FROM StageRequests WHERE ReplicaID IN (%s);" % intListToString(replicaIDs)
        res = self._update(req, connection)
        if not res['OK']:
          gLogger.error("%s.%s_DB: problem removing records: %s. %s" %
                        (self._caller(), 'killTasksBySourceTaskID', req, res['Message']))

        req = "DELETE FROM CacheReplicas WHERE ReplicaID in (%s) AND Links=1;" % intListToString(
            replicaIDs)
        res = self._update(req, connection)
        if not res['OK']:
          gLogger.error("%s.%s_DB: problem removing records: %s. %s" %
                        (self._caller(), 'killTasksBySourceTaskID', req, res['Message']))

      # Finally, remove the Task and TaskReplicas entries.
      res = self.removeTasks(taskIDs, connection)
    return res

  def removeStageRequests(self, replicaIDs, connection=False):
    connection = self.__getConnection(connection)
    req = "DELETE FROM StageRequests WHERE ReplicaID in (%s);" % intListToString(replicaIDs)
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error(
          "StorageManagementDB.removeStageRequests. Problem removing entries from StageRequests.")
      return res
    return res

  def removeTasks(self, taskIDs, connection=False):
    """ This will delete the entries from the TaskReplicas for the provided taskIDs. """
    connection = self.__getConnection(connection)
    req = "DELETE FROM TaskReplicas WHERE TaskID IN (%s);" % intListToString(taskIDs)
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error("StorageManagementDB.removeTasks. Problem removing entries from TaskReplicas.")
      return res
    # gLogger.info( "%s_DB:%s" % ('removeTasks',req))
    reqSelect = "SELECT * FROM Tasks WHERE TaskID IN (%s);" % intListToString(taskIDs)
    resSelect = self._query(reqSelect)
    if not resSelect['OK']:
      gLogger.error("%s.%s_DB: problem retrieving record: %s. %s" %
                    (self._caller(), 'removeTasks', reqSelect, resSelect['Message']))
    else:
      for record in resSelect['Value']:
        gLogger.verbose("%s.%s_DB: to_delete Tasks =  %s" % (self._caller(), 'removeTasks', record))

    req = "DELETE FROM Tasks WHERE TaskID in (%s);" % intListToString(taskIDs)
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error("StorageManagementDB.removeTasks. Problem removing entries from Tasks.")
    gLogger.verbose("%s.%s_DB: deleted Tasks" % (self._caller(), 'removeTasks'))
    # gLogger.info( "%s_DB:%s" % ('removeTasks',req))
    return res

  def setOldTasksAsFailed(self, daysOld, connection=False):
    """
    Set Tasks older than "daysOld" number of days to Failed
    These tasks have already been retried every day for staging
    """
    req = "UPDATE Tasks SET Status='Failed' WHERE DATE_ADD(SubmitTime, INTERVAL %s DAY ) < UTC_TIMESTAMP();" % (
        daysOld)
    res = self._update(req, connection)
    if not res['OK']:
      gLogger.error("StorageManagementDB.setOldTasksAsFailed. Problem setting old Tasks to Failed.")
      return res
    return res

  def getCacheReplicasSummary(self, connection=False):
    """
    Reports breakdown of file number/size in different staging states across storage elements
    """
    connection = self.__getConnection(connection)
    req = "SELECT DISTINCT(Status),SE,COUNT(*),sum(size)/(1024*1024*1024) FROM CacheReplicas GROUP BY Status,SE;"
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error("StorageManagementDB.getCacheReplicasSummary failed.")
      return res

    resSummary = {}
    i = 1
    for status, se, numFiles, sumFiles in res['Value']:
      resSummary[i] = {
          'Status': status,
          'SE': se,
          'NumFiles': long(numFiles),
          'SumFiles': float(sumFiles)}
      i += 1
    return S_OK(resSummary)

  def removeUnlinkedReplicas(self, connection=False):
    """ This will remove Replicas from the CacheReplicas that are not associated to any Task.
        If the Replica has been Staged,
        wait until StageRequest.PinExpiryTime and remove the StageRequest and CacheReplicas entries
    """
    connection = self.__getConnection(connection)
    # First, check if there is a StageRequest and PinExpiryTime has arrived
    req = "select SR.ReplicaID from CacheReplicas CR,StageRequests SR WHERE CR.Links = 0 and CR.ReplicaID=SR.ReplicaID group by SR.ReplicaID HAVING max(SR.PinExpiryTime) < UTC_TIMESTAMP();"
    # req = "SELECT ReplicaID from CacheReplicas WHERE Links = 0;"
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          "StorageManagementDB.removeUnlinkedReplicas. Problem selecting entries from CacheReplicas where Links = 0.")
      return res

    replicaIDs = [row[0] for row in res['Value']]
    # Look for Failed CacheReplicas which are not associated to any Task. These have no PinExpiryTime in StageRequests
    # as they were not staged successfully (for various reasons), even though
    # a staging request had been submitted
    req = "SELECT ReplicaID FROM CacheReplicas WHERE Links = 0 AND Status = 'Failed';"
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          "StorageManagementDB.removeUnlinkedReplicas. Problem selecting entries from CacheReplicas where Links = 0 AND Status=Failed.")
    else:
      replicaIDs.extend([row[0] for row in res['Value']])

    if replicaIDs:
      # Removed the entries from the StageRequests table that are expired
      reqSelect = "SELECT * FROM StageRequests WHERE ReplicaID IN (%s);" % intListToString(
          replicaIDs)
      resSelect = self._query(reqSelect)
      if not resSelect['OK']:
        gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                     (self._caller(), 'removeUnlinkedReplicas', reqSelect, resSelect['Message']))
      else:
        for record in resSelect['Value']:
          gLogger.verbose(
              "%s.%s_DB: to_delete StageRequests = %s" %
              (self._caller(), 'removeUnlinkedReplicas', record))

      req = "DELETE FROM StageRequests WHERE ReplicaID IN (%s);" % intListToString(replicaIDs)
      res = self._update(req, connection)
      if not res['OK']:
        gLogger.error(
            "StorageManagementDB.removeUnlinkedReplicas. Problem deleting from StageRequests.")
        return res
      gLogger.verbose(
          "%s.%s_DB: deleted StageRequests" %
          (self._caller(), 'removeUnlinkedReplicas'))

      gLogger.debug(
          "StorageManagementDB.removeUnlinkedReplicas: Successfully removed %s StageRequests entries for ReplicaIDs: %s." %
          (res['Value'], replicaIDs))

    # Second look for CacheReplicas for which there is no entry in StageRequests
    req = 'SELECT ReplicaID FROM CacheReplicas WHERE Links = 0 AND ReplicaID NOT IN ( SELECT DISTINCT( ReplicaID ) FROM StageRequests )'
    res = self._query(req, connection)
    if not res['OK']:
      gLogger.error(
          "StorageManagementDB.removeUnlinkedReplicas. Problem selecting entries from CacheReplicas where Links = 0.")
    else:
      replicaIDs.extend([row[0] for row in res['Value']])

    if not replicaIDs:
      return S_OK()

    # Now delete all CacheReplicas
    reqSelect = "SELECT * FROM CacheReplicas WHERE ReplicaID IN (%s);" % intListToString(replicaIDs)
    resSelect = self._query(reqSelect)
    if not resSelect['OK']:
      gLogger.warn("%s.%s_DB: problem retrieving record: %s. %s" %
                   (self._caller(), 'removeUnlinkedReplicas', reqSelect, resSelect['Message']))
    else:
      for record in resSelect['Value']:
        gLogger.verbose(
            "%s.%s_DB: to_delete CacheReplicas =  %s" %
            (self._caller(), 'removeUnlinkedReplicas', record))

    req = "DELETE FROM CacheReplicas WHERE ReplicaID IN (%s) AND Links= 0;" % intListToString(
        replicaIDs)
    res = self._update(req, connection)
    if res['OK']:
      gLogger.verbose(
          "%s.%s_DB: deleted CacheReplicas" %
          (self._caller(), 'removeUnlinkedReplicas'))
      gLogger.debug(
          "StorageManagementDB.removeUnlinkedReplicas: Successfully removed %s CacheReplicas entries for ReplicaIDs: %s." %
          (res['Value'], replicaIDs))
    else:
      gLogger.error(
          "StorageManagementDB.removeUnlinkedReplicas. Problem removing entries from CacheReplicas.")
    return res
