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_get_query_object_by_id(flowmachine_connect):
"""
Test that we can get a query object back out of the database by the query's id
"""
dl = daily_location("2016-01-01").store().result()
retrieved_query = get_query_object_by_id(flowmachine_connect, dl.query_id)
assert dl.query_id == retrieved_query.query_id
assert dl.get_query() == retrieved_query.get_query()
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 get_sql_for_query_id(query_id):
"""
Return the SQL which, when run against flowdb, will
return the result for the query with the given id.
Parameters
----------
query_id : str
The query id
Returns
-------
str
"""
q = get_query_object_by_id(Query.connection, query_id)
sql = q.get_query()
return sql
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__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)
async def test_rerun_query_after_removed_from_cache(dummy_redis, server_config,
real_connections):
"""
Test that a query can be rerun after it has been removed from the cache.
"""
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",
),
)
query_id = msg["payload"]["query_id"]
qsm = QueryStateMachine(get_redis(), query_id, get_db().conn_id)
qsm.wait_until_complete()
query_obj = get_query_object_by_id(get_db(), query_id)
assert query_obj.is_stored
query_obj.invalidate_db_cache()
assert not query_obj.is_stored
assert qsm.is_known
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