N+1 query fix and code cleanup