feat: add rest.WithSSE to build SSE route easier (#4729)

This commit is contained in:
Kevin Wan
2025-03-22 13:38:13 +08:00
committed by GitHub
parent cdb0098b18
commit 6edfce63e3
13 changed files with 106 additions and 34 deletions

View File

@@ -15,6 +15,7 @@ import (
"github.com/zeromicro/go-zero/rest/handler"
"github.com/zeromicro/go-zero/rest/httpx"
"github.com/zeromicro/go-zero/rest/internal"
"github.com/zeromicro/go-zero/rest/internal/header"
"github.com/zeromicro/go-zero/rest/internal/response"
)
@@ -54,6 +55,9 @@ func newEngine(c RestConf) *engine {
}
func (ng *engine) addRoutes(r featuredRoutes) {
if r.sse {
r.routes = buildSSERoutes(r.routes)
}
ng.routes = append(ng.routes, r)
// need to guarantee the timeout is the max of all routes
@@ -63,6 +67,20 @@ func (ng *engine) addRoutes(r featuredRoutes) {
}
}
func buildSSERoutes(routes []Route) []Route {
for i, route := range routes {
h := route.Handler
routes[i].Handler = func(w http.ResponseWriter, r *http.Request) {
w.Header().Set(header.ContentType, header.ContentTypeEventStream)
w.Header().Set(header.CacheControl, header.CacheControlNoCache)
w.Header().Set(header.Connection, header.ConnectionKeepAlive)
h(w, r)
}
}
return routes
}
func (ng *engine) appendAuthHandler(fr featuredRoutes, chn chain.Chain,
verifier func(chain.Chain) chain.Chain) chain.Chain {
if fr.jwt.enabled {