context.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. package context
  2. import (
  3. "context"
  4. )
  5. // 定义全局上下文中的键
  6. type (
  7. transCtx struct{}
  8. transLockCtx struct{}
  9. userIDCtx struct{}
  10. traceIDCtx struct{}
  11. fileSizeLimitKey struct{}
  12. fileNameKey struct{}
  13. fileExpireKey struct{}
  14. fileHashKey struct{}
  15. // FileSizeLimitHandle file size limit
  16. FileSizeLimitHandle func() bool
  17. // FileNameHandle the file name
  18. FileNameHandle func(fileName string) string
  19. )
  20. // NewTrans 创建事务的上下文
  21. func NewTrans(ctx context.Context, trans interface{}) context.Context {
  22. return context.WithValue(ctx, transCtx{}, trans)
  23. }
  24. func NewFileExpireContext(ctx context.Context, v interface{}) context.Context {
  25. return context.WithValue(ctx, fileExpireKey{}, v)
  26. }
  27. func NewFileHashContext(ctx context.Context, v interface{}) context.Context {
  28. return context.WithValue(ctx, fileHashKey{}, v)
  29. }
  30. func FromFileHashContext(ctx context.Context) (interface{}, bool) {
  31. v := ctx.Value(fileHashKey{})
  32. return v, v != nil
  33. }
  34. func FromFileExpireContext(ctx context.Context) (interface{}, bool) {
  35. v := ctx.Value(fileExpireKey{})
  36. return v, v != nil
  37. }
  38. // FromTrans 从上下文中获取事务
  39. func FromTrans(ctx context.Context) (interface{}, bool) {
  40. v := ctx.Value(transCtx{})
  41. return v, v != nil
  42. }
  43. // NewTransLock 创建事务锁的上下文
  44. func NewTransLock(ctx context.Context) context.Context {
  45. return context.WithValue(ctx, transLockCtx{}, struct{}{})
  46. }
  47. // FromTransLock 从上下文中获取事务锁
  48. func FromTransLock(ctx context.Context) bool {
  49. v := ctx.Value(transLockCtx{})
  50. return v != nil
  51. }
  52. // NewUserID 创建用户ID的上下文
  53. func NewUserID(ctx context.Context, userID string) context.Context {
  54. return context.WithValue(ctx, userIDCtx{}, userID)
  55. }
  56. // FromUserID 从上下文中获取用户ID
  57. func FromUserID(ctx context.Context) (string, bool) {
  58. v := ctx.Value(userIDCtx{})
  59. if v != nil {
  60. if s, ok := v.(string); ok {
  61. return s, s != ""
  62. }
  63. }
  64. return "", false
  65. }
  66. // NewTraceID 创建追踪ID的上下文
  67. func NewTraceID(ctx context.Context, traceID string) context.Context {
  68. return context.WithValue(ctx, traceIDCtx{}, traceID)
  69. }
  70. // FromTraceID 从上下文中获取追踪ID
  71. func FromTraceID(ctx context.Context) (string, bool) {
  72. v := ctx.Value(traceIDCtx{})
  73. if v != nil {
  74. if s, ok := v.(string); ok {
  75. return s, s != ""
  76. }
  77. }
  78. return "", false
  79. }
  80. // NewFileSizeLimitContext returns a new Context that carries value fsl.
  81. func NewFileSizeLimitContext(ctx context.Context, fsl FileSizeLimitHandle) context.Context {
  82. return context.WithValue(ctx, fileSizeLimitKey{}, fsl)
  83. }
  84. // FromFileSizeLimitContext returns the FileSizeLimitHandle value stored in ctx, if any.
  85. func FromFileSizeLimitContext(ctx context.Context) (FileSizeLimitHandle, bool) {
  86. handle, ok := ctx.Value(fileSizeLimitKey{}).(FileSizeLimitHandle)
  87. return handle, ok
  88. }
  89. // NewFileNameContext returns a new Context that carries value fn.
  90. func NewFileNameContext(ctx context.Context, fn FileNameHandle) context.Context {
  91. return context.WithValue(ctx, fileNameKey{}, fn)
  92. }
  93. // FromFileNameContext returns the FileNameHandle value stored in ctx, if any.
  94. func FromFileNameContext(ctx context.Context) (FileNameHandle, bool) {
  95. handle, ok := ctx.Value(fileNameKey{}).(FileNameHandle)
  96. return handle, ok
  97. }