order_by("pk") on a union() of a composite pk model crashes or sorts wrongly

[Akshat Gupta]

class Comment(models.Model):
    pk = models.CompositePrimaryKey("tenant_id", "comment_id")
    tenant_id = models.IntegerField()
    comment_id = models.IntegerField()

# rows (tenant_id, comment_id): (1, 2), (2, 1), (1, 1)
qs = Comment.objects.filter(tenant_id=1).union(Comment.objects.filter(tenant_id=2))

Ordering a combined query by a composite primary key fails in two ways.

1. On a normal queryset, order_by("pk") crashes. first() crashes too, because it orders by pk.

qs.first()
qs.order_by("pk")  # when evaluated
# TypeError: 'NoneType' object does not support item assignment

The error is raised in Query.add_select_col() (called from SQLCompiler.get_order_by()), because query.selected is None.

2. With values_list("pk") there is no error, but the order is wrong.

list(qs.values_list("pk").order_by("pk"))
# [((1, 1),), ((2, 1),), ((1, 2),)]
# expected: [((1, 1),), ((1, 2),), ((2, 1),)]

The SQL ends with ORDER BY "col1", and col1 is only the last pk column (comment_id). The same query without union() is ordered correctly.

I think get_order_by() treats the composite pk (ColPairs) as a single column in the combinator branch. It should order by each column of the pk.

Related: #36065, #37354 (other composite pk ordering bugs).