def _ReadCurrentFlowInfo(self, responses, currently_available_requests,
next_request_by_flow, responses_expected_by_request,
current_responses_by_request, cursor):
"""Reads stored data for flows we want to modify."""
flow_conditions = []
flow_args = []
req_conditions = []
req_args = []
for r in responses:
flow_conditions.append("(client_id=%s AND flow_id=%s)")
flow_args.append(mysql_utils.ClientIDToInt(r.client_id))
flow_args.append(mysql_utils.FlowIDToInt(r.flow_id))
req_conditions.append("(client_id=%s AND flow_id=%s AND request_id=%s)")
req_args.append(mysql_utils.ClientIDToInt(r.client_id))
req_args.append(mysql_utils.FlowIDToInt(r.flow_id))
req_args.append(r.request_id)
flow_query = ("SELECT client_id, flow_id, next_request_to_process "
"FROM flows WHERE ")
flow_query += " OR ".join(flow_conditions)
req_query = ("SELECT client_id, flow_id, request_id, responses_expected "
"FROM flow_requests WHERE ")
req_query += " OR ".join(req_conditions)
res_query = ("SELECT client_id, flow_id, request_id, response_id "
"FROM flow_responses WHERE ")
res_query += " OR ".join(req_conditions)
cursor.execute(flow_query, flow_args)
for row in cursor.fetchall():
client_id_int, flow_id_int, next_request_to_process = row
client_id = mysql_utils.IntToClientID(client_id_int)
flow_id = mysql_utils.IntToFlowID(flow_id_int)
next_request_by_flow[(client_id, flow_id)] = next_request_to_process
cursor.execute(req_query, req_args)
for row in cursor.fetchall():
client_id_int, flow_id_int, request_id, responses_expected = row
client_id = mysql_utils.IntToClientID(client_id_int)
flow_id = mysql_utils.IntToFlowID(flow_id_int)
request_key = (client_id, flow_id, request_id)
currently_available_requests.add(request_key)
if responses_expected:
responses_expected_by_request[request_key] = responses_expected
cursor.execute(res_query, req_args)
for row in cursor.fetchall():
client_id_int, flow_id_int, request_id, response_id = row
client_id = mysql_utils.IntToClientID(client_id_int)
flow_id = mysql_utils.IntToFlowID(flow_id_int)
request_key = (client_id, flow_id, request_id)
current_responses_by_request.setdefault(request_key,
set()).add(response_id)
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)