spanner-opencensus.go

// Sample program for OpenCensus tracing in the Cloud Spanner client.
// The database specified by the flags should have two columns,
// Key and Value, both strings.

package main

import (
	"flag"
	"fmt"
	"log"
	"strings"
	"time"
	"cloud.google.com/go/spanner"
	"go.opencensus.io/trace"
	"golang.org/x/net/context"
	"google.golang.org/api/option"
)

var (
	projectID = flag.String("project", "", "project ID")
	instance  = flag.String("instance", "", "spanner instance")
	dbName    = flag.String("db", "", "spanner database")
	keyFile   = flag.String("keyfile", "", "file with credentials")
)

func main() {
	flag.Parse()
	exp := &exporter{startTime: time.Now()}
	trace.RegisterExporter(exp)
	defer exp.Dump()
        trace.SetDefaultSampler(trace.AlwaysSample())
        ctx := context.Background()

	db := fmt.Sprintf("projects/%v/instances/%v/databases/%v", *projectID, *instance, *dbName)
	client, err := spanner.NewClient(ctx, db, option.WithCredentialsFile(*keyFile))
	if err != nil {
		log.Fatal(err)
	}
	defer client.Close()
	const table = "T"
	columns := []string{"Key", "Value"}
	var ms []*spanner.Mutation
	for i := 0; i < 1; i++ {
		ms = append(ms, spanner.InsertOrUpdate(table,
			columns,
			[]interface{}{fmt.Sprintf("k%d", i), fmt.Sprintf("v%d", i)}))
	}
	if _, err := client.Apply(ctx, ms, spanner.ApplyAtLeastOnce()); err != nil {
		log.Fatal(err)
	}
	rng := spanner.KeyRange{Start: spanner.Key{"k1"}, End: spanner.Key{"k2"}}
	err = client.Single().Read(ctx, table, rng, columns).Do(func(r *spanner.Row) error {
		var k, v string
		if err := r.Columns(&k, &v); err != nil {
			return err
		}
		fmt.Printf("%s: %s\n", k, v)
		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
}

type exporter struct {
	startTime time.Time
	data      []*trace.SpanData
}

func (e *exporter) Export(sd *trace.SpanData) {
	e.data = append(e.data, sd)
}

func (e *exporter) Dump() {
	roots := e.spansWithParent(trace.SpanID{})
	for _, r := range roots {
		e.dumpTree(r, 0)
	}
}

func (e *exporter) dumpTree(sd *trace.SpanData, level int) {
	start := int(sd.StartTime.Sub(e.startTime) / time.Millisecond)
	end := int(sd.EndTime.Sub(e.startTime) / time.Millisecond)
	indent := strings.Repeat("    ", level)
	fmt.Printf("%6d\t%s> %s\n", start, indent, sd.Name)
	for _, c := range e.spansWithParent(sd.SpanID) {
		e.dumpTree(c, level+1)
	}
	fmt.Printf("%6d\t%s< %s (%.1fs)\n",
		end, indent, sd.Name, sd.EndTime.Sub(sd.StartTime).Seconds())
}

func (e *exporter) spansWithParent(p trace.SpanID) []*trace.SpanData {
	var res []*trace.SpanData
	for _, sd := range e.data {
		if sd.ParentSpanID == p {
			res = append(res, sd)
		}
	}
	return res
}