85 lines
1.9 KiB
Go
85 lines
1.9 KiB
Go
package api
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"net/http"
|
|
)
|
|
|
|
type ctxKey string
|
|
|
|
var traceKey ctxKey = "trace_id"
|
|
var sqlKey ctxKey = "sql"
|
|
var metaKey ctxKey = "req_meta"
|
|
var payloadKey ctxKey = "payload"
|
|
|
|
func withTrace(h http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
tid := r.Header.Get("X-Request-ID")
|
|
if tid == "" {
|
|
buf := make([]byte, 16)
|
|
_, _ = rand.Read(buf)
|
|
tid = hex.EncodeToString(buf)
|
|
}
|
|
w.Header().Set("X-Request-ID", tid)
|
|
m := ReqMeta{Method: r.Method, Path: r.URL.Path, Query: r.URL.RawQuery, Remote: r.RemoteAddr}
|
|
ctx := context.WithValue(r.Context(), traceKey, tid)
|
|
ctx = context.WithValue(ctx, metaKey, m)
|
|
h.ServeHTTP(w, r.WithContext(ctx))
|
|
})
|
|
}
|
|
|
|
func TraceIDFrom(r *http.Request) string {
|
|
v := r.Context().Value(traceKey)
|
|
if v == nil {
|
|
return ""
|
|
}
|
|
s, _ := v.(string)
|
|
return s
|
|
}
|
|
|
|
func WithSQL(r *http.Request, sql string) *http.Request {
|
|
return r.WithContext(context.WithValue(r.Context(), sqlKey, sql))
|
|
}
|
|
|
|
func SQLFrom(r *http.Request) string {
|
|
v := r.Context().Value(sqlKey)
|
|
if v == nil {
|
|
return ""
|
|
}
|
|
s, _ := v.(string)
|
|
return s
|
|
}
|
|
|
|
type ReqMeta struct {
|
|
Method string
|
|
Path string
|
|
Query string
|
|
Remote string
|
|
}
|
|
|
|
func MetaFrom(r *http.Request) ReqMeta {
|
|
v := r.Context().Value(metaKey)
|
|
if v == nil {
|
|
return ReqMeta{}
|
|
}
|
|
m, _ := v.(ReqMeta)
|
|
return m
|
|
}
|
|
|
|
func WithPayload(r *http.Request, v interface{}) *http.Request {
|
|
b, _ := json.Marshal(v)
|
|
return r.WithContext(context.WithValue(r.Context(), payloadKey, string(b)))
|
|
}
|
|
|
|
func PayloadFrom(r *http.Request) string {
|
|
v := r.Context().Value(payloadKey)
|
|
if v == nil {
|
|
return ""
|
|
}
|
|
s, _ := v.(string)
|
|
return s
|
|
}
|