main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/hexdigest/gowrap/generator"
	"github.com/pkg/errors"
)

var (
	interfaceName = flag.String("i", "", "interface name")
	structName    = flag.String("s", "", "target struct name, default: <interface name>Tracer")
	outputFile    = flag.String("o", "", "output filename")
)

const headerTemplate = `
	// Code generated by http://github.com/gojuno/tracegen. DO NOT EDIT.

	package {{$.Package.Name}}
`

const bodyTemplate = `
	import (
		opentracing "github.com/opentracing/opentracing-go"
		_ext "github.com/opentracing/opentracing-go/ext"
		_log "github.com/opentracing/opentracing-go/log"
	)

	{{ $decorator := (or .Vars.DecoratorName (printf "%sTracer" .Interface.Name)) }}

	// {{$decorator}} instruments {{.Interface.Type}} with opentracing spans
	type {{$decorator}} struct {
		{{.Interface.Type}}
		_instance string
	}

	// New{{$decorator}} returns {{$decorator}}
	func New{{$decorator}} (base {{.Interface.Type}}, instance string) {{$decorator}} {
		return {{$decorator}} {
			{{.Interface.Name}}: base,
			_instance: instance,
		}
	}

	{{range $method := .Interface.Methods}}
		{{if $method.AcceptsContext}}
			// {{$method.Name}} implements {{$.Interface.Type}}
			func (_d {{$decorator}}) {{$method.Declaration}} {
				_span, ctx := opentracing.StartSpanFromContext(ctx, _d._instance + ".{{$.Interface.Type}}.{{$method.Name}}")
				defer func() {
					{{if $method.ReturnsError}}if err != nil {
						_ext.Error.Set(_span, true)
						_span.LogFields(
							_log.String("event", "error"),
							_log.String("message", err.Error()),
						)
					}
					{{end}}
					_span.Finish()
				}()
				{{$method.Pass (printf "_d.%s." $.Interface.Name) }}
			}
		{{end}}
	{{end}}
`

func main() {
	flag.Parse()

	if *interfaceName == "" || *outputFile == "" || flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}

	if *structName == "" {
		*structName = fmt.Sprintf("%vTracer", *interfaceName)
	}

	opts := generator.Options{
		InterfaceName: *interfaceName,
		SourcePackage: flag.Arg(0),
		//SourcePackageAlias: "mm_" + sourcePackage.Name,
		HeaderTemplate: headerTemplate,
		BodyTemplate:   bodyTemplate,
		Vars: map[string]interface{}{
			"DecoratorName": structName,
		},
		OutputFile: *outputFile,
	}

	if err := generate(opts); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(1)
	}
}

func generate(o generator.Options) error {
	g, err := generator.NewGenerator(o)
	if err != nil {
		return err
	}

	buf := bytes.NewBuffer([]byte{})

	if err = g.Generate(buf); err != nil {
		return errors.Wrap(err, "failed to generate decorator")
	}

	return ioutil.WriteFile(o.OutputFile, buf.Bytes(), 0644)
}