Implement FindTopN to limit number of results in find many operations (#37)

internal/mongo/common.go

@@ -8,6 +8,8 @@ import (
 	"github.com/sunboyy/repogen/internal/spec"
 )
 
+var errOccurred = codegen.RawStatement("err != nil")
+
 var returnNilErr = codegen.ReturnStatement{
 	codegen.Identifier("nil"),
 	codegen.Identifier("err"),
@@ -15,7 +17,7 @@ var returnNilErr = codegen.ReturnStatement{
 
 var ifErrReturnNilErr = codegen.IfBlock{
 	Condition: []codegen.Statement{
-		codegen.RawStatement("err != nil"),
+		errOccurred,
 	},
 	Statements: []codegen.Statement{
 		returnNilErr,
@@ -24,7 +26,7 @@ var ifErrReturnNilErr = codegen.IfBlock{
 
 var ifErrReturn0Err = codegen.IfBlock{
 	Condition: []codegen.Statement{
-		codegen.RawStatement("err != nil"),
+		errOccurred,
 	},
 	Statements: []codegen.Statement{
 		codegen.ReturnStatement{
@@ -36,7 +38,7 @@ var ifErrReturn0Err = codegen.IfBlock{
 
 var ifErrReturnFalseErr = codegen.IfBlock{
 	Condition: []codegen.Statement{
-		codegen.RawStatement("err != nil"),
+		errOccurred,
 	},
 	Statements: []codegen.Statement{
 		codegen.ReturnStatement{

internal/mongo/find.go

@@ -1,6 +1,8 @@
 package mongo
 
 import (
+	"strconv"
+
 	"github.com/sunboyy/repogen/internal/code"
 	"github.com/sunboyy/repogen/internal/codegen"
 	"github.com/sunboyy/repogen/internal/spec"
@@ -66,7 +68,7 @@ func (g findBodyGenerator) generateFindOneBody(querySpec querySpec,
 							).Build(),
 					},
 				},
-				codegen.RawStatement("err != nil"),
+				errOccurred,
 			},
 			Statements: []codegen.Statement{
 				returnNilErr,
@@ -91,10 +93,7 @@ func (g findBodyGenerator) generateFindManyBody(querySpec querySpec,
 					Call("Find",
 						codegen.Identifier("arg0"),
 						querySpec.Code(),
-						codegen.NewChainBuilder("options").
-							Call("Find").
-							Call("SetSort", sortsCode).
-							Build(),
+						g.findManyOptions(sortsCode),
 					).Build(),
 			},
 		},
@@ -119,7 +118,7 @@ func (g findBodyGenerator) generateFindManyBody(querySpec querySpec,
 							).Build(),
 					},
 				},
-				codegen.RawStatement("err != nil"),
+				errOccurred,
 			},
 			Statements: []codegen.Statement{
 				returnNilErr,
@@ -132,6 +131,21 @@ func (g findBodyGenerator) generateFindManyBody(querySpec querySpec,
 	}
 }
 
+func (g findBodyGenerator) findManyOptions(
+	sortsCode codegen.MapStatement) codegen.Statement {
+
+	optionsBuilder := codegen.NewChainBuilder("options").
+		Call("Find").
+		Call("SetSort", sortsCode)
+	if g.operation.Limit > 0 {
+		optionsBuilder = optionsBuilder.Call("SetLimit",
+			codegen.Identifier(strconv.Itoa(g.operation.Limit)),
+		)
+	}
+
+	return optionsBuilder.Build()
+}
+
 func (g findBodyGenerator) generateSortMap() (
 	codegen.MapStatement, error) {
 

internal/mongo/find_test.go

@@ -825,6 +825,38 @@ func TestGenerateMethod_Find(t *testing.T) {
 	if err := cursor.All(arg0, &entities); err != nil {
 		return nil, err
 	}
+	return entities, nil`,
+		},
+		{
+			Name: "find with limit",
+			MethodSpec: spec.MethodSpec{
+				Name: "FindTop5AllOrderByAgeDesc",
+				Params: []code.Param{
+					{Name: "ctx", Type: code.ExternalType{PackageAlias: "context", Name: "Context"}},
+				},
+				Returns: []code.Type{
+					code.ArrayType{ContainedType: code.PointerType{ContainedType: code.SimpleType("UserModel")}},
+					code.TypeError,
+				},
+				Operation: spec.FindOperation{
+					Mode: spec.QueryModeMany,
+					Sorts: []spec.Sort{
+						{FieldReference: spec.FieldReference{ageField}, Ordering: spec.OrderingDescending},
+					},
+					Limit: 5,
+				},
+			},
+			ExpectedBody: `	cursor, err := r.collection.Find(arg0, bson.M{
+	}, options.Find().SetSort(bson.M{
+		"age": -1,
+	}).SetLimit(5))
+	if err != nil {
+		return nil, err
+	}
+	var entities []*UserModel
+	if err := cursor.All(arg0, &entities); err != nil {
+		return nil, err
+	}
 	return entities, nil`,
 		},
 	}