Refactor Query generation logic