def get_reply_for_message(msg_str: str) -> ZMQReply:
"""
Parse the zmq message string, perform the desired action and return the result in JSON format.
Parameters
----------
msg_str : str
The message string as received from zmq. See the docstring of `parse_zmq_message` for valid structure.
Returns
-------
dict
The reply in JSON format.
"""
try:
action_request = ActionRequest().loads(msg_str)
query_run_log.info(
f"Attempting to perform action: '{action_request.action}'",
request_id=action_request.request_id,
action=action_request.action,
params=action_request.params,
)
reply = perform_action(action_request.action, action_request.params)
query_run_log.info(
f"Action completed with status: '{reply.status}'",
request_id=action_request.request_id,
action=action_request.action,
params=action_request.params,
reply_status=reply.status,
reply_msg=reply.msg,
reply_payload=reply.payload,
)
except FlowmachineServerError as exc:
return ZMQReply(status="error", msg=exc.error_msg)
except ValidationError as exc:
# The dictionary of marshmallow errors can contain integers as keys,
# which will raise an error when converting to JSON (where the keys
# must be strings). Therefore we transform the keys to strings here.
error_msg = "Invalid action request."
validation_error_messages = convert_dict_keys_to_strings(exc.messages)
return ZMQReply(status="error",
msg=error_msg,
payload=validation_error_messages)
except JSONDecodeError as exc:
return ZMQReply(status="error",
msg="Invalid JSON.",
payload={"decode_error": exc.msg})
# Return the reply (in JSON format)
return reply
def action_handler__get_geography(aggregation_unit: str) -> ZMQReply:
"""
Handler for the 'get_query_geography' action.
Returns SQL to get geography for the given `aggregation_unit` as GeoJSON.
"""
try:
try:
try:
query_obj = GeographySchema().load(
{"aggregation_unit": aggregation_unit}
)
except TypeError as exc:
# We need to catch TypeError here, otherwise they propagate up to
# perform_action() and result in a very misleading error message.
orig_error_msg = exc.args[0]
error_msg = (
f"Internal flowmachine server error: could not create query object using query schema. "
f"The original error was: '{orig_error_msg}'"
)
return ZMQReply(
status="error",
msg=error_msg,
payload={
"params": {"aggregation_unit": aggregation_unit},
"orig_error_msg": orig_error_msg,
},
)
except ValidationError as exc:
# The dictionary of marshmallow errors can contain integers as keys,
# which will raise an error when converting to JSON (where the keys
# must be strings). Therefore we transform the keys to strings here.
error_msg = "Parameter validation failed."
validation_error_messages = convert_dict_keys_to_strings(exc.messages)
return ZMQReply(
status="error", msg=error_msg, payload=validation_error_messages
)
# We don't cache the query, because it just selects columns from a
# geography table. If we expose an aggregation unit which relies on another
# query to create the geometry (e.g. grid), we may want to reconsider this
# decision.
sql = query_obj.geojson_sql
# TODO: put query_run_log back in!
# query_run_log.info("get_geography", **run_log_dict)
payload = {"query_state": QueryState.COMPLETED, "sql": sql}
return ZMQReply(status="success", payload=payload)
except Exception as exc:
# If we don't catch exceptions here, the server will die and FlowAPI will hang indefinitely.
error_msg = f"Internal flowmachine server error: '{exc.args[0]}'"
return ZMQReply(
status="error", msg=error_msg, payload={"error_msg": exc.args[0]}
)
def action_handler__run_query(**action_params):
"""
Handler for the 'run_query' action.
Constructs a flowmachine query object, sets it running and returns the query_id.
For this action handler the `action_params` are exactly the query kind plus the
parameters needed to construct the query.
"""
try:
query_obj = FlowmachineQuerySchema().load(action_params)
except ValidationError as exc:
# The dictionary of marshmallow errors can contain integers as keys,
# which will raise an error when converting to JSON (where the keys
# must be strings). Therefore we transform the keys to strings here.
error_msg = "Parameter validation failed."
validation_error_messages = convert_dict_keys_to_strings(exc.messages)
return ZMQReply(status="error",
msg=error_msg,
payload=validation_error_messages)
q_info_lookup = QueryInfoLookup(Query.redis)
try:
query_id = q_info_lookup.get_query_id(action_params)
except QueryInfoLookupError:
# Set the query running (it's safe to call this even if the query was set running before)
try:
query_id = query_obj.store_async()
except Exception as e:
return ZMQReply(
status="error",
msg="Unable to create query object.",
payload={"exception": str(e)},
)
# Register the query as "known" (so that we can later look up the query kind
# and its parameters from the query_id).
q_info_lookup.register_query(query_id, action_params)
return ZMQReply(status="success", payload={"query_id": query_id})
async def action_handler__run_query(config: "FlowmachineServerConfig",
**action_params: dict) -> ZMQReply:
"""
Handler for the 'run_query' action.
Constructs a flowmachine query object, sets it running and returns the query_id.
For this action handler the `action_params` are exactly the query kind plus the
parameters needed to construct the query.
"""
try:
query_obj = FlowmachineQuerySchema().load(action_params)
except TypeError as exc:
# We need to catch TypeError here, otherwise they propagate up to
# perform_action() and result in a very misleading error message.
orig_error_msg = exc.args[0]
error_msg = (
f"Internal flowmachine server error: could not create query object using query schema. "
f"The original error was: '{orig_error_msg}'")
return ZMQReply(
status="error",
msg=error_msg,
payload={
"params": action_params,
"orig_error_msg": orig_error_msg
},
)
except ValidationError as exc:
# The dictionary of marshmallow errors can contain integers as keys,
# which will raise an error when converting to JSON (where the keys
# must be strings). Therefore we transform the keys to strings here.
validation_error_messages = convert_dict_keys_to_strings(exc.messages)
action_params_as_text = textwrap.indent(
json.dumps(action_params, indent=2), " ")
validation_errors_as_text = textwrap.indent(
json.dumps(validation_error_messages, indent=2), " ")
error_msg = (
"Parameter validation failed.\n\n"
f"The action parameters were:\n{action_params_as_text}.\n\n"
f"Validation error messages:\n{validation_errors_as_text}.\n\n")
payload = {"validation_error_messages": validation_error_messages}
return ZMQReply(status="error", msg=error_msg, payload=payload)
q_info_lookup = QueryInfoLookup(get_redis())
try:
query_id = q_info_lookup.get_query_id(action_params)
qsm = QueryStateMachine(query_id=query_id,
redis_client=get_redis(),
db_id=get_db().conn_id)
if qsm.current_query_state in [
QueryState.CANCELLED,
QueryState.KNOWN,
]: # Start queries running even if they've been cancelled or reset
if qsm.is_cancelled:
reset = qsm.reset()
finish = qsm.finish_resetting()
raise QueryInfoLookupError
except QueryInfoLookupError:
try:
# Set the query running (it's safe to call this even if the query was set running before)
query_id = await asyncio.get_running_loop().run_in_executor(
executor=config.server_thread_pool,
func=partial(
copy_context().run,
partial(
query_obj.store_async,
store_dependencies=config.store_dependencies,
),
),
)
except Exception as e:
return ZMQReply(
status="error",
msg="Unable to create query object.",
payload={"exception": str(e)},
)
# Register the query as "known" (so that we can later look up the query kind
# and its parameters from the query_id).
q_info_lookup.register_query(query_id, action_params)
return ZMQReply(
status="success",
payload={
"query_id": query_id,
"progress": query_progress(query_obj._flowmachine_query_obj),
},
)
def action_handler__run_query(**action_params: dict) -> ZMQReply:
"""
Handler for the 'run_query' action.
Constructs a flowmachine query object, sets it running and returns the query_id.
For this action handler the `action_params` are exactly the query kind plus the
parameters needed to construct the query.
"""
try:
try:
query_obj = FlowmachineQuerySchema().load(action_params)
except TypeError as exc:
# We need to catch TypeError here, otherwise they propagate up to
# perform_action() and result in a very misleading error message.
orig_error_msg = exc.args[0]
error_msg = (
f"Internal flowmachine server error: could not create query object using query schema. "
f"The original error was: '{orig_error_msg}'"
)
return ZMQReply(
status="error",
msg=error_msg,
payload={"params": action_params, "orig_error_msg": orig_error_msg},
)
except ValidationError as exc:
# The dictionary of marshmallow errors can contain integers as keys,
# which will raise an error when converting to JSON (where the keys
# must be strings). Therefore we transform the keys to strings here.
validation_error_messages = convert_dict_keys_to_strings(exc.messages)
action_params_as_text = textwrap.indent(
json.dumps(action_params, indent=2), " "
)
validation_errors_as_text = textwrap.indent(
json.dumps(validation_error_messages, indent=2), " "
)
error_msg = (
"Parameter validation failed.\n\n"
f"The action parameters were:\n{action_params_as_text}.\n\n"
f"Validation error messages:\n{validation_errors_as_text}.\n\n"
)
payload = {"validation_error_messages": validation_error_messages}
return ZMQReply(status="error", msg=error_msg, payload=payload)
q_info_lookup = QueryInfoLookup(Query.redis)
try:
query_id = q_info_lookup.get_query_id(action_params)
except QueryInfoLookupError:
# Set the query running (it's safe to call this even if the query was set running before)
try:
query_id = query_obj.store_async()
except Exception as e:
return ZMQReply(
status="error",
msg="Unable to create query object.",
payload={"exception": str(e)},
)
# Register the query as "known" (so that we can later look up the query kind
# and its parameters from the query_id).
q_info_lookup.register_query(query_id, action_params)
return ZMQReply(status="success", payload={"query_id": query_id})
async def get_reply_for_message(*, msg_str: str,
config: "FlowmachineServerConfig") -> ZMQReply:
"""
Parse the zmq message string, perform the desired action and return the result in JSON format.
Parameters
----------
msg_str : str
The message string as received from zmq. See the docstring of `parse_zmq_message` for valid structure.
config : FlowmachineServerConfig
Server config options
Returns
-------
dict
The reply in JSON format.
"""
try:
try:
action_request = ActionRequest().loads(msg_str)
except ValidationError as exc:
# The dictionary of marshmallow errors can contain integers as keys,
# which will raise an error when converting to JSON (where the keys
# must be strings). Therefore we transform the keys to strings here.
error_msg = "Invalid action request."
validation_error_messages = convert_dict_keys_to_strings(
exc.messages)
logger.error(
"Invalid action request while getting reply for ZMQ message.",
**validation_error_messages,
)
return ZMQReply(status="error",
msg=error_msg,
payload=validation_error_messages)
with action_request_context(action_request):
query_run_log.info(
f"Attempting to perform action: '{action_request.action}'",
params=action_request.params,
)
reply = await perform_action(action_request.action,
action_request.params,
config=config)
query_run_log.info(
f"Action completed with status: '{reply.status}'",
params=action_request.params,
reply_status=reply.status,
reply_msg=reply.msg,
reply_payload=reply.payload,
)
except FlowmachineServerError as exc:
logger.error(
f"Caught Flowmachine server error while getting reply for ZMQ message: {exc.error_msg}"
)
return ZMQReply(status="error", msg=exc.error_msg)
except JSONDecodeError as exc:
logger.error(
f"Invalid JSON while getting reply for ZMQ message: {exc.msg}")
return ZMQReply(status="error",
msg="Invalid JSON.",
payload={"decode_error": exc.msg})
# Return the reply (in JSON format)
return reply