repogen/internal/mongo/templates.go

package mongo

import (
	"fmt"
	"strings"

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

const constructorTemplate = `
import (
	"context"

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

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

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

type mongoConstructorTemplateData struct {
	InterfaceName  string
	ImplStructName string
}

const methodTemplate = `
func (r *{{.StructName}}) {{.MethodName}}({{.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 insertOneTemplate = `	result, err := r.collection.InsertOne(arg0, arg1)
	if err != nil {
		return nil, err
	}
	return result.InsertedID, nil`

const insertManyTemplate = `	var entities []interface{}
	for _, model := range arg1 {
		entities = append(entities, model)
	}
	result, err := r.collection.InsertMany(arg0, entities)
	if err != nil {
		return nil, err
	}
	return result.InsertedIDs, nil`

type mongoFindTemplateData struct {
	EntityType string
	QuerySpec  querySpec
	Sorts      []findSort
}

type findSort struct {
	BsonTag  string
	Ordering spec.Ordering
}

func (s findSort) OrderNum() int {
	if s.Ordering == spec.OrderingAscending {
		return 1
	}
	return -1
}

const findOneTemplate = `	var entity {{.EntityType}}
	if err := r.collection.FindOne(arg0, bson.M{
{{.QuerySpec.Code}}
	}, options.FindOne().SetSort(bson.M{
{{range $index, $element := .Sorts}}		"{{$element.BsonTag}}": {{$element.OrderNum}},
{{end}}	})).Decode(&entity); err != nil {
		return nil, err
	}
	return &entity, nil`

const findManyTemplate = `	cursor, err := r.collection.Find(arg0, bson.M{
{{.QuerySpec.Code}}
	}, options.Find().SetSort(bson.M{
{{range $index, $element := .Sorts}}		"{{$element.BsonTag}}": {{$element.OrderNum}},
{{end}}	}))
	if err != nil {
		return nil, err
	}
	var entities []*{{.EntityType}}
	if err := cursor.All(arg0, &entities); err != nil {
		return nil, err
	}
	return entities, nil`

type mongoUpdateTemplateData struct {
	Update    update
	QuerySpec querySpec
}

const updateOneTemplate = `	result, err := r.collection.UpdateOne(arg0, bson.M{
{{.QuerySpec.Code}}
	}, bson.M{
{{.Update.Code}}
	})
	if err != nil {
		return false, err
	}
	return result.MatchedCount > 0, err`

const updateManyTemplate = `	result, err := r.collection.UpdateMany(arg0, bson.M{
{{.QuerySpec.Code}}
	}, bson.M{
{{.Update.Code}}
	})
	if err != nil {
		return 0, err
	}
	return int(result.MatchedCount), err`

type mongoDeleteTemplateData struct {
	QuerySpec querySpec
}

const deleteOneTemplate = `	result, err := r.collection.DeleteOne(arg0, bson.M{
{{.QuerySpec.Code}}
	})
	if err != nil {
		return false, err
	}
	return result.DeletedCount > 0, nil`

const deleteManyTemplate = `	result, err := r.collection.DeleteMany(arg0, bson.M{
{{.QuerySpec.Code}}
	})
	if err != nil {
		return 0, err
	}
	return int(result.DeletedCount), nil`

type mongoCountTemplateData struct {
	QuerySpec querySpec
}

const countTemplate = `	count, err := r.collection.CountDocuments(arg0, bson.M{
{{.QuerySpec.Code}}
	})
	if err != nil {
		return 0, err
	}
	return int(count), nil`