Skip to content

Instantly share code, notes, and snippets.

@samthor
Created Apr 10, 2020
Embed
What would you like to do?
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))
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment