Allow sorting by $natural (#3822)

Closes #3638.

integration/query_compat_test.go

@@ -241,6 +241,7 @@ func TestQueryCappedCollectionCompat(t *testing.T) {
 	for name, tc := range map[string]struct {
 		filter bson.D
 		sort   bson.D
+		skip   string
 
 		sortPushdown resultPushdown
 	}{
@@ -264,12 +265,45 @@ func TestQueryCappedCollectionCompat(t *testing.T) {
 			sort:         bson.D{{"v", 1}, {"_id", int32(-1)}},
 			sortPushdown: noPushdown,
 		},
+
+		"SortNaturalAsc": {
+			sort:         bson.D{{"$natural", int32(1)}},
+			sortPushdown: allPushdown,
+		},
+		"SortNaturalDesc": {
+			sort:         bson.D{{"$natural", int32(-1)}},
+			sortPushdown: allPushdown,
+		},
+		"SortNaturalInt64": {
+			sort:         bson.D{{"$natural", int64(1)}},
+			sortPushdown: allPushdown,
+		},
+
+		"SortNaturalZero": {
+			skip:         "https://github.com/FerretDB/FerretDB/issues/3638",
+			sort:         bson.D{{"$natural", int32(0)}},
+			sortPushdown: noPushdown,
+		},
+		"SortNaturalString": {
+			skip:         "https://github.com/FerretDB/FerretDB/issues/3638",
+			sort:         bson.D{{"$natural", "foo"}},
+			sortPushdown: noPushdown,
+		},
+		"SortNaturalMultipleSorts": {
+			skip:         "https://github.com/FerretDB/FerretDB/issues/3638",
+			sort:         bson.D{{"$natural", int32(1)}, {"v", int32(1)}},
+			sortPushdown: noPushdown,
+		},
 	} {
 		name, tc := name, tc
 
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
+			if tc.skip != "" {
+				t.Skip(tc.skip)
+			}
+
 			explainQuery := bson.D{
 				{"find", targetCollection.Name()},
 			}
@@ -282,13 +316,6 @@ func TestQueryCappedCollectionCompat(t *testing.T) {
 				explainQuery = append(explainQuery, bson.E{Key: "sort", Value: tc.sort})
 			}
 
-			var explainRes bson.D
-			require.NoError(t, targetCollection.Database().RunCommand(ctx, bson.D{{"explain", explainQuery}}).Decode(&explainRes))
-
-			doc := ConvertDocument(t, explainRes)
-			sortPushdown, _ := doc.Get("sortPushdown")
-			assert.Equal(t, tc.sortPushdown.PushdownExpected(t), sortPushdown)
-
 			findOpts := options.Find()
 			if tc.sort != nil {
 				findOpts.SetSort(tc.sort)
@@ -300,10 +327,17 @@ func TestQueryCappedCollectionCompat(t *testing.T) {
 			}
 
 			targetCursor, targetErr := targetCollection.Find(ctx, filter, findOpts)
-			require.NoError(t, targetErr)
-
 			compatCursor, compatErr := compatCollection.Find(ctx, filter, findOpts)
-			require.NoError(t, compatErr)
+			if targetErr != nil {
+				t.Logf("Target error: %v", targetErr)
+				t.Logf("Compat error: %v", compatErr)
+
+				// error messages are intentionally not compared
+				AssertMatchesCommandError(t, compatErr, targetErr)
+
+				return
+			}
+			require.NoError(t, compatErr, "compat error; target returned no error")
 
 			var targetFindRes []bson.D
 			targetErr = targetCursor.All(ctx, &targetFindRes)
@@ -320,6 +354,13 @@ func TestQueryCappedCollectionCompat(t *testing.T) {
 			for i := range compatFindRes {
 				AssertEqualDocuments(t, compatFindRes[i], targetFindRes[i])
 			}
+
+			var explainRes bson.D
+			require.NoError(t, targetCollection.Database().RunCommand(ctx, bson.D{{"explain", explainQuery}}).Decode(&explainRes))
+
+			doc := ConvertDocument(t, explainRes)
+			sortPushdown, _ := doc.Get("sortPushdown")
+			assert.Equal(t, tc.sortPushdown.PushdownExpected(t), sortPushdown)
 		})
 	}
 }

internal/backends/collection.go

@@ -17,13 +17,10 @@ package backends
 import (
 	"cmp"
 	"context"
-	"errors"
 	"slices"
 	"time"
 
 	"github.com/FerretDB/FerretDB/internal/types"
-	"github.com/FerretDB/FerretDB/internal/util/iterator"
-	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
 	"github.com/FerretDB/FerretDB/internal/util/must"
 	"github.com/FerretDB/FerretDB/internal/util/observability"
 )
@@ -161,24 +158,11 @@ func (cc *collectionContract) Explain(ctx context.Context, params *ExplainParams
 	}
 
 	if params.Sort.Len() != 0 {
-		iter := params.Sort.Iterator()
-		defer iter.Close()
-
-		for {
-			_, v, err := iter.Next()
-			if err != nil {
-				if errors.Is(err, iterator.ErrIteratorDone) {
-					break
-				}
-
-				return nil, lazyerrors.Error(err)
-			}
-
-			sortValue := v.(int64)
+		must.BeTrue(params.Sort.Len() == 1)
+		sortValue := params.Sort.Map()["$natural"].(int64)
 
-			if sortValue != -1 && sortValue != 1 {
-				panic("sort key ordering must be 1 (for ascending) or -1 (for descending)")
-			}
+		if sortValue != -1 && sortValue != 1 {
+			panic("sort value must be 1 (for ascending) or -1 (for descending)")
 		}
 	}
 

internal/backends/collection_test.go

@@ -73,6 +73,11 @@ func TestCollectionInsertAllQueryExplain(t *testing.T) {
 				must.NotFail(types.NewDocument("_id", types.ObjectID{1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})),
 			}
 
+			invertedDocs := make([]*types.Document, len(insertDocs))
+			copy(invertedDocs, insertDocs)
+
+			slices.Reverse(invertedDocs)
+
 			_, err = coll.InsertAll(ctx, &backends.InsertAllParams{Docs: insertDocs})
 			require.NoError(t, err)
 
@@ -102,6 +107,30 @@ func TestCollectionInsertAllQueryExplain(t *testing.T) {
 				assert.True(t, explainRes.SortPushdown)
 			})
 
+			t.Run("CappedCollectionDesc", func(t *testing.T) {
+				t.Parallel()
+
+				sort := must.NotFail(types.NewDocument("$natural", int64(-1)))
+
+				queryRes, err := cappedColl.Query(ctx, &backends.QueryParams{
+					Sort: sort,
+				})
+				require.NoError(t, err)
+
+				docs, err := iterator.ConsumeValues[struct{}, *types.Document](queryRes.Iter)
+				require.NoError(t, err)
+				require.Len(t, docs, len(invertedDocs))
+				testutil.AssertEqualSlices(t, invertedDocs, docs)
+
+				assertEqualRecordID(t, invertedDocs, docs)
+
+				explainRes, err := cappedColl.Explain(ctx, &backends.ExplainParams{
+					Sort: sort,
+				})
+				require.NoError(t, err)
+				assert.True(t, explainRes.SortPushdown)
+			})
+
 			t.Run("CappedCollectionOnlyRecordIDs", func(t *testing.T) {
 				t.Parallel()
 

internal/backends/postgresql/query.go

@@ -237,14 +237,18 @@ func prepareOrderByClause(p *metadata.Placeholder, sort *types.Document) (string
 	}
 
 	v := must.NotFail(sort.Get("$natural"))
+	var order string
 
-	// TODO https://github.com/FerretDB/FerretDB/issues/3638
-	sortOrder := v.(int64)
-	if sortOrder != 1 {
-		return "", nil
+	switch v.(int64) {
+	case 1:
+		// Ascending order
+	case -1:
+		order = " DESC"
+	default:
+		panic("not reachable")
 	}
 
-	return fmt.Sprintf(" ORDER BY %s", metadata.RecordIDColumn), nil
+	return fmt.Sprintf(" ORDER BY %s%s", metadata.RecordIDColumn, order), nil
 }
 
 // filterEqual returns the proper SQL filter with arguments that filters documents

internal/backends/postgresql/query_test.go

@@ -354,9 +354,8 @@ func TestPrepareOrderByClause(t *testing.T) {
 			orderBy: ` ORDER BY _ferretdb_record_id`,
 		},
 		"NaturalDescending": {
-			skip:    "https://github.com/FerretDB/FerretDB/issues/3638",
 			sort:    must.NotFail(types.NewDocument("$natural", int64(-1))),
-			orderBy: "",
+			orderBy: ` ORDER BY _ferretdb_record_id DESC`,
 		},
 	} {
 		name, tc := name, tc

internal/backends/sqlite/query.go

@@ -56,12 +56,16 @@ func prepareOrderByClause(sort *types.Document) string {
 	}
 
 	v := must.NotFail(sort.Get("$natural"))
+	var order string
 
-	// TODO https://github.com/FerretDB/FerretDB/issues/3638
-	sortOrder := v.(int64)
-	if sortOrder != 1 {
-		return ""
+	switch v.(int64) {
+	case 1:
+		// Ascending order
+	case -1:
+		order = " DESC"
+	default:
+		panic("not reachable")
 	}
 
-	return fmt.Sprintf(` ORDER BY %s`, metadata.RecordIDColumn)
+	return fmt.Sprintf(" ORDER BY %s%s", metadata.RecordIDColumn, order)
 }

internal/backends/sqlite/query_test.go

@@ -102,7 +102,7 @@ func TestPrepareOrderByClause(t *testing.T) {
 		},
 		"NaturalDescending": {
 			sort:    must.NotFail(types.NewDocument("$natural", int64(-1))),
-			orderBy: "",
+			orderBy: ` ORDER BY _ferretdb_record_id DESC`,
 		},
 	} {
 		name, tc := name, tc

internal/handler/common/sort.go

@@ -43,13 +43,19 @@ func SortDocuments(docs []*types.Document, sortDoc *types.Document) error {
 
 	for i, sortKey := range sortDoc.Keys() {
 		fields := strings.Split(sortKey, ".")
-		for _, field := range fields {
-			if strings.HasPrefix(field, "$") {
-				return handlererrors.NewCommandErrorMsgWithArgument(
-					handlererrors.ErrFieldPathInvalidName,
-					"FieldPath field names may not start with '$'. Consider using $getField or $setField.",
-					"sort",
-				)
+
+		switch {
+		case sortKey == "$natural":
+		default:
+			// TODO https://github.com/FerretDB/FerretDB/issues/3127
+			for _, field := range fields {
+				if strings.HasPrefix(field, "$") {
+					return handlererrors.NewCommandErrorMsgWithArgument(
+						handlererrors.ErrFieldPathInvalidName,
+						"FieldPath field names may not start with '$'. Consider using $getField or $setField.",
+						"sort",
+					)
+				}
 			}
 		}
 
@@ -94,13 +100,19 @@ func ValidateSortDocument(sortDoc *types.Document) (*types.Document, error) {
 
 	for _, sortKey := range sortDoc.Keys() {
 		fields := strings.Split(sortKey, ".")
-		for _, field := range fields {
-			if strings.HasPrefix(field, "$") {
-				return nil, handlererrors.NewCommandErrorMsgWithArgument(
-					handlererrors.ErrFieldPathInvalidName,
-					"FieldPath field names may not start with '$'. Consider using $getField or $setField.",
-					"sort",
-				)
+
+		switch {
+		case sortKey == "$natural":
+		default:
+			// TODO https://github.com/FerretDB/FerretDB/issues/3127
+			for _, field := range fields {
+				if strings.HasPrefix(field, "$") {
+					return nil, handlererrors.NewCommandErrorMsgWithArgument(
+						handlererrors.ErrFieldPathInvalidName,
+						"FieldPath field names may not start with '$'. Consider using $getField or $setField.",
+						"sort",
+					)
+				}
 			}
 		}
 

internal/handler/msg_aggregate.go

@@ -300,14 +300,26 @@ func (h *Handler) MsgAggregate(ctx context.Context, msg *wire.OpMsg) (*wire.OpMs
 			cInfo = cList.Collections[i]
 		}
 
-		capped := cInfo.Capped()
-
 		switch {
 		case h.DisablePushdown:
 			// Pushdown disabled
-		case sort.Len() == 0 && capped:
+		case sort.Len() == 0 && cInfo.Capped():
 			// Pushdown default recordID sorting for capped collections
 			qp.Sort = must.NotFail(types.NewDocument("$natural", int64(1)))
+		case sort.Len() == 1:
+			if sort.Keys()[0] != "$natural" {
+				break
+			}
+
+			if !cInfo.Capped() {
+				return nil, handlererrors.NewCommandErrorMsgWithArgument(
+					handlererrors.ErrNotImplemented,
+					"$natural sort for non-capped collection is not supported.",
+					"aggregate",
+				)
+			}
+
+			qp.Sort = sort
 		}
 
 		iter, err = processStagesDocuments(ctx, closer, &stagesDocumentsParams{c, qp, stagesDocuments})

internal/handler/msg_explain.go

@@ -120,14 +120,26 @@ func (h *Handler) MsgExplain(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg,
 		cInfo = cList.Collections[i]
 	}
 
-	capped := cInfo.Capped()
-
 	switch {
 	case h.DisablePushdown:
 		// Pushdown disabled
-	case params.Sort.Len() == 0 && capped:
+	case params.Sort.Len() == 0 && cInfo.Capped():
 		// Pushdown default recordID sorting for capped collections
 		qp.Sort = must.NotFail(types.NewDocument("$natural", int64(1)))
+	case params.Sort.Len() == 1:
+		if params.Sort.Keys()[0] != "$natural" {
+			break
+		}
+
+		if !cInfo.Capped() {
+			return nil, handlererrors.NewCommandErrorMsgWithArgument(
+				handlererrors.ErrNotImplemented,
+				"$natural sort for non-capped collection is not supported.",
+				"explain",
+			)
+		}
+
+		qp.Sort = params.Sort
 	}
 
 	// Limit pushdown is not applied if:

internal/handler/msg_find.go

@@ -216,6 +216,20 @@ func (h *Handler) makeFindQueryParams(params *common.FindParams, cInfo *backends
 	case params.Sort.Len() == 0 && cInfo.Capped():
 		// Pushdown default recordID sorting for capped collections
 		qp.Sort = must.NotFail(types.NewDocument("$natural", int64(1)))
+	case params.Sort.Len() == 1:
+		if params.Sort.Keys()[0] != "$natural" {
+			break
+		}
+
+		if !cInfo.Capped() {
+			return nil, handlererrors.NewCommandErrorMsgWithArgument(
+				handlererrors.ErrNotImplemented,
+				"$natural sort for non-capped collection is not supported.",
+				"find",
+			)
+		}
+
+		qp.Sort = params.Sort
 	}
 
 	// Limit pushdown is not applied if: