context.go 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. package context
  2. import (
  3. "context"
  4. )
  5. // 定义全局上下文中的键
  6. type (
  7. transCtx struct{}
  8. transLockCtx struct{}
  9. userIDCtx struct{}
  10. traceIDCtx struct{}
  11. )
  12. // NewTrans 创建事务的上下文
  13. func NewTrans(ctx context.Context, trans interface{}) context.Context {
  14. return context.WithValue(ctx, transCtx{}, trans)
  15. }
  16. // FromTrans 从上下文中获取事务
  17. func FromTrans(ctx context.Context) (interface{}, bool) {
  18. v := ctx.Value(transCtx{})
  19. return v, v != nil
  20. }
  21. // NewTransLock 创建事务锁的上下文
  22. func NewTransLock(ctx context.Context) context.Context {
  23. return context.WithValue(ctx, transLockCtx{}, struct{}{})
  24. }
  25. // FromTransLock 从上下文中获取事务锁
  26. func FromTransLock(ctx context.Context) bool {
  27. v := ctx.Value(transLockCtx{})
  28. return v != nil
  29. }
  30. // NewUserID 创建用户ID的上下文
  31. func NewUserID(ctx context.Context, userID string) context.Context {
  32. return context.WithValue(ctx, userIDCtx{}, userID)
  33. }
  34. // FromUserID 从上下文中获取用户ID
  35. func FromUserID(ctx context.Context) (string, bool) {
  36. v := ctx.Value(userIDCtx{})
  37. if v != nil {
  38. if s, ok := v.(string); ok {
  39. return s, s != ""
  40. }
  41. }
  42. return "", false
  43. }
  44. // NewTraceID 创建追踪ID的上下文
  45. func NewTraceID(ctx context.Context, traceID string) context.Context {
  46. return context.WithValue(ctx, traceIDCtx{}, traceID)
  47. }
  48. // FromTraceID 从上下文中获取追踪ID
  49. func FromTraceID(ctx context.Context) (string, bool) {
  50. v := ctx.Value(traceIDCtx{})
  51. if v != nil {
  52. if s, ok := v.(string); ok {
  53. return s, s != ""
  54. }
  55. }
  56. return "", false
  57. }