wrap.go

package loggerwrap

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"strings"

	"cloud.google.com/go/logging"
	mrpb "google.golang.org/genproto/googleapis/api/monitoredres"
)

var (
	projectName   = os.Getenv("GOOGLE_CLOUD_PROJECT")
	isProd        = (projectName != "")
	loggingClient *logging.Client
	pendingError  error

	requestLogger *logging.Logger
	eventLogger   *logging.Logger
)

const (
	loggerKey = "it's a gae logger because they removed the handy one"
)

func init() {
	if isProd {
		ctx := context.Background()
		loggingClient, pendingError = logging.NewClient(ctx, "projects/"+projectName)

		res := logging.CommonResource(&mrpb.MonitoredResource{
			Type: "gae_app",
		})

		requestLogger = loggingClient.Logger("request", res)
		eventLogger = loggingClient.Logger("event", res)
	}
}

type interceptingResponseWriter struct {
	w     http.ResponseWriter
	total int
}

func (irw *interceptingResponseWriter) Header() http.Header {
	return irw.w.Header()
}

func (irw *interceptingResponseWriter) Write(b []byte) (int, error) {
	count, err := irw.w.Write(b)
	if count > 0 {
		irw.total += count
	}
	return count, err
}

func (irw *interceptingResponseWriter) WriteHeader(statusCode int) {
	irw.w.WriteHeader(statusCode)
}

func internalLog(c context.Context, severity logging.Severity, format string, rest ...interface{}) {
	requestLogger.Log(logging.Entry{
		Severity: severity,
		Payload:  fmt.Sprintf(format, rest...),
	})
}

// Infof logs a message with the "INFO" level.
func Infof(c context.Context, format string, rest ...interface{}) {
	internalLog(c, logging.Info, format, rest...)
}

// Wrap wraps the given http.HandlerFunc so that logging calls into this package will operate correctly on the active trace.
func Wrap(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		traceID := strings.Split(r.Header.Get("X-Cloud-Trace-Context"), "/")
		ctx := context.WithValue(r.Context(), loggerKey, traceID)
		irw := &interceptingResponseWriter{w: w}
		handler(irw, r.WithContext(ctx))
	}
}