Add status code recorder to middleware

Author: simondanielsson

Makefile

@@ -32,4 +32,4 @@ clean:
 .PHONY: ci
 ci: clean lint test
 
-.DEFAULT_GOAL := build
+.DEFAULT_GOAL := api

cmd/internal/constants.go

@@ -0,0 +1,14 @@
+package constants
+
+// Context types
+type RepositoryKeyType string
+
+const RepositoryKey RepositoryKeyType = "repo_key"
+
+type DBConnPoolKey string
+
+const DBConnPool DBConnPoolKey = "db_pool_key"
+
+type StatusCodeKeyType string
+
+const StatusCodeKey StatusCodeKeyType = "status_code_key"

cmd/internal/db/middleware.go

@@ -2,47 +2,54 @@ package db
 
 import (
 	"context"
+	"fmt"
 	"log"
 	"net/http"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	constants "github.com/simondanielsson/recite/cmd/internal"
 	"github.com/simondanielsson/recite/cmd/internal/queries"
 )
 
-type RepositoryKeyType string
-
-const RepositoryKey RepositoryKeyType = "repo_key"
-
-type statusRecorder struct {
-	http.ResponseWriter
-	status int
-}
-
 func AddDatabaseMiddleware(handler http.Handler, pool *pgxpool.Pool, logger *log.Logger) http.Handler {
 	return http.HandlerFunc(
 		func(w http.ResponseWriter, r *http.Request) {
 			ctx := r.Context()
-			tx, err := pool.BeginTx(ctx, pgx.TxOptions{
-				IsoLevel:   pgx.Serializable,
-				AccessMode: pgx.ReadWrite,
-			})
+			repository, tx, err := NewRepositoryWithTx(ctx, pool)
 			if err != nil {
-				http.Error(w, "failed to start transaction", http.StatusInternalServerError)
+				http.Error(w, err.Error(), http.StatusInternalServerError)
 			}
-			repository := queries.New(pool).WithTx(tx)
 
-			ctx = context.WithValue(ctx, RepositoryKey, repository)
+			// Attach transaction-scoped repository, and pool for transactions to be launched outside request context
+			ctx = context.WithValue(ctx, constants.RepositoryKey, repository)
+			ctx = context.WithValue(ctx, constants.DBConnPool, pool)
 
-			rec := statusRecorder{ResponseWriter: w, status: http.StatusOK}
-			handler.ServeHTTP(rec, r.WithContext(ctx))
+			handler.ServeHTTP(w, r.WithContext(ctx))
 
-			if rec.status >= http.StatusBadRequest {
+			status, ok := ctx.Value(constants.StatusCodeKey).(int)
+			if !ok {
+				_ = tx.Rollback(ctx)
+				return
+			}
+
+			if status >= http.StatusBadRequest {
 				_ = tx.Rollback(ctx)
 			} else if err := tx.Commit(ctx); err != nil {
 				// Response has already been sent - just log
-				logger.Printf("commit tx failed: %v", err)
+				logger.Printf("tx commit failed: %v", err)
 			}
 		},
 	)
 }
+
+func NewRepositoryWithTx(ctx context.Context, pool *pgxpool.Pool) (*queries.Queries, pgx.Tx, error) {
+	tx, err := pool.BeginTx(ctx, pgx.TxOptions{
+		IsoLevel:   pgx.Serializable,
+		AccessMode: pgx.ReadWrite,
+	})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to start transaction: %w", err)
+	}
+	return queries.New(pool).WithTx(tx), tx, nil
+}

cmd/internal/logging/logging.go

@@ -3,21 +3,12 @@ package logging
 import (
 	"io"
 	"log"
-	"net/http"
-	"time"
 )
 
+const DateFormat string = "Mon Jan 2 15:04:05 MST 2006"
+
 func NewLogger(w io.Writer) *log.Logger {
 	logger := log.Logger{}
 	logger.SetOutput(w)
 	return &logger
 }
-
-func AddLoggingMiddleware(handler http.Handler, logger *log.Logger) http.Handler {
-	return http.HandlerFunc(
-		func(w http.ResponseWriter, r *http.Request) {
-			handler.ServeHTTP(w, r)
-			logger.Printf("%s | %s %s\n", time.Now().Format("Mon Jan 2 15:04:05 MST 2006"), r.Method, r.URL)
-		},
-	)
-}

cmd/internal/logging/middleware.go

@@ -0,0 +1,24 @@
+package logging
+
+import (
+	"log"
+	"net/http"
+	"time"
+
+	constants "github.com/simondanielsson/recite/cmd/internal"
+)
+
+func AddLoggingMiddleware(handler http.Handler, logger *log.Logger) http.Handler {
+	return http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			handler.ServeHTTP(w, r)
+			status, ok := r.Context().Value(constants.StatusCodeKey).(int)
+			if !ok {
+				// TODO: write proper error, or assume it failed
+				panic("could not get status")
+			}
+
+			logger.Printf("%s | %s %s - %d\n", time.Now().Format(DateFormat), r.Method, r.URL, status)
+		},
+	)
+}

cmd/internal/marshal/marshal.go

@@ -1,14 +1,22 @@
 package marshal
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
+
+	constants "github.com/simondanielsson/recite/cmd/internal"
 )
 
-func Encode[T any](w http.ResponseWriter, r *http.Request, status int, v T) error {
+func Encode[T any](ctx context.Context, w http.ResponseWriter, r *http.Request, status int, v T) error {
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(status)
+
+	// A bit hacky way to write override the existing context
+	ctx = context.WithValue(ctx, constants.StatusCodeKey, status)
+	*r = *(r.WithContext(ctx))
+
 	if err := json.NewEncoder(w).Encode(v); err != nil {
 		return fmt.Errorf("encode json: %w", err)
 	}

cmd/internal/routes/routes.go

@@ -1,12 +1,14 @@
 package routes
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
 	"strconv"
 
-	"github.com/simondanielsson/recite/cmd/internal/db"
+	"github.com/jackc/pgx/v5/pgxpool"
+	constants "github.com/simondanielsson/recite/cmd/internal"
 	"github.com/simondanielsson/recite/cmd/internal/marshal"
 	"github.com/simondanielsson/recite/cmd/internal/queries"
 	"github.com/simondanielsson/recite/cmd/internal/services"
@@ -28,6 +30,7 @@ func rootGetHandler(logger *log.Logger) http.Handler {
 		func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusOK)
 			if _, err := w.Write([]byte("ok\n")); err != nil {
+				logger.Println(err)
 				w.WriteHeader(http.StatusInternalServerError)
 			}
 		},
@@ -49,7 +52,7 @@ func recitalPostHandler(logger *log.Logger) http.Handler {
 				res := messageResponse{
 					Message: fmt.Sprintf("bad request, should contain url. Got %s", req.Url),
 				}
-				if err := marshal.Encode(w, r, http.StatusBadRequest, res); err != nil {
+				if err := marshal.Encode(r.Context(), w, r, http.StatusBadRequest, res); err != nil {
 					writeErrHeader(w, err, logger)
 				}
 				return
@@ -62,71 +65,69 @@ func recitalPostHandler(logger *log.Logger) http.Handler {
 			// 4. Build a simple voice streamer in a basic htmx web page
 
 			ctx := r.Context()
-			repository, ok := ctx.Value(db.RepositoryKey).(*queries.Queries)
+			repository, ok := ctx.Value(constants.RepositoryKey).(*queries.Queries)
 			if !ok {
-				res := messageResponse{Message: "Something went wrong."}
-				logger.Print(err)
-				if err := marshal.Encode(w, r, int(http.StatusInternalServerError), res); err != nil {
-					writeErrHeader(w, err, logger)
-					return
-				}
+				logGenericInternalServiceError(ctx, w, r, fmt.Errorf("failed loading repository"), logger)
+				return
+			}
+			pool, ok := ctx.Value(constants.DBConnPool).(*pgxpool.Pool)
+			if !ok {
+				logGenericInternalServiceError(ctx, w, r, fmt.Errorf("failed loading db connection pool"), logger)
+				return
 			}
 
-			id, err := services.CreateRecital(ctx, req.Url, repository, logger)
+			id, err := services.CreateRecital(ctx, req.Url, repository, pool, logger)
 			if err != nil {
-				res := messageResponse{Message: "Something went wrong."}
-				logger.Print(err)
-				if err := marshal.Encode(w, r, int(http.StatusInternalServerError), res); err != nil {
-					writeErrHeader(w, err, logger)
-					return
-				}
+				logGenericInternalServiceError(ctx, w, r, err, logger)
+				return
 			}
 
 			res := response{
 				Id: id,
 			}
-			if err := marshal.Encode(w, r, http.StatusCreated, res); err != nil {
+			if err := marshal.Encode(ctx, w, r, http.StatusCreated, res); err != nil {
 				writeErrHeader(w, err, logger)
 				return
 			}
 		},
 	)
 }
 
-// TODO: move
-type GenerationStatus string
-
-const (
-	InProgress GenerationStatus = "in progress"
-	Completed  GenerationStatus = "completed"
-	Failed     GenerationStatus = "failed"
-)
-
 func recitalGetHandler(logger *log.Logger) http.Handler {
 	type response struct {
-		Status GenerationStatus
+		queries.Recital
 	}
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		id, err := strconv.Atoi(r.PathValue("id"))
 		if err != nil {
 			res := messageResponse{Message: "invalid id"}
-			if err := marshal.Encode(w, r, http.StatusBadRequest, res); err != nil {
+			if err := marshal.Encode(r.Context(), w, r, http.StatusBadRequest, res); err != nil {
 				writeErrHeader(w, err, logger)
-				return
 			}
+			return
 		}
 
-		_ = id
-		// TODO: unmock
-		res := response{Status: Completed}
-		if err := marshal.Encode(w, r, http.StatusOK, res); err != nil {
-			res := messageResponse{"failed to encode response"}
-			if err := marshal.Encode(w, r, http.StatusInternalServerError, res); err != nil {
-				writeErrHeader(w, err, logger)
+		ctx := r.Context()
+		repository, ok := ctx.Value(constants.RepositoryKey).(*queries.Queries)
+		if !ok {
+			logGenericInternalServiceError(ctx, w, r, fmt.Errorf("failed loading repository"), logger)
+		}
+
+		recital, err := services.GetRecital(ctx, int32(id), repository, logger)
+		if err != nil {
+			res := messageResponse{Message: fmt.Sprintf("Could not find recital with id %d", id)}
+			if err := marshal.Encode(ctx, w, r, http.StatusNotFound, res); err != nil {
+				logFailedEncodingResponse(ctx, w, r, err, logger)
 			}
+			return
 		}
 
-		return
+		res := response{
+			Recital: recital,
+		}
+		if err := marshal.Encode(ctx, w, r, http.StatusOK, res); err != nil {
+			logFailedEncodingResponse(ctx, w, r, err, logger)
+		}
 	})
 }
 
@@ -136,3 +137,20 @@ func writeErrHeader(w http.ResponseWriter, err error, logger *log.Logger) {
 		logger.Print("failed writing error")
 	}
 }
+
+func logGenericInternalServiceError(ctx context.Context, w http.ResponseWriter, r *http.Request, err error, logger *log.Logger) {
+	res := messageResponse{Message: "Something went wrong."}
+	logger.Print(err)
+	if err := marshal.Encode(ctx, w, r, int(http.StatusInternalServerError), res); err != nil {
+		writeErrHeader(w, err, logger)
+		return
+	}
+}
+
+func logFailedEncodingResponse(ctx context.Context, w http.ResponseWriter, r *http.Request, err error, logger *log.Logger) {
+	logger.Println(err)
+	res := messageResponse{"failed to encode response"}
+	if err := marshal.Encode(ctx, w, r, http.StatusInternalServerError, res); err != nil {
+		writeErrHeader(w, err, logger)
+	}
+}

cmd/internal/services/constants.go

@@ -5,3 +5,11 @@ import "path"
 const maxArticleContentCharLength int = 3500
 
 var baseOutputPath string = path.Join("data", "output", "generations", "audio")
+
+type Status string
+
+const (
+	InProgress Status = "in progress"
+	Completed  Status = "completed"
+	Failed     Status = "failed"
+)