dialect/sql/sqlgraph: allow selecting order by terms (#3439)

dialect/sql/sqlgraph/graph.go

@@ -312,6 +312,9 @@ type (
 		// Terms used for non-aggregation ordering.
 		// See, OrderByNeighborTerms for more info.
 		Terms []OrderByTerm
+		// Selected indicates that the order terms
+		// should be appended to the query selection.
+		Selected bool
 	}
 	// OrderByInfo holds the information done by the OrderBy functions.
 	OrderByInfo struct {
@@ -383,6 +386,13 @@ func OrderByColumnDesc(c string) OrderByOption {
 	}
 }
 
+// OrderBySelected appends the ordered columns or terms with aliases to the query selection.
+func OrderBySelected() OrderByOption {
+	return func(opts *OrderByOptions) {
+		opts.Selected = true
+	}
+}
+
 // NewOrderBy gets list of options and returns a configured order-by.
 //
 //	NewOrderBy(
@@ -469,9 +479,9 @@ func OrderByNeighborsCount(q *sql.Selector, opts *OrderByOptions) {
 				q.C(s.From.Column),
 				join.C(pk1),
 			)
-		orderTerms(q, join, []OrderByTerm{
-			{Column: countC, Type: field.TypeInt, Desc: desc},
-		})
+		opts.Terms = []OrderByTerm{{As: countC, Type: field.TypeInt, Desc: desc}}
+		orderTerms(q, join, opts.Terms)
+		appendTerms(q, opts)
 	case s.ToEdgeOwner():
 		countC := countAlias(q, opts)
 		edgeT := build.Table(s.Edge.Table).Schema(s.Edge.Schema)
@@ -486,9 +496,9 @@ func OrderByNeighborsCount(q *sql.Selector, opts *OrderByOptions) {
 				q.C(s.From.Column),
 				join.C(s.Edge.Columns[0]),
 			)
-		orderTerms(q, join, []OrderByTerm{
-			{Column: countC, Type: field.TypeInt, Desc: desc},
-		})
+		opts.Terms = []OrderByTerm{{As: countC, Type: field.TypeInt, Desc: desc}}
+		orderTerms(q, join, opts.Terms)
+		appendTerms(q, opts)
 	}
 }
 
@@ -531,6 +541,20 @@ func selectTerms(q *sql.Selector, ts []OrderByTerm) {
 	}
 }
 
+func appendTerms(q *sql.Selector, opts *OrderByOptions) {
+	if !opts.Selected {
+		return
+	}
+	for _, t := range opts.Terms {
+		switch {
+		case t.As != "":
+			q.AppendSelect(t.As)
+		case t.Column != "":
+			q.AppendSelect(q.C(t.Column))
+		}
+	}
+}
+
 // OrderByNeighborTerms appends ordering based on the number of neighbors.
 // For example, order users by their number of posts.
 func OrderByNeighborTerms(q *sql.Selector, opts *OrderByOptions) {
@@ -571,6 +595,7 @@ func OrderByNeighborTerms(q *sql.Selector, opts *OrderByOptions) {
 			On(q.C(s.From.Column), join.C(s.Edge.Columns[0]))
 	}
 	orderTerms(q, join, opts.Terms)
+	appendTerms(q, opts)
 }
 
 type (

dialect/sql/sqlgraph/graph_test.go

@@ -969,6 +969,21 @@ func TestOrderByNeighborsCount(t *testing.T) {
 		require.Empty(t, args)
 		require.Equal(t, `SELECT "users"."name" FROM "users" ORDER BY "owner_id" IS NOT NULL`, query)
 	})
+	t.Run("Selected", func(t *testing.T) {
+		s := s.Clone()
+		OrderByNeighborsCount(s, NewOrderBy(
+			NewStep(
+				From("users", "id"),
+				To("pets", "owner_id"),
+				Edge(O2M, false, "pets", "owner_id"),
+			),
+			OrderByExprDesc(nil, "count_pets"),
+			OrderBySelected(),
+		))
+		query, args := s.Query()
+		require.Empty(t, args)
+		require.Equal(t, `SELECT "users"."name", "count_pets" FROM "users" LEFT JOIN (SELECT "pets"."owner_id", COUNT(*) AS "count_pets" FROM "pets" GROUP BY "pets"."owner_id") AS "t1" ON "users"."id" = "t1"."owner_id" ORDER BY "t1"."count_pets" DESC NULLS LAST`, query)
+	})
 }
 
 func TestOrderByNeighborTerms(t *testing.T) {
@@ -1028,6 +1043,28 @@ func TestOrderByNeighborTerms(t *testing.T) {
 		require.Empty(t, args)
 		require.Equal(t, `SELECT "users"."name" FROM "users" LEFT JOIN (SELECT "user_id", (SUM("num_users")) AS "total_users" FROM "group" JOIN "user_groups" AS "t1" ON "group"."id" = "t1"."group_id" GROUP BY "user_id") AS "t1" ON "users"."id" = "t1"."user_id" ORDER BY "t1"."total_users" NULLS FIRST`, query)
 	})
+	t.Run("Selected", func(t *testing.T) {
+		t1 := build.Table("users")
+		s := build.Select(t1.C("name")).
+			From(t1)
+		OrderByNeighborTerms(s, NewOrderBy(
+			NewStep(
+				From("users", "id"),
+				To("group", "id"),
+				Edge(M2M, false, "user_groups", "user_id", "group_id"),
+			),
+			OrderByExpr(
+				sql.ExprFunc(func(b *sql.Builder) {
+					b.S("SUM(").Ident("num_users").S(")")
+				}),
+				"total_users",
+			),
+			OrderBySelected(),
+		))
+		query, args := s.Query()
+		require.Empty(t, args)
+		require.Equal(t, `SELECT "users"."name", "total_users" FROM "users" LEFT JOIN (SELECT "user_id", (SUM("num_users")) AS "total_users" FROM "group" JOIN "user_groups" AS "t1" ON "group"."id" = "t1"."group_id" GROUP BY "user_id") AS "t1" ON "users"."id" = "t1"."user_id" ORDER BY "t1"."total_users" NULLS FIRST`, query)
+	})
 }
 
 func TestCreateNode(t *testing.T) {