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