Repository: danielgtaylor/huma
Pull request #369 (2024-04-10 16:09:23)

Add the ability to get ctx.Status()
For middlewares that needs to log the status result of a request, or needs to commit/rollback transactions based on the success of the handler, it'd be handy to be able to query the response status. 

In a [recent comment](https://github.com/danielgtaylor/huma/issues/302#issuecomment-1993236171) you mentioned another solution, 
```go
type MyContext struct {
  huma.Context
  status int
}

func (c *MyContext) SetStatus(code int) {
  c.status = code
  c.Context.SetStatus(code)
}
```

Except this doesn't work because `huma.Context` conflicts with the method named `Context()` in the interface definition of `huma.Context`. It gives the error: `cannot use myCtx (variable of type *MyContext) as huma.Context value in argument to next: *MyContext does not implement huma.Context (*MyContext.Context is a field, not a method)`. Perhaps type aliasing or assigning huma.Context to a new type might be a work-around, but we can make this easy for all Huma users.

This PR adds `Status()` to the interface and updates all the adapter contexts.



Patch:
diff --git a/adapters/humabunrouter/humabunrouter.go b/adapters/humabunrouter/humabunrouter.go
index aad5fde..0887535 100644
--- a/adapters/humabunrouter/humabunrouter.go
+++ b/adapters/humabunrouter/humabunrouter.go
@@ -20,11 +20,15 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type bunContext struct {
-	op *huma.Operation
-	r  bunrouter.Request
-	w  http.ResponseWriter
+	op     *huma.Operation
+	r      bunrouter.Request
+	w      http.ResponseWriter
+	status int
 }
 
+// check that bunContext implements huma.Context
+var _ huma.Context = &bunContext{}
+
 func (c *bunContext) Operation() *huma.Operation {
 	return c.op
 }
@@ -79,9 +83,14 @@ func (c *bunContext) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *bunContext) SetStatus(code int) {
+	c.status = code
 	c.w.WriteHeader(code)
 }
 
+func (c *bunContext) Status() int {
+	return c.status
+}
+
 func (c *bunContext) AppendHeader(name string, value string) {
 	c.w.Header().Add(name, value)
 }
@@ -100,9 +109,10 @@ func NewContext(op *huma.Operation, r bunrouter.Request, w http.ResponseWriter)
 }
 
 type bunCompatContext struct {
-	op *huma.Operation
-	r  *http.Request
-	w  http.ResponseWriter
+	op     *huma.Operation
+	r      *http.Request
+	w      http.ResponseWriter
+	status int
 }
 
 func (c *bunCompatContext) Operation() *huma.Operation {
@@ -160,9 +170,14 @@ func (c *bunCompatContext) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *bunCompatContext) SetStatus(code int) {
+	c.status = code
 	c.w.WriteHeader(code)
 }
 
+func (c *bunCompatContext) Status() int {
+	return c.status
+}
+
 func (c *bunCompatContext) AppendHeader(name string, value string) {
 	c.w.Header().Add(name, value)
 }
diff --git a/adapters/humachi/humachi.go b/adapters/humachi/humachi.go
index 547820a..cb42d9e 100644
--- a/adapters/humachi/humachi.go
+++ b/adapters/humachi/humachi.go
@@ -19,12 +19,16 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type chiContext struct {
-	op *huma.Operation
-	r  *http.Request
-	w  http.ResponseWriter
-	v4 bool
+	op     *huma.Operation
+	r      *http.Request
+	w      http.ResponseWriter
+	v4     bool
+	status int
 }
 
+// check that chiContext implements huma.Context
+var _ huma.Context = &chiContext{}
+
 func (c *chiContext) Operation() *huma.Operation {
 	return c.op
 }
@@ -83,9 +87,14 @@ func (c *chiContext) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *chiContext) SetStatus(code int) {
+	c.status = code
 	c.w.WriteHeader(code)
 }
 
+func (c *chiContext) Status() int {
+	return c.status
+}
+
 func (c *chiContext) AppendHeader(name string, value string) {
 	c.w.Header().Add(name, value)
 }
diff --git a/adapters/humaecho/humaecho.go b/adapters/humaecho/humaecho.go
index 3642b57..c956add 100644
--- a/adapters/humaecho/humaecho.go
+++ b/adapters/humaecho/humaecho.go
@@ -18,10 +18,14 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type echoCtx struct {
-	op   *huma.Operation
-	orig echo.Context
+	op     *huma.Operation
+	orig   echo.Context
+	status int
 }
 
+// check that echoCtx implements huma.Context
+var _ huma.Context = &echoCtx{}
+
 func (c *echoCtx) Operation() *huma.Operation {
 	return c.op
 }
@@ -76,9 +80,14 @@ func (c *echoCtx) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *echoCtx) SetStatus(code int) {
+	c.status = code
 	c.orig.Response().WriteHeader(code)
 }
 
+func (c *echoCtx) Status() int {
+	return c.status
+}
+
 func (c *echoCtx) AppendHeader(name, value string) {
 	c.orig.Response().Header().Add(name, value)
 }
diff --git a/adapters/humafiber/humafiber.go b/adapters/humafiber/humafiber.go
index c8d3207..b4e6d8d 100644
--- a/adapters/humafiber/humafiber.go
+++ b/adapters/humafiber/humafiber.go
@@ -15,10 +15,14 @@ import (
 )
 
 type fiberCtx struct {
-	op   *huma.Operation
-	orig *fiber.Ctx
+	op     *huma.Operation
+	orig   *fiber.Ctx
+	status int
 }
 
+// check that fiberCtx implements huma.Context
+var _ huma.Context = &fiberCtx{}
+
 func (c *fiberCtx) Operation() *huma.Operation {
 	return c.op
 }
@@ -84,9 +88,13 @@ func (c *fiberCtx) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *fiberCtx) SetStatus(code int) {
+	c.status = code
 	c.orig.Status(code)
 }
 
+func (c *fiberCtx) Status() int {
+	return c.status
+}
 func (c *fiberCtx) AppendHeader(name string, value string) {
 	c.orig.Append(name, value)
 }
diff --git a/adapters/humaflow/humaflow.go b/adapters/humaflow/humaflow.go
index 87f8390..4ee29a4 100644
--- a/adapters/humaflow/humaflow.go
+++ b/adapters/humaflow/humaflow.go
@@ -19,9 +19,10 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type goContext struct {
-	op *huma.Operation
-	r  *http.Request
-	w  http.ResponseWriter
+	op     *huma.Operation
+	r      *http.Request
+	w      http.ResponseWriter
+	status int
 }
 
 func (c *goContext) Operation() *huma.Operation {
@@ -78,9 +79,14 @@ func (c *goContext) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *goContext) SetStatus(code int) {
+	c.status = code
 	c.w.WriteHeader(code)
 }
 
+func (c *goContext) Status() int {
+	return c.status
+}
+
 func (c *goContext) AppendHeader(name string, value string) {
 	c.w.Header().Add(name, value)
 }
diff --git a/adapters/humagin/humagin.go b/adapters/humagin/humagin.go
index c99b052..992ebf6 100644
--- a/adapters/humagin/humagin.go
+++ b/adapters/humagin/humagin.go
@@ -18,10 +18,14 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type ginCtx struct {
-	op   *huma.Operation
-	orig *gin.Context
+	op     *huma.Operation
+	orig   *gin.Context
+	status int
 }
 
+// check that ginCtx implements huma.Context
+var _ huma.Context = &ginCtx{}
+
 func (c *ginCtx) Operation() *huma.Operation {
 	return c.op
 }
@@ -76,9 +80,14 @@ func (c *ginCtx) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *ginCtx) SetStatus(code int) {
+	c.status = code
 	c.orig.Status(code)
 }
 
+func (c *ginCtx) Status() int {
+	return c.status
+}
+
 func (c *ginCtx) AppendHeader(name string, value string) {
 	c.orig.Writer.Header().Add(name, value)
 }
diff --git a/adapters/humago/humago.go b/adapters/humago/humago.go
index 6a3ca67..5fa7d7b 100644
--- a/adapters/humago/humago.go
+++ b/adapters/humago/humago.go
@@ -18,11 +18,15 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type goContext struct {
-	op *huma.Operation
-	r  *http.Request
-	w  http.ResponseWriter
+	op     *huma.Operation
+	r      *http.Request
+	w      http.ResponseWriter
+	status int
 }
 
+// check that goContext implements huma.Context
+var _ huma.Context = &goContext{}
+
 func (c *goContext) Operation() *huma.Operation {
 	return c.op
 }
@@ -84,9 +88,14 @@ func (c *goContext) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *goContext) SetStatus(code int) {
+	c.status = code
 	c.w.WriteHeader(code)
 }
 
+func (c *goContext) Status() int {
+	return c.status
+}
+
 func (c *goContext) AppendHeader(name string, value string) {
 	c.w.Header().Add(name, value)
 }
diff --git a/adapters/humahttprouter/humahttprouter.go b/adapters/humahttprouter/humahttprouter.go
index 46f9530..6b99a8e 100644
--- a/adapters/humahttprouter/humahttprouter.go
+++ b/adapters/humahttprouter/humahttprouter.go
@@ -19,12 +19,16 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type httprouterContext struct {
-	op *huma.Operation
-	r  *http.Request
-	w  http.ResponseWriter
-	ps httprouter.Params
+	op     *huma.Operation
+	r      *http.Request
+	w      http.ResponseWriter
+	ps     httprouter.Params
+	status int
 }
 
+// check that httprouterContext implements huma.Context
+var _ huma.Context = &httprouterContext{}
+
 func (c *httprouterContext) Operation() *huma.Operation {
 	return c.op
 }
@@ -79,9 +83,14 @@ func (c *httprouterContext) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *httprouterContext) SetStatus(code int) {
+	c.status = code
 	c.w.WriteHeader(code)
 }
 
+func (c *httprouterContext) Status() int {
+	return c.status
+}
+
 func (c *httprouterContext) AppendHeader(name string, value string) {
 	c.w.Header().Add(name, value)
 }
diff --git a/adapters/humamux/humamux.go b/adapters/humamux/humamux.go
index 0bdfda6..108b61a 100644
--- a/adapters/humamux/humamux.go
+++ b/adapters/humamux/humamux.go
@@ -18,11 +18,15 @@ import (
 var MultipartMaxMemory int64 = 8 * 1024
 
 type gmuxContext struct {
-	op *huma.Operation
-	r  *http.Request
-	w  http.ResponseWriter
+	op     *huma.Operation
+	r      *http.Request
+	w      http.ResponseWriter
+	status int
 }
 
+// check that gmuxContext implements huma.Context
+var _ huma.Context = &gmuxContext{}
+
 func (c *gmuxContext) Operation() *huma.Operation {
 	return c.op
 }
@@ -77,9 +81,14 @@ func (c *gmuxContext) SetReadDeadline(deadline time.Time) error {
 }
 
 func (c *gmuxContext) SetStatus(code int) {
+	c.status = code
 	c.w.WriteHeader(code)
 }
 
+func (c *gmuxContext) Status() int {
+	return c.status
+}
+
 func (c *gmuxContext) AppendHeader(name string, value string) {
 	c.w.Header().Add(name, value)
 }
diff --git a/api.go b/api.go
index 0972ff3..d14d47a 100644
--- a/api.go
+++ b/api.go
@@ -100,6 +100,9 @@ type Context interface {
 	// SetStatus sets the HTTP status code for the response.
 	SetStatus(code int)
 
+	// Status returns the HTTP status code for the response.
+	Status() int
+
 	// SetHeader sets the given header to the given value, overwriting any
 	// existing value. Use `AppendHeader` to append a value instead.
 	SetHeader(name, value string)


Test patch:
diff --git a/api_test.go b/api_test.go
index 3eeca4d..87c9e18 100644
--- a/api_test.go
+++ b/api_test.go
@@ -77,6 +77,7 @@ func TestContextValue(t *testing.T) {
 		// Make an updated context available to the handler.
 		ctx = huma.WithValue(ctx, "foo", "bar")
 		next(ctx)
+		assert.Equal(t, http.StatusNoContent, ctx.Status())
 	})
 
 	// Register a simple hello world operation in the API.