Added partial handling for associations