async def action_handler__get_sql(config: "FlowmachineServerConfig",
query_id: str) -> ZMQReply:
"""
Handler for the 'get_sql' action.
Returns a SQL string which can be run against flowdb to obtain
the result of the query with given `query_id`.
"""
# TODO: currently we can't use QueryStateMachine to determine whether
# the query_id belongs to a valid query object, so we need to check it
# manually. Would be good to add a QueryState.UNKNOWN so that we can
# avoid this separate treatment.
q_info_lookup = QueryInfoLookup(get_redis())
if not q_info_lookup.query_is_known(query_id):
msg = f"Unknown query id: '{query_id}'"
payload = {"query_id": query_id, "query_state": "awol"}
return ZMQReply(status="error", msg=msg, payload=payload)
query_state = QueryStateMachine(get_redis(), query_id,
get_db().conn_id).current_query_state
if query_state == QueryState.COMPLETED:
q = get_query_object_by_id(get_db(), query_id)
sql = q.get_query()
payload = {
"query_id": query_id,
"query_state": query_state,
"sql": sql
}
return ZMQReply(status="success", payload=payload)
else:
msg = f"Query with id '{query_id}' {query_state.description}."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
def test_retrieve_query_kind_fails_for_unknown_query(dummy_redis):
"""
Test that retrieving a query's parameters raises an error if it has not been registered.
"""
q_info_lookup = QueryInfoLookup(dummy_redis)
with pytest.raises(QueryInfoLookupError, match="Unknown query_id: 'dummy_id'"):
q_info_lookup.get_query_kind("dummy_id")
def action_handler__get_sql(query_id):
"""
Handler for the 'get_sql' action.
Returns a SQL string which can be run against flowdb to obtain
the result of the query with given `query_id`.
"""
# TODO: currently we can't use QueryStateMachine to determine whether
# the query_id belongs to a valid query object, so we need to check it
# manually. Would be good to add a QueryState.UNKNOWN so that we can
# avoid this separate treatment.
q_info_lookup = QueryInfoLookup(Query.redis)
if not q_info_lookup.query_is_known(query_id):
msg = f"Unknown query id: '{query_id}'"
payload = {"query_id": query_id, "query_state": "awol"}
return ZMQReply(status="error", msg=msg, payload=payload)
query_state = QueryStateMachine(Query.redis, query_id).current_query_state
if query_state == QueryState.COMPLETED:
q = get_query_object_by_id(Query.connection, query_id)
sql = q.get_query()
payload = {
"query_id": query_id,
"query_state": query_state,
"sql": sql
}
return ZMQReply(status="success", payload=payload)
elif query_state == QueryState.EXECUTING:
msg = f"Query with id '{query_id}' is still running."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
elif query_state == QueryState.QUEUED:
msg = f"Query with id '{query_id}' is still queued."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
elif query_state == QueryState.ERRORED:
msg = f"Query with id '{query_id}' is failed."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
elif query_state == QueryState.CANCELLED:
msg = f"Query with id '{query_id}' was cancelled."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
elif query_state == QueryState.RESETTING:
msg = f"Query with id '{query_id}' is being removed from cache."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
elif query_state == QueryState.KNOWN:
msg = f"Query with id '{query_id}' has not been run yet, or was reset."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
else:
msg = f"Unknown state for query with id '{query_id}'. Got {query_state}."
return ZMQReply(status="error", msg=msg)
def test_retrieve_query_kind(dummy_redis):
"""
Test that we can retrieve the query kind after registering a query.
"""
q_info_lookup = QueryInfoLookup(dummy_redis)
q_info_lookup.register_query(
query_id="dummy_id",
query_params={"query_kind": "dummy_query", "dummy_param": "some_value"},
)
assert "dummy_query" == q_info_lookup.get_query_kind("dummy_id")
def test_retrieve_query_params(dummy_redis):
"""
Test that we can retrieve a query's parameters after registering it.
"""
q_info_lookup = QueryInfoLookup(dummy_redis)
q_info_lookup.register_query(
query_id="dummy_id",
query_params={"query_kind": "dummy_query", "dummy_param": "some_value"},
)
expected_query_params = {"query_kind": "dummy_query", "dummy_param": "some_value"}
assert expected_query_params == q_info_lookup.get_query_params("dummy_id")
def test_cannot_register_query_params_without_query_kind(dummy_redis):
"""
Test that registering a query fails if the parameters do not contain the 'query_kind' key.
"""
q_info_lookup = QueryInfoLookup(dummy_redis)
with pytest.raises(
QueryInfoLookupError, match="Query params must contain a 'query_kind' entry."
):
q_info_lookup.register_query(
query_id="dummy_id", query_params={"dummy_param": "some_value"}
)
async def action_handler__poll_query(config: "FlowmachineServerConfig",
query_id: str) -> ZMQReply:
"""
Handler for the 'poll_query' action.
Returns the status of the query with the given `query_id`.
"""
query_kind = _get_query_kind_for_query_id(query_id)
# TODO: we should probably be able to use the QueryStateMachine to determine
# whether the query already exists.
if query_kind is None:
payload = {"query_id": query_id, "query_state": "awol"}
return ZMQReply(status="error",
msg=f"Unknown query id: '{query_id}'",
payload=payload)
else:
q_state_machine = QueryStateMachine(get_redis(), query_id,
get_db().conn_id)
payload = {
"query_id":
query_id,
"query_kind":
query_kind,
"query_state":
q_state_machine.current_query_state,
"progress":
query_progress(FlowmachineQuerySchema().load(
QueryInfoLookup(get_redis()).get_query_params(
query_id))._flowmachine_query_obj),
}
return ZMQReply(status="success", payload=payload)
async def test_rerun_query_after_cancelled(server_config, real_connections):
"""
Test that a query can be rerun after it has been cancelled.
"""
query_obj = (FlowmachineQuerySchema().load(
dict(
query_kind="spatial_aggregate",
locations=dict(
query_kind="daily_location",
date="2016-01-01",
method="last",
aggregation_unit="admin3",
),
))._flowmachine_query_obj)
query_id = query_obj.query_id
qsm = QueryStateMachine(get_redis(), query_id, get_db().conn_id)
qsm.enqueue()
qsm.cancel()
assert not query_obj.is_stored
assert qsm.is_cancelled
query_info_lookup = QueryInfoLookup(get_redis())
query_info_lookup.register_query(
query_id,
dict(
query_kind="spatial_aggregate",
locations=dict(
query_kind="daily_location",
date="2016-01-01",
method="last",
aggregation_unit="admin3",
),
),
)
msg = await action_handler__run_query(
config=server_config,
query_kind="spatial_aggregate",
locations=dict(
query_kind="daily_location",
date="2016-01-01",
method="last",
aggregation_unit="admin3",
),
)
assert msg["status"] == ZMQReplyStatus.SUCCESS
qsm.wait_until_complete()
assert query_obj.is_stored
def action_handler__get_query_params(query_id: str) -> ZMQReply:
"""
Handler for the 'get_query_params' action.
Returns query parameters of the query with the given `query_id`.
"""
q_info_lookup = QueryInfoLookup(Query.redis)
try:
query_params = q_info_lookup.get_query_params(query_id)
except UnkownQueryIdError:
payload = {"query_id": query_id, "query_state": "awol"}
return ZMQReply(
status="error", msg=f"Unknown query id: '{query_id}'", payload=payload
)
payload = {"query_id": query_id, "query_params": query_params}
return ZMQReply(status="success", payload=payload)
def _get_query_kind_for_query_id(query_id: str) -> Union[None, str]:
"""
Helper function to look up the query kind corresponding to the
given query id. Returns `None` if the query_id does not exist.
Parameters
----------
query_id : str
Identifier of the query.
Returns
-------
str or None
The query kind associated with this query_id (or None
if no query with this query_id exists).
"""
q_info_lookup = QueryInfoLookup(get_redis())
try:
return q_info_lookup.get_query_kind(query_id)
except UnkownQueryIdError:
return None
async def action_handler__get_geo_sql(config: "FlowmachineServerConfig",
query_id: str) -> ZMQReply:
"""
Handler for the 'get_sql' action.
Returns a SQL string which can be run against flowdb to obtain
the result of the query with given `query_id`.
"""
# TODO: currently we can't use QueryStateMachine to determine whether
# the query_id belongs to a valid query object, so we need to check it
# manually. Would be good to add a QueryState.UNKNOWN so that we can
# avoid this separate treatment.
q_info_lookup = QueryInfoLookup(get_redis())
if not q_info_lookup.query_is_known(query_id):
msg = f"Unknown query id: '{query_id}'"
payload = {"query_id": query_id, "query_state": "awol"}
return ZMQReply(status="error", msg=msg, payload=payload)
query_state = QueryStateMachine(get_redis(), query_id,
get_db().conn_id).current_query_state
if query_state == QueryState.COMPLETED:
q = get_query_object_by_id(get_db(), query_id)
try:
sql = q.geojson_query()
payload = {
"query_id": query_id,
"query_state": query_state,
"sql": sql,
"aggregation_unit": q.spatial_unit.canonical_name,
}
return ZMQReply(status="success", payload=payload)
except AttributeError:
msg = f"Query with id '{query_id}' has no geojson compatible representation." # TODO: This codepath is untested because all queries right now have geography
payload = {"query_id": query_id, "query_state": "errored"}
return ZMQReply(status="error", msg=msg, payload=payload)
else:
msg = f"Query with id '{query_id}' {query_state.description}."
payload = {"query_id": query_id, "query_state": query_state}
return ZMQReply(status="error", msg=msg, payload=payload)
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})
def __call__(self, value) -> Union[None, str]:
from flowmachine.core.server.query_schemas import FlowmachineQuerySchema
if (value is not None) and (value is not missing):
try:
(FlowmachineQuerySchema().load(
QueryInfoLookup(get_redis()).get_query_params(
value))._flowmachine_query_obj)
except UnkownQueryIdError:
if not cache_table_exists(get_db(), value):
raise ValidationError("Must be None or a valid query id.")
return value
def test_register_query(dummy_redis):
"""
Test that a query is known after registering it.
"""
q_info_lookup = QueryInfoLookup(dummy_redis)
assert not q_info_lookup.query_is_known("dummy_id")
q_info_lookup.register_query(
query_id="dummy_id",
query_params={"query_kind": "dummy_query", "dummy_param": "some_value"},
)
assert q_info_lookup.query_is_known("dummy_id")
def deserialize(
self,
value: typing.Any,
attr: str = None,
data: typing.Mapping[str, typing.Any] = None,
**kwargs,
) -> Union[None, Table]:
from flowmachine.core.server.query_schemas import FlowmachineQuerySchema
table_name = super().deserialize(value, attr, data, **kwargs)
if (table_name is missing) or (table_name is None):
return table_name
else:
try:
return (FlowmachineQuerySchema().load(
QueryInfoLookup(get_redis()).get_query_params(
value))._flowmachine_query_obj)
except UnkownQueryIdError:
return get_query_object_by_id(get_db(), value)
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})