Use @PartitionKey when loading related entities

Recently in Hibernate 6.2 @PartitionKey was introduced, which will add the partition key column to the where clause of update and delete statements. Is it also possible to add the partition key column to the where clause when loading a related entity or entity collection?

See the two entities below with one-to-many and many-to-one relationships:

@Entity
@Table(name = "contacts")
public class Contact {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "id")
    private Long id;

    @NotNull
    @PartitionKey
    @Column(name = "tenant_id", updatable = false)
    private Long tenantId;

    @NotNull
    @Column(name = "name")
    private String name;

    @OneToMany(mappedBy = "contact")
    private Set<ContactEmail> emails = new HashSet<>();

    public Set<ContactEmail> getEmails() {
        return emails;
    }
}

@Entity
@Table(name = "contact_emails")
public class ContactEmail {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "id")
    private Long id;

    @NotNull
    @PartitionKey
    @Column(name = "tenant_id", updatable = false)
    private Long tenantId;

    @NotNull
    @Column(name = "email")
    private String email;

    @NotNull
    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "contact_id")
    private Contact contact;

    public Contact getContact() {
        return contact;
    }
}

Example of accessing the entity collection:

// accessing the getter
contact.getEmails();

// today generates the following sql:
SELECT id, tenant_id, email, contact_id FROM contact_emails WHERE contact_id = ?;

// desired sql with partition key in where clause:
SELECT id, tenant_id, email, contact_id FROM contact_emails WHERE contact_id = ? AND tenant_id = ?;

Example of accessing the entity:

// accessing the getter
contactEmail.getContact();

// today generates the following sql:
SELECT id, tenant_id, name FROM contacts WHERE id = ?;

// desired sql with partition key in where clause:
SELECT id, tenant_id, name FROM contacts WHERE id = ? AND tenant_id = ?;

That’s not possible and probably also not very useful. If you want this to reliably work for joins and subselects, you should rather use the tenant_id as join column on your association.

Thanks for the quick response.

My use case: without the partition key in the where clause, Vitess will perform a scatter query, which means it queries all shards because it doesn’t know which specific shard to query. I am currently using hibernate, but in cases like this I have to avoid using the getters so I can avoid the scatter queries.

Also, why is it impossible?

If you want this to reliably work, add the tenant_id to your association’s join columns.

@jcarrig I need the same thing as you. How did you do it?