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.AtLeastOneUnknownFlowError([(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()
request_dict[
request.request_id].timestamp = rdfvalue.RDFDatetime.Now()
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))
if flow_processing_requests:
self.WriteFlowProcessingRequests(flow_processing_requests)
def WriteFlowRequests(self, requests, cursor=None):
"""Writes a list of flow requests to the database."""
args = []
templates = []
flow_keys = []
needs_processing = {}
now_str = mysql_utils.RDFDatetimeToMysqlString(
rdfvalue.RDFDatetime.Now())
for r in requests:
if r.needs_processing:
needs_processing.setdefault((r.client_id, r.flow_id),
[]).append(r.request_id)
flow_keys.append((r.client_id, r.flow_id))
templates.append("(%s, %s, %s, %s, %s, %s)")
args.extend([
mysql_utils.ClientIDToInt(r.client_id),
mysql_utils.FlowIDToInt(r.flow_id), r.request_id,
r.needs_processing,
r.SerializeToString(), now_str
])
if needs_processing:
flow_processing_requests = []
nr_conditions = []
nr_args = []
for client_id, flow_id in needs_processing:
nr_conditions.append("(client_id=%s AND flow_id=%s)")
nr_args.append(mysql_utils.ClientIDToInt(client_id))
nr_args.append(mysql_utils.FlowIDToInt(flow_id))
nr_query = ("SELECT client_id, flow_id, next_request_to_process "
"FROM flows WHERE ")
nr_query += " OR ".join(nr_conditions)
cursor.execute(nr_query, nr_args)
db_result = cursor.fetchall()
for client_id_int, flow_id_int, next_request_to_process in db_result:
client_id = mysql_utils.IntToClientID(client_id_int)
flow_id = mysql_utils.IntToFlowID(flow_id_int)
if next_request_to_process in needs_processing[(client_id,
flow_id)]:
flow_processing_requests.append(
rdf_flows.FlowProcessingRequest(client_id=client_id,
flow_id=flow_id))
if flow_processing_requests:
self._WriteFlowProcessingRequests(flow_processing_requests,
cursor)
query = ("INSERT INTO flow_requests "
"(client_id, flow_id, request_id, needs_processing, request, "
"timestamp) VALUES ")
query += ", ".join(templates)
try:
cursor.execute(query, args)
except MySQLdb.IntegrityError as e:
raise db.AtLeastOneUnknownFlowError(flow_keys, cause=e)
def WriteFlowLogEntries(self, entries):
"""Writes flow log entries for a given flow."""
flow_ids = [(e.client_id, e.flow_id) for e in entries]
for f in flow_ids:
if f not in self.flows:
raise db.AtLeastOneUnknownFlowError(flow_ids)
for e in entries:
dest = self.flow_log_entries.setdefault((e.client_id, e.flow_id), [])
to_write = e.Copy()
to_write.timestamp = rdfvalue.RDFDatetime.Now()
dest.append(to_write)