call OnRequest before handler is called and handle err

webrpc · Sep 30, 2024 · 1afdb0d · 1afdb0d
1 parent c5c5528
commit 1afdb0d
Show file tree

Hide file tree

Showing 5 changed files with 97 additions and 32 deletions.
diff --git a/_examples/golang-basics/example.gen.go b/_examples/golang-basics/example.gen.go
diff --git a/_examples/golang-imports/api.gen.go b/_examples/golang-imports/api.gen.go
diff --git a/helpers.go.tmpl b/helpers.go.tmpl
@@ -5,6 +5,12 @@
 // Helpers
 //
 
+type MethodCtx struct  {
+  Name string
+  Service string
+  Annotations map[string]string
+}
+
 type contextKey struct {
 	name string
 }
@@ -27,7 +33,7 @@ var (
 
 	MethodNameCtxKey = &contextKey{"MethodName"}
 
-	MethodAnnotationsCtxKey = &contextKey{"MethodAnnotations"}
+	methodAnnotationsCtxKey = &contextKey{"MethodAnnotations"}
 )
 
 func ServiceNameFromContext(ctx context.Context) string {
@@ -45,12 +51,19 @@ func RequestFromContext(ctx context.Context) *http.Request {
 	return r
 }
 
-func MethodAnnotationsFromContext(ctx context.Context) map[string]string {
-	annotations, _ := ctx.Value(MethodAnnotationsCtxKey).(map[string]string)
-	return annotations
+func MethodFromContext(ctx context.Context) MethodCtx {
+	name, _ := ctx.Value(MethodNameCtxKey).(string)
+	service, _ := ctx.Value(ServiceNameCtxKey).(string)
+	annotations, _ := ctx.Value(methodAnnotationsCtxKey).(map[string]string)
+
+	return MethodCtx{
+		Name: name,
+		Service: service,
+		Annotations: annotations,
+	}
 }
 
-{{- if $opts.server}}
+{{ if $opts.server}}
 func ResponseWriterFromContext(ctx context.Context) http.ResponseWriter {
 	w, _ := ctx.Value(HTTPResponseWriterCtxKey).(http.ResponseWriter)
 	return w

diff --git a/server.go.tmpl b/server.go.tmpl
@@ -42,11 +42,7 @@ func (s *{{$serviceName}}) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ctx = context.WithValue(ctx, HTTPResponseWriterCtxKey, w)
 	ctx = context.WithValue(ctx, HTTPRequestCtxKey, r)
 	ctx = context.WithValue(ctx, ServiceNameCtxKey, "{{.Name}}")
-	ctx = context.WithValue(ctx, MethodAnnotationsCtxKey, methodAnnotations[r.URL.Path])
-
-    if s.OnRequest != nil {
-        s.OnRequest(w, r)
-    }
+	ctx = context.WithValue(ctx, methodAnnotationsCtxKey, methodAnnotations[r.URL.Path])
 
 	var handler func(ctx context.Context, w http.ResponseWriter, r *http.Request)
 	switch r.URL.Path {
@@ -75,6 +71,17 @@ func (s *{{$serviceName}}) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	switch contentType {
 	case "application/json":
+        if s.OnRequest != nil {
+            if err := s.OnRequest(w, r); err != nil {
+                rpcErr, ok := err.(WebRPCError)
+                if !ok {
+                    rpcErr = ErrWebrpcEndpoint.WithCause(err)
+                }
+                s.sendErrorJSON(w, r, rpcErr)
+                return
+            }
+        }
+
 		handler(ctx, w, r)
 	default:
 		err := ErrWebrpcBadRequest.WithCause(fmt.Errorf("unsupported Content-Type %q (only application/json is allowed)", r.Header.Get("Content-Type")))

diff --git a/types.go.tmpl b/types.go.tmpl
@@ -29,7 +29,7 @@
 
 {{- end }}
 
-{{- range $_, $service := $services -}}
+{{ range $_, $service := $services -}}
 var (
 	methodAnnotations = map[string]map[string]string{
     	{{- range $_, $method := $service.Methods }}