repogen/internal/mongo/templates.go

package mongo

import (
	"fmt"
	"strings"

	"github.com/sunboyy/repogen/internal/code"
)

const baseTemplate = `// Code generated by repogen. DO NOT EDIT.
package {{.PackageName}}

import (
	"context"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
)

func New{{.InterfaceName}}(collection *mongo.Collection) {{.InterfaceName}} {
	return &{{.StructName}}{
		collection: collection,
	}
}

type {{.StructName}} struct {
	collection *mongo.Collection
}
`

type mongoBaseTemplateData struct {
	PackageName   string
	InterfaceName string
	StructName    string
}

const methodTemplate = `
func (r *{{.StructName}}) {{.MethodName}}(ctx context.Context, {{.Parameters}}){{.Returns}} {
{{.Implementation}}
}
`

type mongoMethodTemplateData struct {
	StructName     string
	MethodName     string
	ParamTypes     []code.Type
	ReturnTypes    []code.Type
	Implementation string
}

func (data mongoMethodTemplateData) Parameters() string {
	var params []string
	for i, paramType := range data.ParamTypes {
		params = append(params, fmt.Sprintf("arg%d %s", i, paramType.Code()))
	}
	return strings.Join(params, ", ")
}

func (data mongoMethodTemplateData) Returns() string {
	if len(data.ReturnTypes) == 0 {
		return ""
	}

	if len(data.ReturnTypes) == 1 {
		return fmt.Sprintf(" %s", data.ReturnTypes[0].Code())
	}

	var returns []string
	for _, returnType := range data.ReturnTypes {
		returns = append(returns, returnType.Code())
	}
	return fmt.Sprintf(" (%s)", strings.Join(returns, ", "))
}

const findOneTemplate = `	var entity {{.EntityType}}
	if err := r.collection.FindOne(ctx, bson.M{
{{range $index, $field := .Predicates}}		{{$field.Code $index}},
{{end}}	}).Decode(&entity); err != nil {
		return nil, err
	}
	return &entity, nil`

type mongoFindTemplateData struct {
	EntityType string
	Predicates []predicate
}

const findManyTemplate = `	cursor, err := r.collection.Find(ctx, bson.M{
{{range $index, $field := .Predicates}}		{{$field.Code $index}},
{{end}}	})
	if err != nil {
		return nil, err
	}
	var entities []*{{.EntityType}}
	if err := cursor.All(ctx, &entities); err != nil {
		return nil, err
	}
	return entities, nil`