From 860d8f4fe812dafce31cb76c11b649bc607d748c Mon Sep 17 00:00:00 2001
From: Alice Wasko <alicewasko@datawire.io>
Date: Fri, 23 Feb 2024 08:27:21 -0800
Subject: [PATCH] update go generator template (#3)

* add openapi code gen config

Signed-off-by: Alice Wasko <alicewasko@datawire.io>

* update log middleware to log response codes

Signed-off-by: Alice Wasko <alicewasko@datawire.io>

* add codegen injection

Signed-off-by: Alice Wasko <alicewasko@datawire.io>

* update go generator version

Signed-off-by: Alice Wasko <alicewasko@datawire.io>

---------

Signed-off-by: Alice Wasko <alicewasko@datawire.io>
---
 generators.json                               |  2 +-
 .../base/internal/binary/handlers.go.template | 46 +++++++++++++++++++
 .../base/internal/binary/server.go.template   | 19 +++-----
 generators/go/generator-config.yaml           |  7 ++-
 4 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/generators.json b/generators.json
index c9a1146..e4c610a 100644
--- a/generators.json
+++ b/generators.json
@@ -3,7 +3,7 @@
     {
       "name": "Ambassador Labs Go Generator",
       "description": "Generates a new project for an API server writen in go",
-      "version": "1.0.0",
+      "version": "0.0.1",
       "languages": ["go"],
       "dir_name": "go"
     }]
diff --git a/generators/go/base/internal/binary/handlers.go.template b/generators/go/base/internal/binary/handlers.go.template
index 7494680..79bb1fa 100644
--- a/generators/go/base/internal/binary/handlers.go.template
+++ b/generators/go/base/internal/binary/handlers.go.template
@@ -3,6 +3,7 @@ package api
 import (
 	"log"
 	"net/http"
+	"time"
 )
 
 // Handles returning a 404 to the client and logging a message for unsupported routes
@@ -16,3 +17,48 @@ func handleAbout(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusOK)
 	w.Write([]byte(" {{{ .ProjectName }}}: server generated by Ambassador Labs"))
 }
+
+
+// responseWriter is a custom http.ResponseWriter that captures the status code.
+type responseWriter struct {
+	http.ResponseWriter
+	statusCode int
+	written    bool
+}
+
+// NewResponseWriter creates a new response writer.
+func NewResponseWriter(w http.ResponseWriter) *responseWriter {
+	return &responseWriter{ResponseWriter: w, statusCode: http.StatusOK}
+}
+
+// WriteHeader captures the status code and delegates to the original ResponseWriter, ensuring it's only called once.
+func (rw *responseWriter) WriteHeader(code int) {
+	if rw.written {
+		return // Prevent multiple WriteHeader calls
+	}
+	rw.statusCode = code
+	rw.ResponseWriter.WriteHeader(code)
+	rw.written = true
+}
+
+// Write ensures WriteHeader is called with the default status code before any response body is written.
+func (rw *responseWriter) Write(b []byte) (int, error) {
+	if !rw.written {
+		rw.WriteHeader(http.StatusOK) // Ensure WriteHeader is called at least once
+	}
+	return rw.ResponseWriter.Write(b)
+}
+
+// logMiddleware logs the request method, URL path, duration, and status code.
+func logMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		start := time.Now()
+
+		// Use the responseWriter wrapper to capture the status code and handle Write calls properly.
+		wrappedWriter := NewResponseWriter(w)
+		next.ServeHTTP(wrappedWriter, r)
+
+		// Log with the captured status code.
+		log.Printf("%s %s %d %s\n", r.Method, r.URL.Path, wrappedWriter.statusCode, time.Since(start))
+	})
+}
diff --git a/generators/go/base/internal/binary/server.go.template b/generators/go/base/internal/binary/server.go.template
index b9aa1d2..953d314 100644
--- a/generators/go/base/internal/binary/server.go.template
+++ b/generators/go/base/internal/binary/server.go.template
@@ -5,6 +5,7 @@ import (
 	"log"
 	"net/http"
 
+  "github.com/gorilla/mux"
 	"golang.org/x/sync/errgroup"
 	ctrl "sigs.k8s.io/controller-runtime"
 )
@@ -45,12 +46,14 @@ func (s *apiService) Start() {
 
 // Starts the sub server for the api endpoints
 func (s *apiService) Serve(ctx context.Context) error {
-	mux := http.NewServeMux()
+	router := mux.NewRouter().StrictSlash(true)
 
-	mux.HandleFunc("/", handle404)
-	mux.HandleFunc("/about", handleAbout)
+	router.HandleFunc("/", handle404)
+	router.HandleFunc("/about", handleAbout)
 
-	loggedMux := logMiddleware(mux)
+  {{{ codeGenGorillaMux "router" }}}
+
+	loggedMux := logMiddleware(router)
 
 	server := &http.Server{
 		Addr:    s.addr,
@@ -75,11 +78,3 @@ func (s *apiService) Serve(ctx context.Context) error {
 		return err
 	}
 }
-
-// Lightweight middleware for logging basic info about each request
-func logMiddleware(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		log.Printf("%s %s\n", r.Method, r.URL.Path) // TODO: log status code
-		next.ServeHTTP(w, r)
-	})
-}
diff --git a/generators/go/generator-config.yaml b/generators/go/generator-config.yaml
index 260096e..7342b3b 100644
--- a/generators/go/generator-config.yaml
+++ b/generators/go/generator-config.yaml
@@ -1,7 +1,7 @@
 metadata:
   name: "Ambassador Labs Go Generator"
   description: "Generates a new project for an API server writen in go"
-  version: "1.0.0"
+  version: "0.0.1"
   languages: ["go"]
 
 # Defines the base directory within the template folder.
@@ -11,6 +11,11 @@ metadata:
 # that will be built later.
 baseDirectory: base
 
+codeGeneration:
+  codeGenLanguage: "go"
+  outputDirectory: "base/internal/api"
+  indentStyle: "tab"
+
 # Delimiters define the characters that surround a template variable that needs to be built
 # These will only be evaluated when they are in .template files
 leftDelimiter: "{{{"