def ReadFlowRequestsReadyForProcessing(self, client_id, flow_id): """Reads all requests for a flow that can be processed by the worker.""" try: flow_obj = self.flows[(client_id, flow_id)] except KeyError: raise db.UnknownFlowError(client_id, flow_id) next_request_to_process = flow_obj.next_request_to_process request_dict = self.flow_requests.get((client_id, flow_id), {}) response_dict = self.flow_responses.get((client_id, flow_id), {}) res = {} for request_id in sorted(request_dict): # Ignore outdated requests. if request_id < next_request_to_process: continue # The request we are currently looking for is not in yet, we are done. if request_id != next_request_to_process: break request = request_dict[request_id] if not request.needs_processing: break responses = sorted(itervalues(response_dict.get(request_id, {})), key=lambda response: response.response_id) res[request_id] = (request, responses) next_request_to_process += 1 return res
def UpdateFlow(self, client_id, flow_id, flow_obj=db.Database.unchanged, client_crash_info=db.Database.unchanged, pending_termination=db.Database.unchanged, processing_on=db.Database.unchanged, processing_since=db.Database.unchanged, processing_deadline=db.Database.unchanged): """Updates flow objects in the database.""" try: flow = self.flows[(client_id, flow_id)] except KeyError: raise db.UnknownFlowError(client_id, flow_id) if flow_obj != db.Database.unchanged: self.flows[(client_id, flow_id)] = flow_obj flow = flow_obj if client_crash_info != db.Database.unchanged: flow.client_crash_info = client_crash_info if pending_termination != db.Database.unchanged: flow.pending_termination = pending_termination if processing_on != db.Database.unchanged: flow.processing_on = processing_on if processing_since != db.Database.unchanged: flow.processing_since = processing_since if processing_deadline != db.Database.unchanged: flow.processing_deadline = processing_deadline
def WriteFlowResponses(self, responses): """Writes a list of flow responses to the database.""" status_available = set() requests_updated = set() for response in responses: flow_key = (response.client_id, response.flow_id) if flow_key not in self.flows: raise db.UnknownFlowError(response.client_id, response.flow_id) request_dict = self.flow_requests.get(flow_key, {}) if response.request_id not in request_dict: logging.error( "Received response for unknown request %s, %s, %d.", response.client_id, response.flow_id, response.request_id) continue response_dict = self.flow_responses.setdefault(flow_key, {}) response_dict.setdefault( response.request_id, {})[response.response_id] = response.Copy() if isinstance(response, rdf_flow_objects.FlowStatus): status_available.add(response) requests_updated.add( (response.client_id, response.flow_id, response.request_id)) # Every time we get a status we store how many responses are expected. for status in status_available: request_dict = self.flow_requests[(status.client_id, status.flow_id)] request = request_dict[status.request_id] request.nr_responses_expected = status.response_id # And we check for all updated requests if we need to process them. needs_processing = [] for client_id, flow_id, request_id in requests_updated: flow_key = (client_id, flow_id) request_dict = self.flow_requests[flow_key] request = request_dict[request_id] if request.nr_responses_expected and not request.needs_processing: response_dict = self.flow_responses.setdefault(flow_key, {}) responses = response_dict.get(request_id, {}) if len(responses) == request.nr_responses_expected: request.needs_processing = True flow = self.flows[flow_key] if flow.next_request_to_process == request_id: needs_processing.append( rdf_flows.FlowProcessingRequest( client_id=client_id, flow_id=flow_id, request_id=request_id)) if needs_processing: self.WriteFlowProcessingRequests(needs_processing)
def ReadFlowObject(self, client_id, flow_id, cursor=None): """Reads a flow object from the database.""" query = ("SELECT " + self.FLOW_DB_FIELDS + "FROM flows WHERE client_id=%s AND flow_id=%s") cursor.execute(query, [ mysql_utils.ClientIDToInt(client_id), mysql_utils.FlowIDToInt(flow_id) ]) result = cursor.fetchall() if not result: raise db.UnknownFlowError(client_id, flow_id) row, = result return self._FlowObjectFromRow(row)
def UpdateFlow(self, client_id, flow_id, flow_obj=db.Database.unchanged, flow_state=db.Database.unchanged, client_crash_info=db.Database.unchanged, pending_termination=db.Database.unchanged, processing_on=db.Database.unchanged, processing_since=db.Database.unchanged, processing_deadline=db.Database.unchanged, cursor=None): """Updates flow objects in the database.""" updates = [] args = [] if flow_obj != db.Database.unchanged: updates.append("flow=%s") args.append(flow_obj.SerializeToString()) updates.append("flow_state=%s") args.append(int(flow_obj.flow_state)) if flow_state != db.Database.unchanged: updates.append("flow_state=%s") args.append(int(flow_state)) if client_crash_info != db.Database.unchanged: updates.append("client_crash_info=%s") args.append(client_crash_info.SerializeToString()) if pending_termination != db.Database.unchanged: updates.append("pending_termination=%s") args.append(pending_termination.SerializeToString()) if processing_on != db.Database.unchanged: updates.append("processing_on=%s") args.append(processing_on) if processing_since != db.Database.unchanged: updates.append("processing_since=%s") args.append(mysql_utils.RDFDatetimeToMysqlString(processing_since)) if processing_deadline != db.Database.unchanged: updates.append("processing_deadline=%s") args.append(mysql_utils.RDFDatetimeToMysqlString(processing_deadline)) if not updates: return query = "UPDATE flows SET " query += ", ".join(updates) query += " WHERE client_id=%s AND flow_id=%s" args.append(mysql_utils.ClientIDToInt(client_id)) args.append(mysql_utils.FlowIDToInt(flow_id)) updated = cursor.execute(query, args) if updated == 0: raise db.UnknownFlowError(client_id, flow_id)
def ReadAllFlowRequestsAndResponses(self, client_id, flow_id): """Reads all requests and responses for a given flow from the database.""" flow_key = (client_id, flow_id) try: self.flows[flow_key] except KeyError: raise db.UnknownFlowError(client_id, flow_id) request_dict = self.flow_requests.get(flow_key, {}) response_dict = self.flow_responses.get(flow_key, {}) res = [] for request_id in sorted(request_dict): res.append( (request_dict[request_id], response_dict.get(request_id, []))) return res
def DeleteAllFlowRequestsAndResponses(self, client_id, flow_id): """Deletes all requests and responses for a given flow from the database.""" flow_key = (client_id, flow_id) try: self.flows[flow_key] except KeyError: raise db.UnknownFlowError(client_id, flow_id) try: del self.flow_requests[flow_key] except KeyError: pass try: del self.flow_responses[flow_key] except KeyError: pass
def ReadFlowForProcessing(self, client_id, flow_id, processing_time, cursor=None): """Marks a flow as being processed on this worker and returns it.""" query = ("SELECT " + self.FLOW_DB_FIELDS + "FROM flows WHERE client_id=%s AND flow_id=%s") cursor.execute(query, [ mysql_utils.ClientIDToInt(client_id), mysql_utils.FlowIDToInt(flow_id) ]) response = cursor.fetchall() if not response: raise db.UnknownFlowError(client_id, flow_id) row, = response rdf_flow = self._FlowObjectFromRow(row) now = rdfvalue.RDFDatetime.Now() if rdf_flow.processing_on and rdf_flow.processing_deadline > now: raise ValueError( "Flow %s on client %s is already being processed." % (client_id, flow_id)) update_query = ( "UPDATE flows SET processing_on=%s, processing_since=%s, " "processing_deadline=%s WHERE client_id=%s and flow_id=%s") processing_deadline = now + processing_time process_id_string = utils.ProcessIdString() args = [ process_id_string, mysql_utils.RDFDatetimeToMysqlString(now), mysql_utils.RDFDatetimeToMysqlString(processing_deadline), mysql_utils.ClientIDToInt(client_id), mysql_utils.FlowIDToInt(flow_id) ] cursor.execute(update_query, args) # This needs to happen after we are sure that the write has succeeded. rdf_flow.processing_on = process_id_string rdf_flow.processing_since = now rdf_flow.processing_deadline = processing_deadline return rdf_flow
def DeleteFlowRequests(self, requests): """Deletes a list of flow requests from the database.""" for request in requests: if (request.client_id, request.flow_id) not in self.flows: raise db.UnknownFlowError(request.client_id, request.flow_id) for request in requests: key = (request.client_id, request.flow_id) request_dict = self.flow_requests.get(key, {}) try: del request_dict[request.request_id] except KeyError: raise db.UnknownFlowRequestError(request.client_id, request.flow_id, request.request_id) response_dict = self.flow_responses.get(key, {}) try: del response_dict[request.request_id] except KeyError: pass
def WriteFlowRequests(self, requests): """Writes a list of flow requests to the database.""" flow_processing_requests = [] for request in requests: if (request.client_id, request.flow_id) not in self.flows: raise db.UnknownFlowError(request.client_id, request.flow_id) for request in requests: key = (request.client_id, request.flow_id) request_dict = self.flow_requests.setdefault(key, {}) request_dict[request.request_id] = request.Copy() if request.needs_processing: flow = self.flows[(request.client_id, request.flow_id)] if flow.next_request_to_process == request.request_id: flow_processing_requests.append( rdf_flows.FlowProcessingRequest( client_id=request.client_id, flow_id=request.flow_id, request_id=request.request_id)) if flow_processing_requests: self.WriteFlowProcessingRequests(flow_processing_requests)
def ReadFlowObject(self, client_id, flow_id): """Reads a flow object from the database.""" try: return self.flows[(client_id, flow_id)].Copy() except KeyError: raise db.UnknownFlowError(client_id, flow_id)