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)
