Add query for flags and also clean up the code slightly