Adds RequestIDHandler function to RequestID middleware

labstack · Jun 22, 2021 · 78a0ea2 · 78a0ea2
1 parent 1ac4a8f
commit 78a0ea2
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/middleware/request_id.go b/middleware/request_id.go
@@ -14,6 +14,9 @@ type (
  // Generator defines a function to generate an ID.
  // Optional. Default value random.String(32).
  Generator func() string
+
+ // RequestIDHandler defines a function which is executed for a request id.
+ RequestIDHandler func(echo.Context, string)
  }
 )
 
@@ -53,6 +56,9 @@ func RequestIDWithConfig(config RequestIDConfig) echo.MiddlewareFunc {
  rid = config.Generator()
  }
  res.Header().Set(echo.HeaderXRequestID, rid)
+ if config.RequestIDHandler != nil {
+ config.RequestIDHandler(c, rid)
+ }
 
  return next(c)
  }

diff --git a/middleware/request_id_test.go b/middleware/request_id_test.go
@@ -23,13 +23,20 @@ func TestRequestID(t *testing.T) {
  h(c)
  assert.Len(t, rec.Header().Get(echo.HeaderXRequestID), 32)
 
- // Custom generator
+ // Custom generator and handler
+ customID := "customGenerator"
+ calledHandler := false
  rid = RequestIDWithConfig(RequestIDConfig{
- Generator: func() string { return "customGenerator" },
+ Generator: func() string { return customID },
+ RequestIDHandler: func(_ echo.Context, id string) {
+ calledHandler = true
+ assert.Equal(t, customID, id)
+ },
  })
  h = rid(handler)
  h(c)
  assert.Equal(t, rec.Header().Get(echo.HeaderXRequestID), "customGenerator")
+ assert.True(t, calledHandler)
 }
 
 func TestRequestID_IDNotAltered(t *testing.T) {